deepstrike_core/context/
utility.rs

1//! Deterministic value-aware selection over indivisible context units.
2
3use std::cmp::Ordering;
4use std::collections::BTreeSet;
5use std::ops::Range;
6
7use super::token_engine::ContextTokenEngine;
8use super::units::unit_boundaries;
9use crate::lexical::{overlap_count, terms};
10use crate::types::message::{Content, ContentPart, Message};
11
12pub struct UtilitySelectionContext<'a> {
13    pub goal: &'a str,
14    pub criteria: &'a [String],
15    pub preserved_refs: &'a [String],
16    pub active_directives: &'a [String],
17}
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct UtilityUnitScore {
21    pub range: Range<usize>,
22    pub tokens: u32,
23    pub mandatory: bool,
24    pub goal_overlap: u32,
25    pub has_unresolved: bool,
26    pub referenced_later: bool,
27    pub is_error_or_decision: bool,
28    pub recency: u32,
29    pub token_cost: u32,
30    pub prefix_invalidation_cost: u32,
31    pub utility: i64,
32}
33
34#[derive(Debug, Clone, Default, PartialEq, Eq)]
35pub struct UtilityArchivePlan {
36    pub archived_ranges: Vec<Range<usize>>,
37    pub retained_ranges: Vec<Range<usize>>,
38    pub archived_tokens: u32,
39    pub retained_tokens: u32,
40    pub scores: Vec<UtilityUnitScore>,
41}
42
43/// Select complete units to retain under `target_tokens`.
44///
45/// Mandatory dependencies are retained even when they alone exceed the target;
46/// callers can then escalate pressure honestly instead of silently deleting the
47/// evidence required to continue the task.
48pub fn plan_utility_archive(
49    messages: &[Message],
50    total_tokens: u32,
51    target_tokens: u32,
52    preserve_recent_units: usize,
53    engine: &ContextTokenEngine,
54    context: &UtilitySelectionContext<'_>,
55) -> UtilityArchivePlan {
56    let ranges = unit_boundaries(messages);
57    if ranges.is_empty() {
58        return UtilityArchivePlan::default();
59    }
60    let unit_texts = ranges
61        .iter()
62        .map(|range| unit_text(&messages[range.clone()]))
63        .collect::<Vec<_>>();
64    let goal_terms = terms(
65        std::iter::once(context.goal)
66            .chain(context.criteria.iter().map(String::as_str))
67            .collect::<Vec<_>>()
68            .join(" ")
69            .as_str(),
70    );
71    let recent_start = ranges.len().saturating_sub(preserve_recent_units);
72    let denominator = total_tokens.max(1);
73    let unit_count = ranges.len().max(1) as u32;
74    let mut scores = Vec::with_capacity(ranges.len());
75
76    for (index, range) in ranges.iter().enumerate() {
77        let slice = &messages[range.clone()];
78        let text = &unit_texts[index];
79        let tokens = slice
80            .iter()
81            .map(|message| {
82                message
83                    .token_count
84                    .unwrap_or_else(|| engine.count_message(message))
85            })
86            .sum::<u32>();
87        let goal_overlap = overlap_count(&terms(text), &goal_terms);
88        let has_unresolved = has_unresolved(slice, text);
89        let referenced_later = unit_referenced_later(slice, text, &unit_texts[index + 1..]);
90        let is_error_or_decision = is_error_or_decision(slice, text);
91        let dependency = context
92            .preserved_refs
93            .iter()
94            .any(|reference| contains_folded(text, reference))
95            || context
96                .active_directives
97                .iter()
98                .any(|directive| directive_dependency(text, directive));
99        let mandatory = index >= recent_start || has_unresolved || dependency;
100        let recency = ((index as u64 + 1) * 1_000 / u64::from(unit_count)) as u32;
101        let token_cost = (u64::from(tokens) * 1_000 / u64::from(denominator)) as u32;
102        let prefix_invalidation_cost =
103            ((ranges.len() - index) as u64 * 1_000 / u64::from(unit_count)) as u32;
104        let utility = i64::from(goal_overlap) * 4_000
105            + if has_unresolved { 20_000 } else { 0 }
106            + if referenced_later { 5_000 } else { 0 }
107            + if is_error_or_decision { 6_000 } else { 0 }
108            + i64::from(recency) * 2
109            - i64::from(token_cost) * 2
110            - i64::from(prefix_invalidation_cost);
111        scores.push(UtilityUnitScore {
112            range: range.clone(),
113            tokens,
114            mandatory,
115            goal_overlap,
116            has_unresolved,
117            referenced_later,
118            is_error_or_decision,
119            recency,
120            token_cost,
121            prefix_invalidation_cost,
122            utility,
123        });
124    }
125
126    if total_tokens <= target_tokens {
127        return UtilityArchivePlan {
128            archived_ranges: Vec::new(),
129            retained_ranges: ranges,
130            archived_tokens: 0,
131            retained_tokens: scores.iter().map(|score| score.tokens).sum(),
132            scores,
133        };
134    }
135
136    let mut retained = scores
137        .iter()
138        .enumerate()
139        .filter_map(|(index, score)| score.mandatory.then_some(index))
140        .collect::<BTreeSet<_>>();
141    let mut retained_tokens = retained
142        .iter()
143        .map(|index| scores[*index].tokens)
144        .sum::<u32>();
145    let mut optional = scores
146        .iter()
147        .enumerate()
148        .filter_map(|(index, score)| (!score.mandatory).then_some(index))
149        .collect::<Vec<_>>();
150    optional.sort_by(|left, right| compare_density(&scores[*right], &scores[*left]));
151    for index in optional {
152        let tokens = scores[index].tokens;
153        if retained_tokens.saturating_add(tokens) <= target_tokens {
154            retained.insert(index);
155            retained_tokens = retained_tokens.saturating_add(tokens);
156        }
157    }
158
159    let retained_ranges = ranges
160        .iter()
161        .enumerate()
162        .filter_map(|(index, range)| retained.contains(&index).then_some(range.clone()))
163        .collect::<Vec<_>>();
164    let archived_ranges = ranges
165        .iter()
166        .enumerate()
167        .filter_map(|(index, range)| (!retained.contains(&index)).then_some(range.clone()))
168        .collect::<Vec<_>>();
169    let archived_tokens = scores
170        .iter()
171        .enumerate()
172        .filter_map(|(index, score)| (!retained.contains(&index)).then_some(score.tokens))
173        .sum();
174    UtilityArchivePlan {
175        archived_ranges,
176        retained_ranges,
177        archived_tokens,
178        retained_tokens,
179        scores,
180    }
181}
182
183fn compare_density(left: &UtilityUnitScore, right: &UtilityUnitScore) -> Ordering {
184    let left_density = i128::from(left.utility) * i128::from(right.tokens.max(1));
185    let right_density = i128::from(right.utility) * i128::from(left.tokens.max(1));
186    left_density
187        .cmp(&right_density)
188        .then_with(|| left.utility.cmp(&right.utility))
189        .then_with(|| left.range.start.cmp(&right.range.start))
190}
191
192fn unit_text(messages: &[Message]) -> String {
193    let mut parts = Vec::new();
194    for message in messages {
195        match &message.content {
196            Content::Text(text) => parts.push(text.clone()),
197            Content::Parts(content_parts) => {
198                for part in content_parts {
199                    match part {
200                        ContentPart::Text { text } => parts.push(text.clone()),
201                        ContentPart::ToolResult {
202                            call_id, output, ..
203                        } => parts.push(format!("{call_id} {output}")),
204                        ContentPart::Image { url, .. } => {
205                            parts.push(url.clone().unwrap_or_default())
206                        }
207                        ContentPart::Audio { .. } => parts.push("audio".into()),
208                    }
209                }
210            }
211        }
212        for call in &message.tool_calls {
213            parts.push(format!("{} {} {}", call.id, call.name, call.arguments));
214        }
215    }
216    parts.join("\n")
217}
218
219fn contains_folded(text: &str, pattern: &str) -> bool {
220    !pattern.trim().is_empty() && text.to_lowercase().contains(&pattern.to_lowercase())
221}
222
223fn directive_dependency(text: &str, directive: &str) -> bool {
224    if contains_folded(text, directive) {
225        return true;
226    }
227    let directive_terms = terms(directive);
228    if directive_terms.is_empty() {
229        return false;
230    }
231    let threshold = directive_terms.len().min(2);
232    terms(text).intersection(&directive_terms).count() >= threshold
233}
234
235fn has_unresolved(messages: &[Message], text: &str) -> bool {
236    let mut opened = BTreeSet::new();
237    let mut resolved = BTreeSet::new();
238    for message in messages {
239        for call in &message.tool_calls {
240            opened.insert(call.id.to_string());
241        }
242        if let Content::Parts(parts) = &message.content {
243            for part in parts {
244                if let ContentPart::ToolResult {
245                    call_id, is_error, ..
246                } = part
247                {
248                    if *is_error {
249                        return true;
250                    }
251                    resolved.insert(call_id.to_string());
252                }
253            }
254        }
255    }
256    opened.iter().any(|call_id| !resolved.contains(call_id))
257        || marker(
258            text,
259            &[
260                "unresolved",
261                "open question",
262                "retry",
263                "blocked",
264                "待确认",
265                "未解决",
266                "重试",
267                "阻塞",
268            ],
269        )
270}
271
272fn is_error_or_decision(messages: &[Message], text: &str) -> bool {
273    messages.iter().any(|message| {
274        matches!(&message.content, Content::Parts(parts) if parts.iter().any(|part| matches!(part, ContentPart::ToolResult { is_error: true, .. })))
275    }) || marker(
276        text,
277        &[
278            "error", "failed", "failure", "exception", "decision", "decided", "must", "should",
279            "错误", "失败", "异常", "决定", "选择", "必须", "应当",
280        ],
281    )
282}
283
284fn marker(text: &str, markers: &[&str]) -> bool {
285    markers.iter().any(|marker| contains_folded(text, marker))
286}
287
288fn unit_referenced_later(messages: &[Message], text: &str, later: &[String]) -> bool {
289    let mut references = messages
290        .iter()
291        .flat_map(|message| message.tool_calls.iter().map(|call| call.id.to_string()))
292        .collect::<BTreeSet<_>>();
293    references.extend(
294        text.split_whitespace()
295            .map(|token| token.trim_matches(|character: char| character.is_ascii_punctuation()))
296            .filter(|token| token.contains('/') || token.contains("://"))
297            .filter(|token| token.len() > 3)
298            .map(str::to_string),
299    );
300    references.iter().any(|reference| {
301        later
302            .iter()
303            .any(|later_text| contains_folded(later_text, reference))
304    })
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use crate::types::message::{ContentPart, ToolCall};
311
312    #[test]
313    fn unresolved_tool_unit_is_mandatory() {
314        let mut call = Message::assistant("working");
315        call.tool_calls.push(ToolCall {
316            id: "call-1".into(),
317            name: "read".into(),
318            arguments: serde_json::json!({"path": "/work/a"}),
319        });
320        call.token_count = Some(20);
321        let mut recent = Message::user("recent");
322        recent.token_count = Some(20);
323        let messages = vec![call, recent];
324        let plan = plan_utility_archive(
325            &messages,
326            40,
327            20,
328            1,
329            &ContextTokenEngine::char_approx(),
330            &UtilitySelectionContext {
331                goal: "",
332                criteria: &[],
333                preserved_refs: &[],
334                active_directives: &[],
335            },
336        );
337        assert!(plan.scores[0].mandatory);
338        assert!(plan.scores[0].has_unresolved);
339        assert_eq!(plan.retained_tokens, 40);
340    }
341
342    #[test]
343    fn chinese_directive_dependency_requires_bigram_overlap_not_shared_characters() {
344        // Under the old per-character CJK vocabulary, any Chinese unit sharing two
345        // common characters (中/文/回…) with an active directive was marked mandatory,
346        // so compression could never archive unrelated Chinese history.
347        let mut unrelated = Message::assistant("我们在文中回顾了天气");
348        unrelated.token_count = Some(30);
349        let mut on_topic = Message::user("已按要求保持中文回答");
350        on_topic.token_count = Some(30);
351        let mut recent = Message::user("recent");
352        recent.token_count = Some(10);
353        let messages = vec![unrelated, on_topic, recent];
354        let plan = plan_utility_archive(
355            &messages,
356            70,
357            10,
358            1,
359            &ContextTokenEngine::char_approx(),
360            &UtilitySelectionContext {
361                goal: "",
362                criteria: &[],
363                preserved_refs: &[],
364                active_directives: &["必须用中文回答".into()],
365            },
366        );
367        assert!(
368            !plan.scores[0].mandatory,
369            "unrelated Chinese text must not bind to the directive"
370        );
371        assert!(
372            plan.scores[1].mandatory,
373            "text restating the directive must stay mandatory"
374        );
375    }
376
377    #[test]
378    fn preserved_ref_keeps_complete_tool_unit() {
379        let mut call = Message::assistant("read artifact");
380        call.tool_calls.push(ToolCall {
381            id: "call-keep".into(),
382            name: "read".into(),
383            arguments: serde_json::json!({}),
384        });
385        call.token_count = Some(20);
386        let mut result = Message::tool(vec![ContentPart::ToolResult {
387            call_id: "call-keep".into(),
388            output: "artifact".into(),
389            is_error: false,
390        }]);
391        result.token_count = Some(20);
392        let messages = vec![call, result];
393        let plan = plan_utility_archive(
394            &messages,
395            40,
396            0,
397            0,
398            &ContextTokenEngine::char_approx(),
399            &UtilitySelectionContext {
400                goal: "",
401                criteria: &[],
402                preserved_refs: &["call-keep".into()],
403                active_directives: &[],
404            },
405        );
406        assert!(plan.scores[0].mandatory);
407        assert_eq!(plan.archived_ranges, Vec::<Range<usize>>::new());
408    }
409}
deepstrike_core/context/utility.rs

deepstrike_core/context/
utility.rs
