deepstrike_core/context/

utility.rs

//! Deterministic value-aware selection over indivisible context units.

use std::cmp::Ordering;
use std::collections::BTreeSet;
use std::ops::Range;

use super::measurement::TokenMeasurement;
use super::token_engine::ContextTokenEngine;
use super::units::unit_boundaries;
use crate::lexical::{overlap_count, terms};
use crate::types::message::{Content, ContentPart, CoreMessage};

pub struct UtilitySelectionContext<'a> {
    pub goal: &'a str,
    pub criteria: &'a [String],
    pub preserved_refs: &'a [String],
    pub active_directives: &'a [String],
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UtilityUnitScore {
    pub range: Range<usize>,
    pub tokens: u32,
    pub mandatory: bool,
    pub goal_overlap: u32,
    pub has_unresolved: bool,
    pub referenced_later: bool,
    pub is_error_or_decision: bool,
    pub recency: u32,
    pub token_cost: u32,
    pub prefix_invalidation_cost: u32,
    pub utility: i64,
}

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct UtilityArchivePlan {
    pub archived_ranges: Vec<Range<usize>>,
    pub retained_ranges: Vec<Range<usize>>,
    pub archived_tokens: u32,
    pub retained_tokens: u32,
    pub scores: Vec<UtilityUnitScore>,
}

/// Select complete units to retain under `target_tokens`.
///
/// Mandatory dependencies are retained even when they alone exceed the target;
/// callers can then escalate pressure honestly instead of silently deleting the
/// evidence required to continue the task.
pub fn plan_utility_archive(
    messages: &[CoreMessage],
    total_tokens: u32,
    target_tokens: u32,
    preserve_recent_units: usize,
    engine: &ContextTokenEngine,
    context: &UtilitySelectionContext<'_>,
) -> UtilityArchivePlan {
    plan_utility_archive_with_measurements(
        messages,
        &[],
        total_tokens,
        target_tokens,
        preserve_recent_units,
        engine,
        context,
    )
}

/// Measurement-aware planner entry point used by context partitions. The parallel slice is
/// host-owned evidence aligned with `messages`; an empty or short slice falls back to deterministic
/// engine recomputation for the missing entries.
pub fn plan_utility_archive_with_measurements(
    messages: &[CoreMessage],
    measurements: &[TokenMeasurement],
    total_tokens: u32,
    target_tokens: u32,
    preserve_recent_units: usize,
    engine: &ContextTokenEngine,
    context: &UtilitySelectionContext<'_>,
) -> UtilityArchivePlan {
    let ranges = unit_boundaries(messages);
    if ranges.is_empty() {
        return UtilityArchivePlan::default();
    }
    let unit_texts = ranges
        .iter()
        .map(|range| unit_text(&messages[range.clone()]))
        .collect::<Vec<_>>();
    let goal_terms = terms(
        std::iter::once(context.goal)
            .chain(context.criteria.iter().map(String::as_str))
            .collect::<Vec<_>>()
            .join(" ")
            .as_str(),
    );
    let recent_start = ranges.len().saturating_sub(preserve_recent_units);
    let denominator = total_tokens.max(1);
    let unit_count = ranges.len().max(1) as u32;
    let mut scores = Vec::with_capacity(ranges.len());

    for (index, range) in ranges.iter().enumerate() {
        let slice = &messages[range.clone()];
        let text = &unit_texts[index];
        let folded_text = text.to_lowercase();
        let tokens = range
            .clone()
            .map(|message_index| {
                measurements
                    .get(message_index)
                    .map(|measurement| measurement.tokens)
                    .unwrap_or_else(|| {
                        let message = &messages[message_index];
                        engine.count_message(message)
                    })
            })
            .sum::<u32>();
        let goal_overlap = if goal_terms.is_empty() {
            0
        } else {
            overlap_count(&terms(text), &goal_terms)
        };
        let has_unresolved = has_unresolved(slice, &folded_text);
        let referenced_later = unit_referenced_later(slice, text, &unit_texts[index + 1..]);
        let is_error_or_decision = is_error_or_decision(slice, &folded_text);
        let dependency = context
            .preserved_refs
            .iter()
            .any(|reference| contains_folded(text, reference))
            || context
                .active_directives
                .iter()
                .any(|directive| directive_dependency(text, directive));
        let mandatory = index >= recent_start || has_unresolved || dependency;
        let recency = ((index as u64 + 1) * 1_000 / u64::from(unit_count)) as u32;
        let token_cost = (u64::from(tokens) * 1_000 / u64::from(denominator)) as u32;
        let prefix_invalidation_cost =
            ((ranges.len() - index) as u64 * 1_000 / u64::from(unit_count)) as u32;
        let utility = i64::from(goal_overlap) * 4_000
            + if has_unresolved { 20_000 } else { 0 }
            + if referenced_later { 5_000 } else { 0 }
            + if is_error_or_decision { 6_000 } else { 0 }
            + i64::from(recency) * 2
            - i64::from(token_cost) * 2
            - i64::from(prefix_invalidation_cost);
        scores.push(UtilityUnitScore {
            range: range.clone(),
            tokens,
            mandatory,
            goal_overlap,
            has_unresolved,
            referenced_later,
            is_error_or_decision,
            recency,
            token_cost,
            prefix_invalidation_cost,
            utility,
        });
    }

    if total_tokens <= target_tokens {
        return UtilityArchivePlan {
            archived_ranges: Vec::new(),
            retained_ranges: ranges,
            archived_tokens: 0,
            retained_tokens: scores.iter().map(|score| score.tokens).sum(),
            scores,
        };
    }

    let mut retained = scores
        .iter()
        .enumerate()
        .filter_map(|(index, score)| score.mandatory.then_some(index))
        .collect::<BTreeSet<_>>();
    let mut retained_tokens = retained
        .iter()
        .map(|index| scores[*index].tokens)
        .sum::<u32>();
    let mut optional = scores
        .iter()
        .enumerate()
        .filter_map(|(index, score)| (!score.mandatory).then_some(index))
        .collect::<Vec<_>>();
    optional.sort_by(|left, right| compare_density(&scores[*right], &scores[*left]));
    for index in optional {
        let tokens = scores[index].tokens;
        if retained_tokens.saturating_add(tokens) <= target_tokens {
            retained.insert(index);
            retained_tokens = retained_tokens.saturating_add(tokens);
        }
    }

    let retained_ranges = ranges
        .iter()
        .enumerate()
        .filter_map(|(index, range)| retained.contains(&index).then_some(range.clone()))
        .collect::<Vec<_>>();
    let archived_ranges = ranges
        .iter()
        .enumerate()
        .filter_map(|(index, range)| (!retained.contains(&index)).then_some(range.clone()))
        .collect::<Vec<_>>();
    let archived_tokens = scores
        .iter()
        .enumerate()
        .filter_map(|(index, score)| (!retained.contains(&index)).then_some(score.tokens))
        .sum();
    UtilityArchivePlan {
        archived_ranges,
        retained_ranges,
        archived_tokens,
        retained_tokens,
        scores,
    }
}

fn compare_density(left: &UtilityUnitScore, right: &UtilityUnitScore) -> Ordering {
    let left_density = i128::from(left.utility) * i128::from(right.tokens.max(1));
    let right_density = i128::from(right.utility) * i128::from(left.tokens.max(1));
    left_density
        .cmp(&right_density)
        .then_with(|| left.utility.cmp(&right.utility))
        .then_with(|| left.range.start.cmp(&right.range.start))
}

fn unit_text(messages: &[CoreMessage]) -> String {
    let mut text = String::new();
    let mut first_part = true;
    for message in messages {
        match &message.content {
            Content::Text(content) => append_unit_part(&mut text, &mut first_part, content),
            Content::Parts(content_parts) => {
                for part in content_parts {
                    match part {
                        ContentPart::Text { text: content } => {
                            append_unit_part(&mut text, &mut first_part, content)
                        }
                        ContentPart::ToolResult {
                            call_id, output, ..
                        } => {
                            append_unit_part(&mut text, &mut first_part, call_id.as_str());
                            text.push(' ');
                            text.push_str(output);
                        }
                        ContentPart::Image { source, .. } => append_unit_part(
                            &mut text,
                            &mut first_part,
                            match source {
                                crate::types::durable_content::DurableSource::Url { url } => url,
                                _ => "[image]",
                            },
                        ),
                        ContentPart::Audio { .. } => {
                            append_unit_part(&mut text, &mut first_part, "audio")
                        }
                    }
                }
            }
        }
        for call in &message.tool_calls {
            append_unit_part(&mut text, &mut first_part, call.id.as_str());
            text.push(' ');
            text.push_str(call.name.as_str());
            text.push(' ');
            text.push_str(&call.arguments.to_string());
        }
    }
    text
}

fn append_unit_part(text: &mut String, first_part: &mut bool, part: &str) {
    if !*first_part {
        text.push('\n');
    }
    *first_part = false;
    text.push_str(part);
}

fn contains_folded(text: &str, pattern: &str) -> bool {
    !pattern.trim().is_empty() && text.to_lowercase().contains(&pattern.to_lowercase())
}

fn directive_dependency(text: &str, directive: &str) -> bool {
    if contains_folded(text, directive) {
        return true;
    }
    let directive_terms = terms(directive);
    if directive_terms.is_empty() {
        return false;
    }
    let threshold = directive_terms.len().min(2);
    terms(text).intersection(&directive_terms).count() >= threshold
}

fn has_unresolved(messages: &[CoreMessage], folded_text: &str) -> bool {
    let mut opened = BTreeSet::new();
    let mut resolved = BTreeSet::new();
    for message in messages {
        for call in &message.tool_calls {
            opened.insert(call.id.to_string());
        }
        if let Content::Parts(parts) = &message.content {
            for part in parts {
                if let ContentPart::ToolResult {
                    call_id, is_error, ..
                } = part
                {
                    if *is_error {
                        return true;
                    }
                    resolved.insert(call_id.to_string());
                }
            }
        }
    }
    opened.iter().any(|call_id| !resolved.contains(call_id))
        || marker_folded(
            folded_text,
            &[
                "unresolved",
                "open question",
                "retry",
                "blocked",
                "待确认",
                "未解决",
                "重试",
                "阻塞",
            ],
        )
}

fn is_error_or_decision(messages: &[CoreMessage], folded_text: &str) -> bool {
    messages.iter().any(|message| {
        matches!(&message.content, Content::Parts(parts) if parts.iter().any(|part| matches!(part, ContentPart::ToolResult { is_error: true, .. })))
    }) || marker_folded(
        folded_text,
        &[
            "error", "failed", "failure", "exception", "decision", "decided", "must", "should",
            "错误", "失败", "异常", "决定", "选择", "必须", "应当",
        ],
    )
}

fn marker_folded(folded_text: &str, markers: &[&str]) -> bool {
    markers.iter().any(|marker| folded_text.contains(marker))
}

fn unit_referenced_later(messages: &[CoreMessage], text: &str, later: &[String]) -> bool {
    let mut references = messages
        .iter()
        .flat_map(|message| message.tool_calls.iter().map(|call| call.id.to_string()))
        .collect::<BTreeSet<_>>();
    references.extend(
        text.split_whitespace()
            .map(|token| token.trim_matches(|character: char| character.is_ascii_punctuation()))
            .filter(|token| token.contains('/') || token.contains("://"))
            .filter(|token| token.len() > 3)
            .map(str::to_string),
    );
    references.iter().any(|reference| {
        later
            .iter()
            .any(|later_text| contains_folded(later_text, reference))
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::message::{ContentPart, ToolCall};

    #[test]
    fn unit_text_preserves_empty_part_separators() {
        let messages = vec![CoreMessage::user(""), CoreMessage::user("next")];
        assert_eq!(unit_text(&messages), "\nnext");
    }

    #[test]
    fn unresolved_tool_unit_is_mandatory() {
        let mut call = CoreMessage::assistant("working");
        call.tool_calls.push(ToolCall {
            id: "call-1".into(),
            name: "read".into(),
            arguments: serde_json::json!({"path": "/work/a"}),
        });
        let mut recent = CoreMessage::user("recent");
        let messages = vec![call, recent];
        let engine = ContextTokenEngine::char_approx();
        let plan = plan_utility_archive(
            &messages,
            40,
            20,
            1,
            &engine,
            &UtilitySelectionContext {
                goal: "",
                criteria: &[],
                preserved_refs: &[],
                active_directives: &[],
            },
        );
        assert!(plan.scores[0].mandatory);
        assert!(plan.scores[0].has_unresolved);
        assert_eq!(plan.retained_tokens, 2);
    }

    #[test]
    fn chinese_directive_dependency_requires_bigram_overlap_not_shared_characters() {
        // Under the old per-character CJK vocabulary, any Chinese unit sharing two
        // common characters (中/文/回…) with an active directive was marked mandatory,
        // so compression could never archive unrelated Chinese history.
        let mut unrelated = CoreMessage::assistant("我们在文中回顾了天气");
        let mut on_topic = CoreMessage::user("已按要求保持中文回答");
        let mut recent = CoreMessage::user("recent");
        let messages = vec![unrelated, on_topic, recent];
        let plan = plan_utility_archive(
            &messages,
            70,
            10,
            1,
            &ContextTokenEngine::char_approx(),
            &UtilitySelectionContext {
                goal: "",
                criteria: &[],
                preserved_refs: &[],
                active_directives: &["必须用中文回答".into()],
            },
        );
        assert!(
            !plan.scores[0].mandatory,
            "unrelated Chinese text must not bind to the directive"
        );
        assert!(
            plan.scores[1].mandatory,
            "text restating the directive must stay mandatory"
        );
    }

    #[test]
    fn preserved_ref_keeps_complete_tool_unit() {
        let mut call = CoreMessage::assistant("read artifact");
        call.tool_calls.push(ToolCall {
            id: "call-keep".into(),
            name: "read".into(),
            arguments: serde_json::json!({}),
        });
        let mut result = CoreMessage::tool(vec![ContentPart::ToolResult {
            call_id: "call-keep".into(),
            output: "artifact".into(),
            is_error: false,
            durable_content: None,
        }]);
        let messages = vec![call, result];
        let plan = plan_utility_archive(
            &messages,
            40,
            0,
            0,
            &ContextTokenEngine::char_approx(),
            &UtilitySelectionContext {
                goal: "",
                criteria: &[],
                preserved_refs: &["call-keep".into()],
                active_directives: &[],
            },
        );
        assert!(plan.scores[0].mandatory);
        assert_eq!(plan.archived_ranges, Vec::<Range<usize>>::new());
    }

    #[test]
    fn measurement_aware_planner_ignores_stale_message_projection() {
        let mut message = CoreMessage::user("short");
        let engine = ContextTokenEngine::char_approx();
        let measurements = vec![TokenMeasurement::for_message(&message, 2)];
        let plan = plan_utility_archive_with_measurements(
            &[message],
            &measurements,
            2,
            1,
            0,
            &engine,
            &UtilitySelectionContext {
                goal: "",
                criteria: &[],
                preserved_refs: &[],
                active_directives: &[],
            },
        );
        assert_eq!(plan.scores[0].tokens, 2);
    }
}