(MOT-4014) fix: guarantee context assembly stays within budget

context-manager/README.md

@@ -67,6 +67,14 @@ system prompt under `# Conversation summary`, and any further compaction
 *updates* it instead of starting over. Callers that skip persistence stay
 correct at the cost of one summariser call per over-budget request.
 
+Every successful `context::assemble` response satisfies `token_count <= usable`.
+Callers should include the complete `tools` array and set
+`options.request_overhead_tokens` for response-format and provider-specific request
+fields. If ordinary pruning and compaction are insufficient, assembly replaces
+oversized function results in its model-facing copy with bounded references to the
+full result retained in the session transcript. Requests that still cannot fit fail
+with `context/overflow`; callers must not issue a provider request in that case.
+
 The other three functions: `context::count-tokens` (estimate messages + tools +
 system prompt vs a model), `context::prune` (replace verbose function outputs
 with `[output pruned: was ~N tokens]` placeholders, no LLM involved), and

context-manager/src/core/prune.rs

@@ -14,15 +14,25 @@
 //! - outputs whose text is at or under `max_output_chars` are not
 //!   "verbose" and stay (pruning them frees almost nothing);
 //! - when everything prunable frees under `min_free_tokens`, nothing is
-//!   touched at all (a no-op beats a destroyed-but-still-over context).
+//!   touched at all. `context::assemble` may subsequently use the
+//!   unconditional emergency pass to enforce its hard budget.
 
 use crate::core::estimate::Estimator;
 use crate::types::{AgentMessage, ContentBlock, Role};
+use serde_json::{json, Value};
+use sha2::{Digest, Sha256};
 
 /// Recent user turns that are always exempt, independent of the token
 /// window (prior-art constant, not operator-tunable).
 const PROTECTED_USER_TURNS: usize = 2;
 
+/// Maximum number of Unicode scalar values copied from an oversized
+/// result into its emergency reference.
+const EMERGENCY_PREVIEW_CHARS: usize = 160;
+
+const EMERGENCY_RETRIEVAL_HINT: &str =
+    "The original result remains in the session transcript; retrieve it by function_call_id.";
+
 #[derive(Debug, Clone)]
 pub struct PruneParams {
     /// Newest function-output tokens kept verbatim.
@@ -137,6 +147,117 @@ pub fn prune(
     }
 }
 
+/// Unconditionally reduce the largest function results until at least
+/// `required_tokens` have been freed or no result can shrink further.
+///
+/// Unlike [`prune`], this emergency pass has no recent-turn window,
+/// protected-function list, size threshold, or minimum-free guard. It
+/// preserves every message and the call/result identity fields while
+/// replacing both rendered content and opaque details with a bounded,
+/// deterministic reference to the original transcript entry.
+pub fn emergency_reduce(
+    messages: &mut [AgentMessage],
+    required_tokens: u64,
+    estimator: &dyn Estimator,
+) -> PruneStats {
+    let mut candidates: Vec<(usize, u64)> = messages
+        .iter()
+        .enumerate()
+        .filter(|(_, message)| matches!(message, AgentMessage::FunctionResult { .. }))
+        .map(|(idx, message)| (idx, estimator.message(message)))
+        .collect();
+
+    // Stable tie-break by transcript order keeps the reduction fully
+    // deterministic when multiple results have the same estimate.
+    candidates.sort_by(|(left_idx, left_tokens), (right_idx, right_tokens)| {
+        right_tokens
+            .cmp(left_tokens)
+            .then_with(|| left_idx.cmp(right_idx))
+    });
+
+    let mut stats = PruneStats {
+        scanned_parts: candidates.len() as u64,
+        ..PruneStats::default()
+    };
+
+    for (idx, original_tokens) in candidates {
+        if stats.pruned_tokens >= required_tokens {
+            break;
+        }
+
+        let replacement = emergency_reference(&messages[idx], original_tokens);
+        let replacement_tokens = estimator.message(&replacement);
+        let freed = original_tokens.saturating_sub(replacement_tokens);
+        if freed == 0 {
+            continue;
+        }
+
+        messages[idx] = replacement;
+        stats.pruned_tokens = stats.pruned_tokens.saturating_add(freed);
+        stats.pruned_parts += 1;
+    }
+
+    stats
+}
+
+fn emergency_reference(message: &AgentMessage, original_tokens: u64) -> AgentMessage {
+    let AgentMessage::FunctionResult {
+        function_call_id,
+        function_id,
+        content,
+        details,
+        is_error,
+        timestamp,
+    } = message
+    else {
+        unreachable!("emergency references are only built for function results");
+    };
+
+    let serialized = serde_json::to_vec(message)
+        .expect("serializing an AgentMessage containing serde_json::Value cannot fail");
+    let sha256 = format!("{:x}", Sha256::digest(&serialized));
+    let preview = emergency_preview(content, details);
+    let reference = json!({
+        "kind": "function_result_reference",
+        "function_id": function_id,
+        "function_call_id": function_call_id,
+        "original_bytes": serialized.len(),
+        "original_estimated_tokens": original_tokens,
+        "sha256": sha256,
+        "preview": preview,
+        "retrieval_hint": EMERGENCY_RETRIEVAL_HINT,
+    });
+    let rendered = format!(
+        "[function result reduced for context budget] {}",
+        serde_json::to_string(&reference)
+            .expect("serializing a function-result reference cannot fail")
+    );
+
+    AgentMessage::FunctionResult {
+        function_call_id: function_call_id.clone(),
+        function_id: function_id.clone(),
+        content: vec![ContentBlock::Text { text: rendered }],
+        details: json!({ "context_reference": reference }),
+        is_error: *is_error,
+        timestamp: *timestamp,
+    }
+}
+
+fn emergency_preview(content: &[ContentBlock], details: &Value) -> String {
+    let text = text_of(content);
+    let source = if text.is_empty() {
+        serde_json::to_string(details).unwrap_or_else(|_| "<unavailable>".into())
+    } else {
+        text
+    };
+    let mut chars = source.chars();
+    let mut preview: String = chars.by_ref().take(EMERGENCY_PREVIEW_CHARS).collect();
+    if chars.next().is_some() {
+        preview.push('\u{2026}');
+    }
+    preview
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -279,4 +400,141 @@ mod tests {
         // The placeholder is tiny, so nothing is verbose any more.
         assert_eq!(second.pruned_parts, 0);
     }
+
+    #[test]
+    fn emergency_reduces_latest_protected_result_and_preserves_pairing() {
+        let call: AgentMessage = serde_json::from_value(json!({
+            "role": "assistant",
+            "content": [{
+                "type": "function_call", "id": "c2",
+                "function_id": "protected::lookup", "arguments": {}
+            }],
+            "stop_reason": "function_call", "model": "m", "provider": "p", "timestamp": 1
+        }))
+        .unwrap();
+        let latest = result("protected::lookup", 40_000, 2);
+        let mut messages = vec![call.clone(), latest];
+
+        let mut normal_params = params();
+        normal_params.protected_functions = vec!["protected::lookup".into()];
+        normal_params.min_free_tokens = u64::MAX;
+        assert_eq!(
+            prune(&mut messages, &normal_params, &HeuristicEstimator).pruned_parts,
+            0
+        );
+
+        let before_len = messages.len();
+        let stats = emergency_reduce(&mut messages, 1, &HeuristicEstimator);
+        assert_eq!(stats.pruned_parts, 1);
+        assert_eq!(messages.len(), before_len);
+        assert_eq!(messages[0], call);
+
+        let AgentMessage::FunctionResult {
+            function_call_id,
+            function_id,
+            details,
+            ..
+        } = &messages[1]
+        else {
+            panic!("result message kind changed");
+        };
+        assert_eq!(function_call_id, "c2");
+        assert_eq!(function_id, "protected::lookup");
+        assert_eq!(details["context_reference"]["function_call_id"], "c2");
+        assert_eq!(
+            details["context_reference"]["function_id"],
+            "protected::lookup"
+        );
+    }
+
+    #[test]
+    fn emergency_reference_bounds_content_and_details() {
+        let mut message: AgentMessage = serde_json::from_value(json!({
+            "role": "function_result",
+            "function_call_id": "call-large",
+            "function_id": "shell::run",
+            "content": [{ "type": "text", "text": "content".repeat(100_000) }],
+            "details": { "stdout": "details".repeat(100_000) },
+            "is_error": false,
+            "timestamp": 9
+        }))
+        .unwrap();
+
+        emergency_reduce(std::slice::from_mut(&mut message), 1, &HeuristicEstimator);
+
+        let serialized = serde_json::to_string(&message).unwrap();
+        assert!(
+            serialized.len() < 2_000,
+            "reference was {} bytes",
+            serialized.len()
+        );
+        let AgentMessage::FunctionResult {
+            content, details, ..
+        } = &message
+        else {
+            panic!("result message kind changed");
+        };
+        let rendered = text_of(content);
+        assert!(rendered.contains("original_estimated_tokens"));
+        assert!(rendered.contains("original_bytes"));
+        assert!(rendered.contains("session transcript"));
+        assert!(
+            details["context_reference"]["preview"]
+                .as_str()
+                .unwrap()
+                .chars()
+                .count()
+                <= EMERGENCY_PREVIEW_CHARS + 1
+        );
+        assert!(details["context_reference"]["sha256"]
+            .as_str()
+            .is_some_and(|hash| hash.len() == 64));
+    }
+
+    #[test]
+    fn emergency_reference_hash_is_deterministic_for_original_message() {
+        let original: AgentMessage = serde_json::from_value(json!({
+            "role": "function_result",
+            "function_call_id": "hash-call",
+            "function_id": "fs::read",
+            "content": [{ "type": "text", "text": "deterministic payload".repeat(1_000) }],
+            "details": { "path": "/tmp/example" },
+            "is_error": false,
+            "timestamp": 7
+        }))
+        .unwrap();
+        let expected = format!(
+            "{:x}",
+            Sha256::digest(serde_json::to_vec(&original).unwrap())
+        );
+        let mut first = original.clone();
+        let mut second = original;
+
+        emergency_reduce(std::slice::from_mut(&mut first), 1, &HeuristicEstimator);
+        emergency_reduce(std::slice::from_mut(&mut second), 1, &HeuristicEstimator);
+
+        assert_eq!(first, second);
+        for reduced in [&first, &second] {
+            let AgentMessage::FunctionResult { details, .. } = reduced else {
+                panic!("result message kind changed");
+            };
+            assert_eq!(details["context_reference"]["sha256"], expected);
+        }
+    }
+
+    #[test]
+    fn emergency_reduces_largest_results_only_as_needed() {
+        let mut messages = vec![
+            result("small", 4_000, 1),
+            result("largest", 40_000, 2),
+            result("middle", 20_000, 3),
+        ];
+
+        let stats = emergency_reduce(&mut messages, 1, &HeuristicEstimator);
+
+        assert_eq!(stats.pruned_parts, 1);
+        assert_eq!(text_of(messages[0].content()).len(), 4_000);
+        assert!(text_of(messages[1].content()).starts_with("[function result reduced"));
+        assert_eq!(text_of(messages[2].content()).len(), 20_000);
+    }
 }

context-manager/src/error.rs

@@ -23,6 +23,13 @@ pub enum ContextError {
     /// A filesystem operation backing the compaction lease failed.
     #[error("context/state: {0}")]
     State(String),
+
+    /// No safe model-facing context can fit within the usable input
+    /// budget after normal pruning, compaction, and emergency reduction.
+    #[error(
+        "context/overflow: assembled context requires {token_count} tokens but usable budget is {usable}"
+    )]
+    Overflow { token_count: u64, usable: u64 },
 }
 
 impl ContextError {
@@ -32,6 +39,7 @@ impl ContextError {
             ContextError::InvalidRequest(_) => "context/invalid_request",
             ContextError::ModelUnresolved(_) => "context/model_unresolved",
             ContextError::State(_) => "context/state",
+            ContextError::Overflow { .. } => "context/overflow",
         }
     }
 }
@@ -68,6 +76,10 @@ mod tests {
             ContextError::InvalidRequest("m".into()),
             ContextError::ModelUnresolved("m".into()),
             ContextError::State("m".into()),
+            ContextError::Overflow {
+                token_count: 101,
+                usable: 100,
+            },
         ];
         for v in variants {
             assert!(
@@ -77,4 +89,25 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn overflow_carries_budget_values_and_stable_code() {
+        let error = ContextError::Overflow {
+            token_count: 12_345,
+            usable: 10_000,
+        };
+
+        assert_eq!(error.code(), "context/overflow");
+        assert_eq!(
+            error.to_string(),
+            "context/overflow: assembled context requires 12345 tokens but usable budget is 10000"
+        );
+        assert!(matches!(
+            error,
+            ContextError::Overflow {
+                token_count: 12_345,
+                usable: 10_000
+            }
+        ));
+    }
 }