rho-coding-agent 0.7.0

A lightweight agent harness inspired by Pi
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161

use crate::model::{estimate_context_usage, ContextUsage, Message, ModelUsage};
use crate::tool::ToolSpec;

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

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.unknown_after_compaction = false;
    }

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

    pub(super) fn before_provider_request(
        &self,
        messages: &[Message],
        specs: &[ToolSpec],
    ) -> Option<ContextUsage> {
        if self.unknown_after_compaction {
            None
        } else {
            Some(estimate_context_usage(
                messages,
                specs,
                self.context_window(),
            ))
        }
    }

    pub(super) fn estimate_for_compaction(
        &self,
        messages: &[Message],
        specs: &[ToolSpec],
    ) -> ContextUsage {
        estimate_context_usage(messages, specs, self.context_window())
    }

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

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

    pub(super) fn context_window(&self) -> Option<u64> {
        self.last_context_window.or(self.configured_context_window)
    }
}

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

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

        let usage = tracker
            .before_provider_request(&[Message::user_text("hello")], &[])
            .unwrap();

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

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

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

        assert_eq!(usage.source, ContextUsageSource::ProviderReported);
        assert_eq!(usage.tokens, Some(150));
        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));
        assert_eq!(
            tracker.before_provider_request(&[Message::user_text("after")], &[]),
            None
        );
    }

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

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

        assert!(tracker
            .before_provider_request(
                &[Message::User(vec![ContentBlock::Text("after".into())])],
                &[]
            )
            .is_some());
    }

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

        tracker.reset();

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