clark_agent/
budget.rs

1//! Default token-budget context transform.
2//!
3//! Full provider history is the source of truth. When the conversation
4//! grows past the budget, we drop the oldest tool results' content
5//! (keeping the call shape) so the model still sees
6//! what it tried, but doesn't pay token cost for stale, re-fetchable
7//! payloads.
8//!
9//! Token counts come from the loop's configured
10//! [`crate::tokens::TokenEstimator`] (default: char heuristic; apps may
11//! plug a real tokenizer). Apps that want an entirely different policy
12//! implement their own `ContextTransform` plugin and skip this one.
13
14use async_trait::async_trait;
15
16use crate::plugin::{ContextTransform, Plugin, PluginCapabilities, TransformContext};
17use crate::tokens::TokenEstimator;
18use crate::types::{AgentMessage, TextContent, ToolResultBlock, ToolResultContent};
19
20/// Configurable token budget. Default `60_000` tokens with a 70% trim
21/// trigger.
22#[derive(Debug, Clone)]
23pub struct TokenBudget {
24    pub max_tokens: usize,
25    /// When total estimate exceeds `trim_trigger * max_tokens`, start
26    /// truncating tool results.
27    pub trim_trigger: f32,
28    /// Replacement text inserted in place of truncated tool result content.
29    pub truncation_marker: String,
30}
31
32impl Default for TokenBudget {
33    fn default() -> Self {
34        Self {
35            max_tokens: 60_000,
36            trim_trigger: 0.7,
37            truncation_marker: "[truncated for context budget — re-run tool to refetch]".into(),
38        }
39    }
40}
41
42impl Plugin for TokenBudget {
43    fn name(&self) -> &'static str {
44        "token_budget"
45    }
46    fn capabilities(&self) -> PluginCapabilities {
47        PluginCapabilities::context_transform()
48    }
49}
50
51#[async_trait]
52impl ContextTransform for TokenBudget {
53    async fn transform(
54        &self,
55        mut messages: Vec<AgentMessage>,
56        cx: &TransformContext<'_>,
57    ) -> Vec<AgentMessage> {
58        let trigger = (self.max_tokens as f32 * self.trim_trigger).round() as usize;
59        let total = cx.estimator.estimate_messages(&messages);
60        if total <= trigger {
61            return messages;
62        }
63
64        // Walk oldest-first (skip the very last user/tool exchange so the
65        // current turn keeps full fidelity). Replace tool result content
66        // with the marker until under budget or nothing left to truncate.
67        let last_idx = messages.len().saturating_sub(2);
68        let mut idx = 0;
69        while idx < last_idx {
70            let truncated = if let AgentMessage::ToolResult { content, .. } = &mut messages[idx] {
71                if !content_already_marker(content, &self.truncation_marker) {
72                    *content = ToolResultContent {
73                        blocks: vec![ToolResultBlock::Text(TextContent {
74                            text: self.truncation_marker.clone(),
75                        })],
76                    };
77                    true
78                } else {
79                    false
80                }
81            } else {
82                false
83            };
84            idx += 1;
85            if truncated {
86                let total = cx.estimator.estimate_messages(&messages);
87                if total <= trigger {
88                    break;
89                }
90            }
91        }
92
93        messages
94    }
95}
96
97/// Char-heuristic estimate kept as a free function for callers that
98/// want a one-off count without holding a `TokenEstimator`. Prefer
99/// [`crate::plugin::TransformContext::estimator`] inside transforms.
100pub fn estimate_tokens(message: &AgentMessage) -> usize {
101    crate::tokens::CHAR_HEURISTIC.estimate_message(message)
102}
103
104fn content_already_marker(content: &ToolResultContent, marker: &str) -> bool {
105    content.blocks.len() == 1
106        && matches!(
107            &content.blocks[0],
108            ToolResultBlock::Text(t) if t.text == marker
109        )
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use crate::types::ToolResultBlock;
116    use tokio_util::sync::CancellationToken;
117
118    #[tokio::test]
119    async fn budget_truncates_old_tool_results() {
120        let budget = TokenBudget {
121            max_tokens: 200,
122            trim_trigger: 0.5,
123            truncation_marker: "[trunc]".into(),
124        };
125        let big = "x".repeat(2000);
126        let messages = vec![
127            AgentMessage::User {
128                content: crate::types::UserContent::Text("start".into()),
129                timestamp: None,
130            },
131            AgentMessage::ToolResult {
132                tool_call_id: "1".into(),
133                tool_name: "shell".into(),
134                content: ToolResultContent::text(big.clone()),
135                is_error: false,
136                narration: None,
137                details: None,
138                timestamp: None,
139            },
140            AgentMessage::User {
141                content: crate::types::UserContent::Text("more".into()),
142                timestamp: None,
143            },
144            AgentMessage::ToolResult {
145                tool_call_id: "2".into(),
146                tool_name: "shell".into(),
147                content: ToolResultContent::text(big),
148                is_error: false,
149                narration: None,
150                details: None,
151                timestamp: None,
152            },
153        ];
154        let token = CancellationToken::new();
155        let cx = TransformContext::for_test(&token);
156        let out = budget.transform(messages, &cx).await;
157        // Oldest tool result should be truncated.
158        let AgentMessage::ToolResult { content, .. } = &out[1] else {
159            panic!("expected tool result");
160        };
161        assert!(matches!(&content.blocks[0], ToolResultBlock::Text(t) if t.text == "[trunc]"));
162        // Last tool result preserved.
163        let AgentMessage::ToolResult { content, .. } = &out[3] else {
164            panic!("expected tool result");
165        };
166        assert!(matches!(&content.blocks[0], ToolResultBlock::Text(t) if t.text != "[trunc]"));
167    }
168}
clark_agent/budget.rs

clark_agent/
budget.rs
