rho-coding-agent 0.26.0

use crate::model::{
    context::{estimate_context_tokens, estimate_message_tokens},
    ContextUsage, Message, ModelUsage,
};
use crate::tool::ToolSpec;

#[derive(Debug, Default)]
pub(super) struct ContextTracker {
    configured_context_window: Option<u64>,
    last_context_window: Option<u64>,
    reported_anchor: Option<ReportedAnchor>,
    unknown_after_compaction: bool,
    cached_message_tokens: u64,
    cached_message_count: usize,
}

/// The local token estimate for one exact message and tool snapshot.
///
/// Keeping this value explicit lets compaction checks, context events, and
/// provider-usage anchoring share one full-history estimation pass.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) struct RequestContextEstimate {
    tokens: u64,
}

/// Pairs a provider-reported input total with the local estimate for the same
/// request, so later estimates can be corrected by the observed difference.
#[derive(Clone, Copy, Debug)]
struct ReportedAnchor {
    reported_tokens: u64,
    estimated_tokens: u64,
}

impl ContextTracker {
    pub(super) fn set_configured_window(&mut self, context_window: Option<u64>) {
        self.configured_context_window = context_window.filter(|window| *window > 0);
    }

    pub(super) fn replace_provider(&mut self) {
        self.last_context_window = None;
        self.reported_anchor = None;
        self.unknown_after_compaction = false;
    }

    pub(super) fn reset(&mut self) {
        self.cached_message_tokens = 0;
        self.cached_message_count = 0;
        self.last_context_window = None;
        self.reported_anchor = None;
        self.unknown_after_compaction = false;
    }

    pub(super) fn history_replaced(&mut self) {
        self.cached_message_tokens = 0;
        self.cached_message_count = 0;
        self.reported_anchor = None;
        self.unknown_after_compaction = false;
    }

    pub(super) fn message_appended(&mut self, message: &Message) {
        self.cached_message_tokens = self
            .cached_message_tokens
            .saturating_add(estimate_message_tokens(message));
        self.cached_message_count += 1;
    }

    pub(super) fn estimate_request(
        &mut self,
        messages: &[Message],
        specs: &[ToolSpec],
    ) -> RequestContextEstimate {
        if self.cached_message_count > messages.len() {
            self.cached_message_tokens = 0;
            self.cached_message_count = 0;
        }
        for message in &messages[self.cached_message_count..] {
            self.message_appended(message);
        }
        let empty_message_tokens = estimate_context_tokens(&[], specs);
        RequestContextEstimate {
            tokens: empty_message_tokens.saturating_add(self.cached_message_tokens),
        }
    }

    pub(super) fn before_provider_request(
        &self,
        estimate: RequestContextEstimate,
    ) -> Option<ContextUsage> {
        if self.unknown_after_compaction {
            None
        } else {
            Some(ContextUsage::estimated(
                estimate.tokens,
                self.context_window(),
            ))
        }
    }

    pub(super) fn estimate_for_compaction(&self, estimate: RequestContextEstimate) -> ContextUsage {
        let tokens = if let Some(anchor) = self.reported_anchor {
            anchor
                .reported_tokens
                .saturating_add(estimate.tokens.saturating_sub(anchor.estimated_tokens))
        } else {
            estimate.tokens
        };
        ContextUsage::estimated(tokens, self.context_window())
    }

    pub(super) fn record_provider_usage(
        &mut self,
        usage: &ModelUsage,
        estimate: RequestContextEstimate,
    ) -> Option<ContextUsage> {
        if let Some(context_window) = usage.context_window {
            self.last_context_window = Some(context_window);
        }
        if let Some(reported_tokens) = usage.total_input_tokens() {
            self.reported_anchor = Some(ReportedAnchor {
                reported_tokens,
                estimated_tokens: estimate.tokens,
            });
        }
        let mut context_usage = ContextUsage::from_model_usage(usage)?;
        context_usage.context_window = self.context_window();
        self.unknown_after_compaction = false;
        Some(context_usage)
    }

    pub(super) fn record_compaction(&mut self) -> ContextUsage {
        self.reported_anchor = None;
        self.unknown_after_compaction = true;
        ContextUsage::unknown_after_compaction(self.context_window())
    }

    pub(super) fn context_window(&self) -> Option<u64> {
        match (self.last_context_window, self.configured_context_window) {
            (Some(reported), Some(configured)) => Some(reported.min(configured)),
            (Some(reported), None) => Some(reported),
            (None, Some(configured)) => Some(configured),
            (None, None) => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::{ContentBlock, ContextUsageSource, Message, ModelUsage};

    fn estimate(tracker: &mut ContextTracker, messages: &[Message]) -> RequestContextEstimate {
        tracker.estimate_request(messages, &[])
    }

    #[test]
    fn one_request_estimate_drives_context_event_and_compaction_check() {
        let mut tracker = ContextTracker::default();
        tracker.set_configured_window(Some(4_000));
        let estimate = estimate(&mut tracker, &[Message::user_text("hello")]);

        let before_request = tracker.before_provider_request(estimate).unwrap();
        let for_compaction = tracker.estimate_for_compaction(estimate);

        assert_eq!(before_request, for_compaction);
        assert_eq!(before_request.source, ContextUsageSource::Estimated);
        assert_eq!(before_request.context_window, Some(4_000));
        assert!(before_request.tokens.unwrap() > 0);
    }

    #[test]
    fn provider_reported_usage_keeps_safer_configured_window() {
        let mut tracker = ContextTracker::default();
        tracker.set_configured_window(Some(4_000));
        let estimate = estimate(&mut tracker, &[]);

        let usage = tracker
            .record_provider_usage(
                &ModelUsage {
                    input_tokens: Some(100),
                    cache_read_tokens: Some(50),
                    context_window: Some(8_000),
                    ..ModelUsage::default()
                },
                estimate,
            )
            .unwrap();

        assert_eq!(usage.source, ContextUsageSource::ProviderReported);
        assert_eq!(usage.tokens, Some(150));
        assert_eq!(usage.context_window, Some(4_000));
        assert_eq!(tracker.context_window(), Some(4_000));
    }

    #[test]
    fn provider_reported_usage_can_lower_configured_window() {
        let mut tracker = ContextTracker::default();
        tracker.set_configured_window(Some(10_000));
        let estimate = estimate(&mut tracker, &[]);

        let usage = tracker
            .record_provider_usage(
                &ModelUsage {
                    input_tokens: Some(100),
                    context_window: Some(8_000),
                    ..ModelUsage::default()
                },
                estimate,
            )
            .unwrap();

        assert_eq!(usage.context_window, Some(8_000));
        assert_eq!(tracker.context_window(), Some(8_000));
    }

    #[test]
    fn unknown_after_compaction_is_not_overwritten_by_later_estimate() {
        let mut tracker = ContextTracker::default();
        tracker.set_configured_window(Some(4_000));

        let usage = tracker.record_compaction();

        assert_eq!(usage.source, ContextUsageSource::UnknownAfterCompaction);
        assert_eq!(usage.context_window, Some(4_000));
        let estimate = estimate(&mut tracker, &[Message::user_text("after")]);
        assert_eq!(tracker.before_provider_request(estimate), None);
    }

    #[test]
    fn provider_usage_clears_unknown_after_compaction() {
        let mut tracker = ContextTracker::default();
        tracker.set_configured_window(Some(4_000));
        tracker.record_compaction();
        let messages = [Message::User(vec![ContentBlock::Text("after".into())])];
        let estimate = estimate(&mut tracker, &messages);

        tracker.record_provider_usage(
            &ModelUsage {
                input_tokens: Some(10),
                output_tokens: Some(2),
                ..ModelUsage::default()
            },
            estimate,
        );

        assert!(tracker.before_provider_request(estimate).is_some());
    }

    #[test]
    fn compaction_estimate_anchors_to_provider_reported_usage() {
        let mut tracker = ContextTracker::default();
        tracker.set_configured_window(Some(10_000));
        let messages = vec![Message::user_text("hello")];
        let estimate_at_report = estimate(&mut tracker, &messages);
        let estimated_tokens = tracker
            .estimate_for_compaction(estimate_at_report)
            .tokens
            .unwrap();

        tracker.record_provider_usage(
            &ModelUsage {
                input_tokens: Some(estimated_tokens + 500),
                ..ModelUsage::default()
            },
            estimate_at_report,
        );

        let mut grown = messages.clone();
        grown.push(Message::assistant_text("a longer assistant reply"));
        let grown_estimate = estimate(&mut tracker, &grown);
        let anchored = tracker
            .estimate_for_compaction(grown_estimate)
            .tokens
            .unwrap();
        assert_eq!(
            anchored,
            estimated_tokens + 500 + (grown_estimate.tokens - estimated_tokens)
        );
    }

    #[test]
    fn compaction_clears_reported_anchor() {
        let mut tracker = ContextTracker::default();
        tracker.set_configured_window(Some(10_000));
        let messages = vec![Message::user_text("hello")];
        let estimate = estimate(&mut tracker, &messages);
        let local_tokens = estimate.tokens;
        tracker.record_provider_usage(
            &ModelUsage {
                input_tokens: Some(local_tokens + 500),
                ..ModelUsage::default()
            },
            estimate,
        );

        tracker.record_compaction();

        assert_eq!(
            tracker.estimate_for_compaction(estimate).tokens,
            Some(local_tokens)
        );
    }

    #[test]
    fn appended_message_cache_matches_fresh_estimate() {
        let specs = Vec::new();
        let mut messages = vec![Message::user_text("first")];
        let mut tracker = ContextTracker::default();
        tracker.estimate_request(&messages, &specs);

        messages.push(Message::assistant_text("second response"));

        assert_eq!(
            tracker.estimate_request(&messages, &specs).tokens,
            estimate_context_tokens(&messages, &specs)
        );
    }

    #[test]
    fn reset_clears_provider_window_and_unknown_state() {
        let mut tracker = ContextTracker::default();
        tracker.set_configured_window(Some(4_000));
        let initial_estimate = estimate(&mut tracker, &[]);
        tracker.record_provider_usage(
            &ModelUsage {
                input_tokens: Some(10),
                context_window: Some(8_000),
                ..ModelUsage::default()
            },
            initial_estimate,
        );
        tracker.record_compaction();

        tracker.reset();

        assert_eq!(tracker.context_window(), Some(4_000));
        let estimate = estimate(&mut tracker, &[Message::user_text("after")]);
        assert!(tracker.before_provider_request(estimate).is_some());
    }
}