dynamo-llm 1.6.2

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::collections::VecDeque;

#[cfg(test)]
use dynamo_kv_router::protocols::KvCacheEventData;
use dynamo_kv_router::{
    protocols::{DpRank, ResetScope, RouterEvent, WorkerId},
    recovery::{CursorObservation, CursorState},
};

pub(super) type RecoveryKey = (WorkerId, DpRank);

const RECOVERY_PENDING_LIVE_EVENT_LIMIT: usize = 1024;

pub(super) enum LiveEventAction {
    Ignore,
    Apply {
        event_id: u64,
        event: RouterEvent,
    },
    Clear {
        event_id: u64,
        event: RouterEvent,
    },
    Recover {
        start_event_id: Option<u64>,
        end_event_id: Option<u64>,
    },
    ResetDegraded {
        event: RouterEvent,
    },
}

#[derive(Clone, Debug, Default)]
pub(super) struct RankState {
    /// NOTE: This coordinator tracks the last event successfully admitted to the indexer's
    /// per-worker FIFO, not confirmed backend application. A trailing `Flush` proves queue
    /// progress, not event success. Stronger applied-state semantics require an explicit
    /// batch-result acknowledgement. Plan without mutating this cursor, advance only after every
    /// event in the commit group is admitted, and fence/reset on partial admission failure.
    pub(super) cursor: CursorState,
    pub(super) recovery_inflight: bool,
    pending_live_events: VecDeque<RouterEvent>,
    max_seen_live_id: Option<u64>,
}

impl RankState {
    pub(super) fn activate(&mut self, recoverable: bool) {
        *self = Self {
            recovery_inflight: recoverable,
            ..Self::default()
        };
    }

    pub(super) fn last_admitted_id(&self) -> Option<u64> {
        self.cursor.last_applied_id()
    }

    #[cfg(feature = "ckf-diagnostics")]
    pub(super) fn pending_live_event_count(&self) -> usize {
        self.pending_live_events.len()
    }

    pub(super) fn observe_live_event(
        &mut self,
        event: RouterEvent,
        recoverable: bool,
    ) -> LiveEventAction {
        let event_id = event.event.event_id;

        if matches!(event.reset_scope(), Ok(Some(ResetScope::All))) {
            if self
                .last_admitted_id()
                .is_some_and(|last_admitted_id| event_id <= last_admitted_id)
            {
                return LiveEventAction::Ignore;
            }
            return LiveEventAction::Clear { event_id, event };
        }

        match self.cursor.observe(event_id) {
            CursorObservation::Stale { .. } => LiveEventAction::Ignore,
            observation if self.recovery_inflight => {
                if matches!(
                    observation,
                    CursorObservation::Initial { .. }
                        | CursorObservation::Contiguous { .. }
                        | CursorObservation::Gap { .. }
                ) {
                    self.observe_and_buffer(event);
                }
                LiveEventAction::Ignore
            }
            CursorObservation::Initial { .. } if recoverable => {
                self.observe_and_buffer(event.clone());
                self.recovery_inflight = true;
                LiveEventAction::Recover {
                    start_event_id: None,
                    end_event_id: None,
                }
            }
            CursorObservation::Gap { expected, .. } if recoverable => {
                // NOTE: KV RECOVERY CONTRACT: Ordinary gaps request the next expected ID
                // and preserve the existing index/cursor. Only the server can decide whether
                // its retained history supports Events or requires TreeDump; never pre-clear
                // the rank or request a snapshot here. Initial recovery and source replacement
                // are separate lifecycle cases. See retained_gap_replays_without_reset and
                // expired_gap_uses_server_selected_snapshot in worker_query.rs.
                self.observe_and_buffer(event);
                self.recovery_inflight = true;
                LiveEventAction::Recover {
                    start_event_id: Some(expected),
                    end_event_id: None,
                }
            }
            CursorObservation::Gap { .. } => LiveEventAction::ResetDegraded { event },
            CursorObservation::Initial { got } | CursorObservation::Contiguous { got } => {
                LiveEventAction::Apply {
                    event_id: got,
                    event,
                }
            }
        }
    }

    pub(super) fn commit_live_admission(&mut self, event_id: u64) {
        self.cursor = self.cursor.advance_to(event_id);
        self.clear_max_seen_if_caught_up(event_id);
    }

    pub(super) fn pending_live_watermark(&self) -> Option<u64> {
        self.max_seen_live_id
    }

    pub(super) fn discard_recovery_before_clear(&mut self) {
        self.recovery_inflight = false;
        self.pending_live_events.clear();
        self.max_seen_live_id = None;
    }

    pub(super) fn finish_failed_recovery(&mut self) {
        self.recovery_inflight = false;
        self.pending_live_events.clear();
        self.max_seen_live_id = None;
    }

    /// Leave authoritative state and the admission cursor unchanged after a
    /// non-authoritative snapshot failure.
    ///
    /// The production fetch path performs bounded retries before completion. This
    /// transition is the defensive fallback for a failure delivered directly to
    /// the state machine; it deliberately waits for the next live event instead of
    /// starting an unbounded autonomous retry loop.
    pub(super) fn retry_after_failed_snapshot(&mut self) {
        self.recovery_inflight = false;
    }

    /// Buffer a live event behind an in-flight source snapshot.
    pub(super) fn buffer_recovery_tail(&mut self, event: RouterEvent) {
        self.recovery_inflight = true;
        self.observe_and_buffer(event);
    }

    /// Drain the buffered suffix after an advisory recovery response.
    ///
    /// Missing IDs do not block the suffix. Worker-query recovery calls this on a
    /// clone and commits it only after queue admission; state-agent recovery owns
    /// its own admission and fencing policy.
    pub(super) fn drain_advisory_tail_after(&mut self, recovered_through: u64) -> Vec<RouterEvent> {
        self.cursor = CursorState::Initial.advance_to(recovered_through);
        let events = self.take_failed_recovery_degraded();
        let last_event_id = events.last().map(|event| event.event.event_id);
        self.commit_failed_recovery_degraded(last_event_id);
        events
    }

    pub(super) fn take_failed_recovery_degraded(&mut self) -> Vec<RouterEvent> {
        let last_admitted_id = self.last_admitted_id().unwrap_or(0);
        let mut events: Vec<_> = self.pending_live_events.drain(..).collect();
        events.sort_unstable_by_key(|event| event.event.event_id);
        events.dedup_by_key(|event| event.event.event_id);
        events.retain(|event| event.event.event_id > last_admitted_id);
        events
    }

    pub(super) fn commit_failed_recovery_degraded(&mut self, last_event_id: Option<u64>) {
        if let Some(last_event_id) = last_event_id {
            self.cursor = self.cursor.advance_to(last_event_id);
        }
        self.recovery_inflight = false;
        self.max_seen_live_id = None;
    }

    fn observe_and_buffer(&mut self, event: RouterEvent) {
        let event_id = event.event.event_id;
        self.max_seen_live_id = Some(self.max_seen_live_id.unwrap_or(0).max(event_id));
        self.pending_live_events.push_back(event);
        while self.pending_live_events.len() > RECOVERY_PENDING_LIVE_EVENT_LIMIT {
            self.pending_live_events.pop_front();
        }
    }

    fn clear_max_seen_if_caught_up(&mut self, last_admitted_id: u64) {
        if self
            .max_seen_live_id
            .is_some_and(|max_seen| max_seen <= last_admitted_id)
        {
            self.max_seen_live_id = None;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use dynamo_kv_router::protocols::{
        ExternalSequenceBlockHash, KvCacheEvent, KvCacheStoreData, KvCacheStoredBlockData,
        LocalBlockHash,
    };

    fn store(event_id: u64) -> RouterEvent {
        RouterEvent::new(
            1,
            KvCacheEvent {
                event_id,
                data: KvCacheEventData::Stored(KvCacheStoreData {
                    parent_hash: None,
                    start_position: None,
                    blocks: vec![KvCacheStoredBlockData {
                        block_hash: ExternalSequenceBlockHash(event_id),
                        tokens_hash: LocalBlockHash(event_id),
                        mm_extra_info: None,
                    }],
                }),
                dp_rank: 0,
            },
        )
    }

    #[test]
    fn live_only_source_accepts_first_event_without_recovery() {
        let mut state = RankState::default();
        let action = state.observe_live_event(store(9), false);
        assert!(matches!(action, LiveEventAction::Apply { event_id: 9, .. }));
        assert_eq!(state.last_admitted_id(), None);
        state.commit_live_admission(9);
        assert_eq!(state.last_admitted_id(), Some(9));
    }

    #[test]
    fn recoverable_source_buffers_until_restore() {
        let mut state = RankState::default();
        assert!(matches!(
            state.observe_live_event(store(9), true),
            LiveEventAction::Recover {
                start_event_id: None,
                end_event_id: None,
            }
        ));
        assert!(state.recovery_inflight);
    }

    #[test]
    fn gap_recovery_buffers_and_drains_live_events_in_event_id_order() {
        let mut state = RankState::default();
        assert!(matches!(
            state.observe_live_event(store(1), false),
            LiveEventAction::Apply { event_id: 1, .. }
        ));
        state.commit_live_admission(1);

        assert!(matches!(
            state.observe_live_event(store(4), true),
            LiveEventAction::Recover {
                start_event_id: Some(2),
                end_event_id: None,
            }
        ));
        assert!(matches!(
            state.observe_live_event(store(3), true),
            LiveEventAction::Ignore
        ));
        assert_eq!(state.last_admitted_id(), Some(1));

        let tail = state.drain_advisory_tail_after(2);
        assert_eq!(
            tail.iter()
                .map(|event| event.event.event_id)
                .collect::<Vec<_>>(),
            vec![3, 4]
        );
        assert_eq!(state.last_admitted_id(), Some(4));
        assert!(!state.recovery_inflight);
    }

    #[test]
    fn advisory_recovery_tail_applies_snapshot_before_ordered_suffix() {
        let mut state = RankState::default();
        state.buffer_recovery_tail(store(5));
        state.buffer_recovery_tail(store(3));
        state.buffer_recovery_tail(store(4));
        state.buffer_recovery_tail(store(4));

        let tail = state.drain_advisory_tail_after(3);
        assert_eq!(
            tail.iter()
                .map(|event| event.event.event_id)
                .collect::<Vec<_>>(),
            vec![4, 5]
        );
        assert_eq!(state.last_admitted_id(), Some(5));
        assert!(!state.recovery_inflight);
    }

    #[test]
    fn only_all_domain_clear_supersedes_same_rank_gap_recovery() {
        let mut state = RankState::default();
        assert!(matches!(
            state.observe_live_event(store(1), false),
            LiveEventAction::Apply { event_id: 1, .. }
        ));
        state.commit_live_admission(1);
        assert!(matches!(
            state.observe_live_event(store(4), true),
            LiveEventAction::Recover {
                start_event_id: Some(2),
                ..
            }
        ));

        let mut worker_clear = store(5);
        worker_clear.event.data = KvCacheEventData::Cleared;
        assert!(matches!(
            state.observe_live_event(worker_clear, true),
            LiveEventAction::Ignore
        ));

        let mut clear = store(6);
        clear.event.data = KvCacheEventData::Cleared;
        clear.residency_domain = Default::default();
        assert!(matches!(
            state.observe_live_event(clear, true),
            LiveEventAction::Clear { event_id: 6, .. }
        ));
        state.discard_recovery_before_clear();
        state.commit_live_admission(6);

        assert!(!state.recovery_inflight);
        assert!(matches!(
            state.observe_live_event(store(7), true),
            LiveEventAction::Apply { event_id: 7, .. }
        ));
    }
}