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 folded_text = text.to_lowercase();
80        let tokens = slice
81            .iter()
82            .map(|message| {
83                message
84                    .token_count
85                    .unwrap_or_else(|| engine.count_message(message))
86            })
87            .sum::<u32>();
88        let goal_overlap = if goal_terms.is_empty() {
89            0
90        } else {
91            overlap_count(&terms(text), &goal_terms)
92        };
93        let has_unresolved = has_unresolved(slice, &folded_text);
94        let referenced_later = unit_referenced_later(slice, text, &unit_texts[index + 1..]);
95        let is_error_or_decision = is_error_or_decision(slice, &folded_text);
96        let dependency = context
97            .preserved_refs
98            .iter()
99            .any(|reference| contains_folded(text, reference))
100            || context
101                .active_directives
102                .iter()
103                .any(|directive| directive_dependency(text, directive));
104        let mandatory = index >= recent_start || has_unresolved || dependency;
105        let recency = ((index as u64 + 1) * 1_000 / u64::from(unit_count)) as u32;
106        let token_cost = (u64::from(tokens) * 1_000 / u64::from(denominator)) as u32;
107        let prefix_invalidation_cost =
108            ((ranges.len() - index) as u64 * 1_000 / u64::from(unit_count)) as u32;
109        let utility = i64::from(goal_overlap) * 4_000
110            + if has_unresolved { 20_000 } else { 0 }
111            + if referenced_later { 5_000 } else { 0 }
112            + if is_error_or_decision { 6_000 } else { 0 }
113            + i64::from(recency) * 2
114            - i64::from(token_cost) * 2
115            - i64::from(prefix_invalidation_cost);
116        scores.push(UtilityUnitScore {
117            range: range.clone(),
118            tokens,
119            mandatory,
120            goal_overlap,
121            has_unresolved,
122            referenced_later,
123            is_error_or_decision,
124            recency,
125            token_cost,
126            prefix_invalidation_cost,
127            utility,
128        });
129    }
130
131    if total_tokens <= target_tokens {
132        return UtilityArchivePlan {
133            archived_ranges: Vec::new(),
134            retained_ranges: ranges,
135            archived_tokens: 0,
136            retained_tokens: scores.iter().map(|score| score.tokens).sum(),
137            scores,
138        };
139    }
140
141    let mut retained = scores
142        .iter()
143        .enumerate()
144        .filter_map(|(index, score)| score.mandatory.then_some(index))
145        .collect::<BTreeSet<_>>();
146    let mut retained_tokens = retained
147        .iter()
148        .map(|index| scores[*index].tokens)
149        .sum::<u32>();
150    let mut optional = scores
151        .iter()
152        .enumerate()
153        .filter_map(|(index, score)| (!score.mandatory).then_some(index))
154        .collect::<Vec<_>>();
155    optional.sort_by(|left, right| compare_density(&scores[*right], &scores[*left]));
156    for index in optional {
157        let tokens = scores[index].tokens;
158        if retained_tokens.saturating_add(tokens) <= target_tokens {
159            retained.insert(index);
160            retained_tokens = retained_tokens.saturating_add(tokens);
161        }
162    }
163
164    let retained_ranges = ranges
165        .iter()
166        .enumerate()
167        .filter_map(|(index, range)| retained.contains(&index).then_some(range.clone()))
168        .collect::<Vec<_>>();
169    let archived_ranges = ranges
170        .iter()
171        .enumerate()
172        .filter_map(|(index, range)| (!retained.contains(&index)).then_some(range.clone()))
173        .collect::<Vec<_>>();
174    let archived_tokens = scores
175        .iter()
176        .enumerate()
177        .filter_map(|(index, score)| (!retained.contains(&index)).then_some(score.tokens))
178        .sum();
179    UtilityArchivePlan {
180        archived_ranges,
181        retained_ranges,
182        archived_tokens,
183        retained_tokens,
184        scores,
185    }
186}
187
188fn compare_density(left: &UtilityUnitScore, right: &UtilityUnitScore) -> Ordering {
189    let left_density = i128::from(left.utility) * i128::from(right.tokens.max(1));
190    let right_density = i128::from(right.utility) * i128::from(left.tokens.max(1));
191    left_density
192        .cmp(&right_density)
193        .then_with(|| left.utility.cmp(&right.utility))
194        .then_with(|| left.range.start.cmp(&right.range.start))
195}
196
197fn unit_text(messages: &[Message]) -> String {
198    let mut text = String::new();
199    let mut first_part = true;
200    for message in messages {
201        match &message.content {
202            Content::Text(content) => append_unit_part(&mut text, &mut first_part, content),
203            Content::Parts(content_parts) => {
204                for part in content_parts {
205                    match part {
206                        ContentPart::Text { text: content } => {
207                            append_unit_part(&mut text, &mut first_part, content)
208                        }
209                        ContentPart::ToolResult {
210                            call_id, output, ..
211                        } => {
212                            append_unit_part(&mut text, &mut first_part, call_id.as_str());
213                            text.push(' ');
214                            text.push_str(output);
215                        }
216                        ContentPart::Image { url, .. } => append_unit_part(
217                            &mut text,
218                            &mut first_part,
219                            url.as_deref().unwrap_or_default(),
220                        ),
221                        ContentPart::Audio { .. } => {
222                            append_unit_part(&mut text, &mut first_part, "audio")
223                        }
224                    }
225                }
226            }
227        }
228        for call in &message.tool_calls {
229            append_unit_part(&mut text, &mut first_part, call.id.as_str());
230            text.push(' ');
231            text.push_str(call.name.as_str());
232            text.push(' ');
233            text.push_str(&call.arguments.to_string());
234        }
235    }
236    text
237}
238
239fn append_unit_part(text: &mut String, first_part: &mut bool, part: &str) {
240    if !*first_part {
241        text.push('\n');
242    }
243    *first_part = false;
244    text.push_str(part);
245}
246
247fn contains_folded(text: &str, pattern: &str) -> bool {
248    !pattern.trim().is_empty() && text.to_lowercase().contains(&pattern.to_lowercase())
249}
250
251fn directive_dependency(text: &str, directive: &str) -> bool {
252    if contains_folded(text, directive) {
253        return true;
254    }
255    let directive_terms = terms(directive);
256    if directive_terms.is_empty() {
257        return false;
258    }
259    let threshold = directive_terms.len().min(2);
260    terms(text).intersection(&directive_terms).count() >= threshold
261}
262
263fn has_unresolved(messages: &[Message], folded_text: &str) -> bool {
264    let mut opened = BTreeSet::new();
265    let mut resolved = BTreeSet::new();
266    for message in messages {
267        for call in &message.tool_calls {
268            opened.insert(call.id.to_string());
269        }
270        if let Content::Parts(parts) = &message.content {
271            for part in parts {
272                if let ContentPart::ToolResult {
273                    call_id, is_error, ..
274                } = part
275                {
276                    if *is_error {
277                        return true;
278                    }
279                    resolved.insert(call_id.to_string());
280                }
281            }
282        }
283    }
284    opened.iter().any(|call_id| !resolved.contains(call_id))
285        || marker_folded(
286            folded_text,
287            &[
288                "unresolved",
289                "open question",
290                "retry",
291                "blocked",
292                "待确认",
293                "未解决",
294                "重试",
295                "阻塞",
296            ],
297        )
298}
299
300fn is_error_or_decision(messages: &[Message], folded_text: &str) -> bool {
301    messages.iter().any(|message| {
302        matches!(&message.content, Content::Parts(parts) if parts.iter().any(|part| matches!(part, ContentPart::ToolResult { is_error: true, .. })))
303    }) || marker_folded(
304        folded_text,
305        &[
306            "error", "failed", "failure", "exception", "decision", "decided", "must", "should",
307            "错误", "失败", "异常", "决定", "选择", "必须", "应当",
308        ],
309    )
310}
311
312fn marker_folded(folded_text: &str, markers: &[&str]) -> bool {
313    markers.iter().any(|marker| folded_text.contains(marker))
314}
315
316fn unit_referenced_later(messages: &[Message], text: &str, later: &[String]) -> bool {
317    let mut references = messages
318        .iter()
319        .flat_map(|message| message.tool_calls.iter().map(|call| call.id.to_string()))
320        .collect::<BTreeSet<_>>();
321    references.extend(
322        text.split_whitespace()
323            .map(|token| token.trim_matches(|character: char| character.is_ascii_punctuation()))
324            .filter(|token| token.contains('/') || token.contains("://"))
325            .filter(|token| token.len() > 3)
326            .map(str::to_string),
327    );
328    references.iter().any(|reference| {
329        later
330            .iter()
331            .any(|later_text| contains_folded(later_text, reference))
332    })
333}
334
335#[cfg(test)]
336mod tests {
337    use super::*;
338    use crate::types::message::{ContentPart, ToolCall};
339
340    #[test]
341    fn unit_text_preserves_empty_part_separators() {
342        let messages = vec![Message::user(""), Message::user("next")];
343        assert_eq!(unit_text(&messages), "\nnext");
344    }
345
346    #[test]
347    fn unresolved_tool_unit_is_mandatory() {
348        let mut call = Message::assistant("working");
349        call.tool_calls.push(ToolCall {
350            id: "call-1".into(),
351            name: "read".into(),
352            arguments: serde_json::json!({"path": "/work/a"}),
353        });
354        call.token_count = Some(20);
355        let mut recent = Message::user("recent");
356        recent.token_count = Some(20);
357        let messages = vec![call, recent];
358        let plan = plan_utility_archive(
359            &messages,
360            40,
361            20,
362            1,
363            &ContextTokenEngine::char_approx(),
364            &UtilitySelectionContext {
365                goal: "",
366                criteria: &[],
367                preserved_refs: &[],
368                active_directives: &[],
369            },
370        );
371        assert!(plan.scores[0].mandatory);
372        assert!(plan.scores[0].has_unresolved);
373        assert_eq!(plan.retained_tokens, 40);
374    }
375
376    #[test]
377    fn chinese_directive_dependency_requires_bigram_overlap_not_shared_characters() {
378        // Under the old per-character CJK vocabulary, any Chinese unit sharing two
379        // common characters (中/文/回…) with an active directive was marked mandatory,
380        // so compression could never archive unrelated Chinese history.
381        let mut unrelated = Message::assistant("我们在文中回顾了天气");
382        unrelated.token_count = Some(30);
383        let mut on_topic = Message::user("已按要求保持中文回答");
384        on_topic.token_count = Some(30);
385        let mut recent = Message::user("recent");
386        recent.token_count = Some(10);
387        let messages = vec![unrelated, on_topic, recent];
388        let plan = plan_utility_archive(
389            &messages,
390            70,
391            10,
392            1,
393            &ContextTokenEngine::char_approx(),
394            &UtilitySelectionContext {
395                goal: "",
396                criteria: &[],
397                preserved_refs: &[],
398                active_directives: &["必须用中文回答".into()],
399            },
400        );
401        assert!(
402            !plan.scores[0].mandatory,
403            "unrelated Chinese text must not bind to the directive"
404        );
405        assert!(
406            plan.scores[1].mandatory,
407            "text restating the directive must stay mandatory"
408        );
409    }
410
411    #[test]
412    fn preserved_ref_keeps_complete_tool_unit() {
413        let mut call = Message::assistant("read artifact");
414        call.tool_calls.push(ToolCall {
415            id: "call-keep".into(),
416            name: "read".into(),
417            arguments: serde_json::json!({}),
418        });
419        call.token_count = Some(20);
420        let mut result = Message::tool(vec![ContentPart::ToolResult {
421            call_id: "call-keep".into(),
422            output: "artifact".into(),
423            is_error: false,
424            durable_content: None,
425        }]);
426        result.token_count = Some(20);
427        let messages = vec![call, result];
428        let plan = plan_utility_archive(
429            &messages,
430            40,
431            0,
432            0,
433            &ContextTokenEngine::char_approx(),
434            &UtilitySelectionContext {
435                goal: "",
436                criteria: &[],
437                preserved_refs: &["call-keep".into()],
438                active_directives: &[],
439            },
440        );
441        assert!(plan.scores[0].mandatory);
442        assert_eq!(plan.archived_ranges, Vec::<Range<usize>>::new());
443    }
444}
deepstrike_core/context/utility.rs

deepstrike_core/context/
utility.rs
