diff --git a/src/agents/hermes/templates/plugin_init.py b/src/agents/hermes/templates/plugin_init.py index d1ab40ca..8fa0447c 100644 --- a/src/agents/hermes/templates/plugin_init.py +++ b/src/agents/hermes/templates/plugin_init.py @@ -2243,6 +2243,116 @@ def _replay_message_list(value): return None return list(value) +def _normalize_replay_tool_pairs(messages): + """Return a replay containing only structurally paired tool events. + + The returned issue list is intentionally separate from the normalized + candidate: callers must reject any replay that required repair instead of + silently persisting a transcript different from the one TraceDecay built. + """ + issues = [] + normalized = [] + index = 0 + messages = list(messages or []) + while index < len(messages): + message = messages[index] + if not isinstance(message, dict): + index += 1 + continue + if message.get("role") == "tool": + issues.append({ + "code": "orphan_tool_result", + "tool_call_id": str(message.get("tool_call_id") or "").strip(), + "message_index": index, + }) + index += 1 + continue + tool_calls = message.get("tool_calls") + if tool_calls is None: + normalized.append(message) + index += 1 + continue + if message.get("role") != "assistant" or not isinstance(tool_calls, list) or not tool_calls: + issues.append({"code": "invalid_tool_calls", "message_index": index}) + cleaned = dict(message) + cleaned.pop("tool_calls", None) + if str(cleaned.get("content") or "").strip(): + normalized.append(cleaned) + index += 1 + continue + + call_ids = [_summary_tool_call_id(call) for call in tool_calls] + result_end = index + 1 + while result_end < len(messages) and isinstance(messages[result_end], dict) \ + and messages[result_end].get("role") == "tool": + result_end += 1 + results = messages[index + 1:result_end] + result_ids = [str(result.get("tool_call_id") or "").strip() for result in results] + trailing_open = ( + result_end == len(messages) + and not results + and all(call_ids) + and len(set(call_ids)) == len(call_ids) + ) + if trailing_open: + normalized.append(message) + index = result_end + continue + complete = ( + all(call_ids) + and len(set(call_ids)) == len(call_ids) + and all(result_ids) + and len(set(result_ids)) == len(result_ids) + and set(result_ids) == set(call_ids) + and len(results) == len(tool_calls) + ) + if complete: + normalized.append(message) + normalized.extend(results) + else: + for tool_call_id in sorted(set(call_ids) | set(result_ids)): + issues.append({ + "code": "unmatched_tool_pair", + "tool_call_id": tool_call_id, + "call_count": call_ids.count(tool_call_id), + "result_count": result_ids.count(tool_call_id), + "message_index": index, + }) + cleaned = dict(message) + cleaned.pop("tool_calls", None) + if str(cleaned.get("content") or "").strip(): + normalized.append(cleaned) + index = result_end + return normalized, issues + +def _compression_boundary_abort( + result, + code, + *, + source_token_estimate, + replay_token_estimate, + reported_source_tokens=None, + reported_replay_tokens=None, + issues=None, +): + payload = dict(result) if isinstance(result, dict) else {} + raw_replay = payload.pop("replay_messages", None) + payload["status"] = "aborted" + payload["reason"] = code + payload["replay_messages"] = [] + payload["compression_diagnostic"] = { + "type": "replay_boundary_rejection", + "code": code, + "defer_preflight_to_real_usage": True, + "source_token_estimate": source_token_estimate, + "replay_token_estimate": replay_token_estimate, + "reported_source_tokens": reported_source_tokens, + "reported_replay_tokens": reported_replay_tokens, + "rejected_replay_message_count": len(raw_replay) if isinstance(raw_replay, list) else 0, + "issues": list(issues or []), + } + return payload + def _compression_replay_is_compacted(result): if not isinstance(result, dict): return False @@ -2911,7 +3021,17 @@ def has_content_to_compress(self, messages, current_tokens=None, **kwargs): return len(non_empty) >= 2 def should_defer_preflight_to_real_usage(self, rough_tokens=None): - if not ( + diagnostic = ( + self.last_compress_result.get("compression_diagnostic") + if isinstance(self.last_compress_result, dict) + else None + ) + replay_boundary_deferred = bool( + isinstance(diagnostic, dict) + and diagnostic.get("type") == "replay_boundary_rejection" + and diagnostic.get("defer_preflight_to_real_usage") + ) + if not replay_boundary_deferred and not ( isinstance(self.last_compress_result, dict) and self.last_compress_result.get("status") == "compressed" ): @@ -3488,6 +3608,51 @@ def compress(self, messages, current_tokens=None, focus_topic=None, **kwargs): return original if replay == original: return original + normalized_replay, replay_issues = _normalize_replay_tool_pairs(replay) + source_token_estimate = _count_messages_tokens(original) + replay_token_estimate = _count_messages_tokens(normalized_replay) + try: + reported_source_tokens = int(current_tokens) if current_tokens is not None else None + except (TypeError, ValueError): + reported_source_tokens = None + try: + reported_replay_tokens = int(result.get("replay_token_estimate")) + except (TypeError, ValueError): + reported_replay_tokens = None + has_reported_estimates = ( + reported_source_tokens is not None + and reported_source_tokens > 0 + and reported_replay_tokens is not None + ) + boundary_code = None + if replay_issues: + boundary_code = "invalid_replay_tool_pairs" + elif ( + has_reported_estimates + and reported_replay_tokens >= reported_source_tokens + ): + boundary_code = "non_shrinking_reported_replay" + elif ( + not has_reported_estimates + and source_token_estimate > 0 + and replay_token_estimate >= source_token_estimate + ): + boundary_code = "non_shrinking_replay" + if boundary_code is not None: + result = _compression_boundary_abort( + result, + boundary_code, + source_token_estimate=source_token_estimate, + replay_token_estimate=replay_token_estimate, + reported_source_tokens=reported_source_tokens, + reported_replay_tokens=reported_replay_tokens, + issues=replay_issues, + ) + self.last_compress_result = result + self._last_compress_aborted = True + self._last_summary_error = boundary_code + return original + replay = normalized_replay if replay != original: self.compression_count += 1 return replay diff --git a/src/sessions/lcm/compression.rs b/src/sessions/lcm/compression.rs index 0de23a4c..8862be7c 100644 --- a/src/sessions/lcm/compression.rs +++ b/src/sessions/lcm/compression.rs @@ -19,12 +19,10 @@ use super::{ LcmLifecycleState, LcmLifecycleUpdate, LcmMaintenanceDebt, LcmPreflightRequest, LcmPreflightResponse, LcmRawMessage, LcmSessionBoundaryRequest, LcmSessionBoundaryResponse, LcmSourceRef, LcmStorageKind, LcmSummaryNode, LcmSummaryNodeDraft, LcmSummaryRequest, dag, - payload, raw, security, util, + payload, raw, replay_transactions, security, util, }; const MAX_FORCED_CATCHUP_PASSES: usize = 4; const MIN_SUMMARY_RESCUE_SOURCE_TOKENS: i64 = 8; -const ACTIVE_REPLAY_METADATA_KEY: &str = "lcm_active_replay"; -const ACTIVE_REPLAY_MESSAGE_KEY: &str = "active_replay"; const PRESERVED_TODO_CONTEXT_PREFIX: &str = "[Your active task list was preserved across context compression]"; const PRESERVED_OBJECTIVE_CONTEXT_PREFIX: &str = @@ -1116,6 +1114,7 @@ fn compression_window( let backlog_len = unsummarized .len() .saturating_sub(effective_fresh_tail_count); + let backlog_len = replay_transactions::atomic_tail_start(&unsummarized, backlog_len); let (older_unsummarized, fresh_tail) = unsummarized.split_at(backlog_len); let fresh_tail_start_store_id = fresh_tail .first() @@ -1163,9 +1162,17 @@ fn replay_without_summary( fresh_tail: &[LcmRawMessage], ) -> Vec { let mut replay_messages = Vec::with_capacity(pinned_anchors.len() + fresh_tail.len()); - replay_messages.extend(pinned_anchors.iter().map(raw_replay_message)); - replay_messages.extend(fresh_tail.iter().map(raw_replay_message)); - replay_messages + replay_messages.extend( + pinned_anchors + .iter() + .map(replay_transactions::raw_replay_message), + ); + replay_messages.extend( + fresh_tail + .iter() + .map(replay_transactions::raw_replay_message), + ); + replay_transactions::normalize_replay_tool_pairs(&replay_messages) } const SUMMARY_REPLAY_PRIORITY: u8 = 0; @@ -1221,13 +1228,18 @@ async fn assemble_overflow_recovery_replay( max_assembly_tokens, ); if candidate.len() == anchors.len() { - if let Some(last) = raws.last() { + if let Some(last_unit) = replay_transactions::replay_units(&raws).last() { let mut replay = anchors .iter() - .map(|message| raw_replay_message(message)) + .map(|message| replay_transactions::raw_replay_message(message)) .collect::>(); - replay.push(raw_replay_message(last)); - return Ok(replay); + replay.extend( + last_unit + .messages + .iter() + .map(|message| replay_transactions::raw_replay_message(message)), + ); + return Ok(replay_transactions::normalize_replay_tool_pairs(&replay)); } } Ok(candidate) @@ -1263,7 +1275,10 @@ fn assemble_replay_messages( let (selected_raws, selected_summaries, preserved_objective_anchor) = match max_assembly_tokens { None => ( - raws.to_vec(), + replay_transactions::replay_units(raws) + .into_iter() + .flat_map(|unit| unit.messages) + .collect(), summaries.iter().collect::>(), latest_user_context_anchor(anchor_source, raws), ), @@ -1307,7 +1322,7 @@ fn assemble_replay_messages( ( message.store_id, RAW_REPLAY_PRIORITY, - raw_replay_message(message), + replay_transactions::raw_replay_message(message), ) })); replay_items.extend(selected_summaries.iter().map(|summary| { @@ -1333,14 +1348,15 @@ fn assemble_replay_messages( ( message.store_id, RAW_REPLAY_PRIORITY, - raw_replay_message(message), + replay_transactions::raw_replay_message(message), ) })); replay_items.sort_by_key(|(store_id, priority, _)| (*store_id, *priority)); - replay_items + let replay = replay_items .into_iter() .map(|(_, _, message)| message) - .collect() + .collect::>(); + replay_transactions::normalize_replay_tool_pairs(&replay) } /// Mirrors hermes-lcm `_assemble_context` tail selection: keep the newest @@ -1355,10 +1371,14 @@ fn select_budget_tail<'a>( let mut kept_reversed = Vec::new(); let mut tail_tokens = 0i64; let mut skipped_gap = false; - for message in raws.iter().rev() { - let message_tokens = estimate_tokens(&message.content); - if used + tail_tokens + message_tokens > cap { - if is_budget_droppable_tail_message(message) { + for unit in replay_transactions::replay_units(raws).iter().rev() { + let unit_tokens = unit.token_count(); + if used + tail_tokens + unit_tokens > cap { + if unit + .messages + .iter() + .all(|message| is_budget_droppable_tail_message(message)) + { skipped_gap = true; continue; } @@ -1367,8 +1387,8 @@ fn select_budget_tail<'a>( if skipped_gap { break; } - kept_reversed.push(*message); - tail_tokens += message_tokens; + kept_reversed.extend(unit.messages.iter().rev().copied()); + tail_tokens += unit_tokens; } kept_reversed.reverse(); (kept_reversed, tail_tokens) @@ -2027,9 +2047,12 @@ fn default_message_kind(role: &str) -> String { fn active_message_metadata(message: &Value, replay: &Value) -> String { let mut metadata = Map::new(); - metadata.insert(ACTIVE_REPLAY_METADATA_KEY.to_string(), Value::Bool(true)); metadata.insert( - ACTIVE_REPLAY_MESSAGE_KEY.to_string(), + replay_transactions::ACTIVE_REPLAY_METADATA_KEY.to_string(), + Value::Bool(true), + ); + metadata.insert( + replay_transactions::ACTIVE_REPLAY_MESSAGE_KEY.to_string(), active_replay_for_metadata(replay), ); if let Some(lcm_ingest) = message.get("lcm_ingest") { @@ -2057,9 +2080,12 @@ fn active_replay_metadata_json(existing_metadata_json: Option<&str>, replay: &Va .and_then(|text| serde_json::from_str::(text).ok()) .and_then(|value| value.as_object().cloned()) .unwrap_or_default(); - metadata.insert(ACTIVE_REPLAY_METADATA_KEY.to_string(), Value::Bool(true)); metadata.insert( - ACTIVE_REPLAY_MESSAGE_KEY.to_string(), + replay_transactions::ACTIVE_REPLAY_METADATA_KEY.to_string(), + Value::Bool(true), + ); + metadata.insert( + replay_transactions::ACTIVE_REPLAY_MESSAGE_KEY.to_string(), active_replay_for_metadata(replay), ); Value::Object(metadata).to_string() @@ -2068,7 +2094,7 @@ fn active_replay_metadata_json(existing_metadata_json: Option<&str>, replay: &Va fn active_replay_for_metadata(replay: &Value) -> Value { let mut replay = replay.clone(); if let Some(object) = replay.as_object_mut() { - strip_disposable_assistant_replay_sidecars(object, ""); + replay_transactions::strip_disposable_assistant_replay_sidecars(object, ""); } replay } @@ -2219,82 +2245,6 @@ fn deterministic_message_id( ) } -fn raw_replay_message(message: &LcmRawMessage) -> Value { - if let Some(mut replay) = active_replay_message_from_metadata(message) { - replay["role"] = Value::String(message.role.clone()); - replay["store_id"] = Value::from(message.store_id); - return replay; - } - json!({ - "role": message.role, - "content": message.content, - "store_id": message.store_id, - }) -} - -fn active_replay_message_from_metadata(message: &LcmRawMessage) -> Option { - let metadata: Value = serde_json::from_str(message.metadata_json.as_deref()?).ok()?; - if metadata - .get(ACTIVE_REPLAY_METADATA_KEY) - .and_then(Value::as_bool) - != Some(true) - { - return None; - } - let mut replay = metadata - .get(ACTIVE_REPLAY_MESSAGE_KEY) - .and_then(Value::as_object) - .cloned() - .or_else(|| legacy_active_replay_message_from_metadata(&metadata))?; - if !replay.contains_key("content") { - replay.insert( - "content".to_string(), - Value::String(message.content.clone()), - ); - } - strip_disposable_assistant_replay_sidecars(&mut replay, &message.role); - Some(Value::Object(replay)) -} - -fn strip_disposable_assistant_replay_sidecars( - replay: &mut Map, - fallback_role: &str, -) { - if !replay - .get("role") - .and_then(Value::as_str) - .unwrap_or(fallback_role) - .eq_ignore_ascii_case("assistant") - { - return; - } - - // Provider replay sidecars are useful before the next API call, but once - // LCM is rebuilding compressed history they become large derived state. - for key in [ - "codex_message_items", - "codex_reasoning_items", - "reasoning", - "reasoning_content", - "reasoning_details", - ] { - replay.remove(key); - } -} - -fn legacy_active_replay_message_from_metadata(metadata: &Value) -> Option> { - let mut replay = metadata.as_object()?.clone(); - replay.remove(ACTIVE_REPLAY_METADATA_KEY); - replay.remove(ACTIVE_REPLAY_MESSAGE_KEY); - replay.remove("ingest_protection"); - replay.remove("external_payload"); - replay.remove("payload_ref"); - replay.remove("byte_count"); - replay.remove("char_count"); - replay.remove("sha256"); - Some(replay) -} - fn summary_replay_message(summary: &LcmSummaryNode) -> Value { json!({ "role": "system", diff --git a/src/sessions/lcm/compression_decision.rs b/src/sessions/lcm/compression_decision.rs index 64f2407a..fea49bfc 100644 --- a/src/sessions/lcm/compression_decision.rs +++ b/src/sessions/lcm/compression_decision.rs @@ -1,5 +1,6 @@ use serde_json::Value; +use super::replay_transactions; use super::summarizer::CompressionSummarizerAdapter; use super::{ LCM_COMPRESSION_BOUNDARY_COOLDOWN_SECONDS, LCM_DEFAULT_SUMMARY_FAN_IN, LcmCompressionRequest, @@ -324,7 +325,7 @@ pub fn bounded_leaf_chunk_len( selected_tokens += message_tokens; selected_len += 1; } - selected_len + replay_transactions::bounded_atomic_prefix_len(backlog, selected_len) } pub fn progress_leaf_chunk_len( @@ -334,7 +335,7 @@ pub fn progress_leaf_chunk_len( ) -> usize { let selected_len = bounded_leaf_chunk_len(backlog, leaf_chunk_tokens, max_source_messages); if selected_len == 0 && !backlog.is_empty() { - 1 + replay_transactions::first_atomic_unit_len(backlog) } else { selected_len } diff --git a/src/sessions/lcm/mod.rs b/src/sessions/lcm/mod.rs index b89935c7..c9114507 100644 --- a/src/sessions/lcm/mod.rs +++ b/src/sessions/lcm/mod.rs @@ -8,6 +8,7 @@ pub mod hermes; pub mod payload; pub mod query; pub mod raw; +mod replay_transactions; pub mod schema; pub mod security; mod summarizer; diff --git a/src/sessions/lcm/replay_transactions.rs b/src/sessions/lcm/replay_transactions.rs new file mode 100644 index 00000000..6630e036 --- /dev/null +++ b/src/sessions/lcm/replay_transactions.rs @@ -0,0 +1,358 @@ +use std::collections::BTreeSet; + +use serde_json::{Map, Value, json}; + +use super::LcmRawMessage; + +pub(crate) const ACTIVE_REPLAY_METADATA_KEY: &str = "lcm_active_replay"; +pub(crate) const ACTIVE_REPLAY_MESSAGE_KEY: &str = "active_replay"; + +#[derive(Debug)] +pub(crate) struct ReplayUnit<'a> { + pub(crate) messages: Vec<&'a LcmRawMessage>, +} + +impl ReplayUnit<'_> { + pub(crate) fn token_count(&self) -> i64 { + self.messages + .iter() + .map(|message| replay_message_tokens(message)) + .sum() + } +} + +/// Back a bounded prefix off when its boundary bisects a tool transaction. +pub(crate) fn bounded_atomic_prefix_len(messages: &[LcmRawMessage], requested_len: usize) -> usize { + let mut selected_len = requested_len.min(messages.len()); + for (start, end) in transaction_ranges(messages.iter()) { + if start < selected_len && selected_len < end { + selected_len = start; + } + } + selected_len +} + +/// Select one complete leading transaction when progress would otherwise be +/// zero. This is the only path allowed to exceed a configured prefix cap. +pub(crate) fn first_atomic_unit_len(messages: &[LcmRawMessage]) -> usize { + transaction_ranges(messages.iter()) + .into_iter() + .find_map(|(start, end)| (start == 0).then_some(end)) + .unwrap_or(1) + .min(messages.len()) +} + +/// Move a suffix boundary backward when it falls inside a valid tool transaction. +pub(crate) fn atomic_tail_start(messages: &[LcmRawMessage], requested_start: usize) -> usize { + let mut start = requested_start.min(messages.len()); + for (transaction_start, transaction_end) in transaction_ranges(messages.iter()) { + if transaction_start < start && start < transaction_end { + start = transaction_start; + } + } + start +} + +/// Build replay-selection units. Valid assistant call + consecutive result +/// transactions are atomic. Legacy orphan results are omitted; an unmatched +/// assistant call remains only when its visible content is useful, and the +/// final value normalizer strips its invalid `tool_calls` field. +pub(crate) fn replay_units<'a>(messages: &[&'a LcmRawMessage]) -> Vec> { + let transaction_ranges = transaction_ranges(messages.iter().copied()); + let mut transaction_by_start = transaction_ranges.into_iter().peekable(); + let mut units = Vec::new(); + let mut index = 0; + while index < messages.len() { + if let Some((start, end)) = transaction_by_start.peek().copied() + && start == index + { + let _ = transaction_by_start.next(); + units.push(ReplayUnit { + messages: messages[index..end].to_vec(), + }); + index = end; + continue; + } + let message = messages[index]; + if message.role == "tool" { + index += 1; + continue; + } + let trailing_open_call = index + 1 == messages.len() && tool_call_ids(message).is_some(); + if tool_call_ids(message).is_some() + && message.content.trim().is_empty() + && !trailing_open_call + { + index += 1; + continue; + } + units.push(ReplayUnit { + messages: vec![message], + }); + index += 1; + } + units +} + +fn transaction_ranges<'a>( + messages: impl IntoIterator, +) -> Vec<(usize, usize)> { + let messages = messages.into_iter().collect::>(); + let mut ranges = Vec::new(); + let mut index = 0; + while index < messages.len() { + let Some(expected_ids) = tool_call_ids(messages[index]) else { + index += 1; + continue; + }; + if messages[index].role != "assistant" || expected_ids.is_empty() { + index += 1; + continue; + } + let mut found_ids = BTreeSet::new(); + let mut end = index + 1; + while end < messages.len() && messages[end].role == "tool" { + if let Some(tool_call_id) = tool_result_id(messages[end]) { + found_ids.insert(tool_call_id); + } + end += 1; + } + if !found_ids.is_empty() + && found_ids.is_subset(&expected_ids) + && end - index - 1 == found_ids.len() + { + ranges.push((index, end)); + index = end; + } else { + index += 1; + } + } + ranges +} + +fn tool_call_ids(message: &LcmRawMessage) -> Option> { + let replay = active_replay_value(message)?; + let tool_calls = replay.get("tool_calls")?.as_array()?; + Some( + tool_calls + .iter() + .filter_map(|tool_call| tool_call.get("id").and_then(Value::as_str)) + .filter(|tool_call_id| !tool_call_id.is_empty()) + .map(str::to_string) + .collect(), + ) +} + +fn tool_result_id(message: &LcmRawMessage) -> Option { + active_replay_value(message)? + .get("tool_call_id")? + .as_str() + .filter(|tool_call_id| !tool_call_id.is_empty()) + .map(str::to_string) +} + +fn active_replay_value(message: &LcmRawMessage) -> Option { + active_replay_message_from_metadata(message) +} + +fn estimate_tokens(text: &str) -> i64 { + text.split_whitespace().count().max(1) as i64 +} + +fn replay_message_tokens(message: &LcmRawMessage) -> i64 { + let mut tokens = estimate_tokens(&message.content); + if let Some(mut replay) = active_replay_value(message) { + if let Some(object) = replay.as_object_mut() { + object.remove("content"); + object.remove("role"); + object.remove("store_id"); + } + let serialized = replay.to_string(); + tokens += serialized.len().div_ceil(4) as i64; + } + tokens +} + +pub(crate) fn normalize_replay_tool_pairs(messages: &[Value]) -> Vec { + let mut normalized = Vec::with_capacity(messages.len()); + let mut index = 0; + while index < messages.len() { + let mut message = messages[index].clone(); + let role = message.get("role").and_then(Value::as_str).unwrap_or(""); + if role == "tool" { + index += 1; + continue; + } + if role != "assistant" || message.get("tool_calls").is_none() { + normalized.push(message); + index += 1; + continue; + } + let tool_calls = message + .get("tool_calls") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let mut result_end = index + 1; + while result_end < messages.len() + && messages[result_end].get("role").and_then(Value::as_str) == Some("tool") + { + result_end += 1; + } + let open_call_ids = tool_calls + .iter() + .filter_map(|tool_call| tool_call.get("id").and_then(Value::as_str)) + .collect::>(); + if result_end == messages.len() + && result_end == index + 1 + && !open_call_ids.is_empty() + && open_call_ids + .iter() + .all(|tool_call_id| !tool_call_id.is_empty()) + && open_call_ids.iter().collect::>().len() == open_call_ids.len() + { + normalized.push(message); + break; + } + let call_id_counts = tool_calls + .iter() + .filter_map(|tool_call| tool_call.get("id").and_then(Value::as_str)) + .fold( + std::collections::BTreeMap::new(), + |mut counts, tool_call_id| { + *counts.entry(tool_call_id.to_string()).or_insert(0usize) += 1; + counts + }, + ); + let result_id_counts = messages[index + 1..result_end] + .iter() + .filter_map(|result| result.get("tool_call_id").and_then(Value::as_str)) + .fold( + std::collections::BTreeMap::new(), + |mut counts, tool_call_id| { + *counts.entry(tool_call_id.to_string()).or_insert(0usize) += 1; + counts + }, + ); + let kept_calls = tool_calls + .into_iter() + .filter(|tool_call| { + tool_call + .get("id") + .and_then(Value::as_str) + .is_some_and(|tool_call_id| { + call_id_counts.get(tool_call_id) == Some(&1) + && result_id_counts.get(tool_call_id) == Some(&1) + }) + }) + .collect::>(); + if kept_calls.is_empty() { + message + .as_object_mut() + .map(|object| object.remove("tool_calls")); + if !replay_content_text(&message).trim().is_empty() { + normalized.push(message); + } + } else { + message["tool_calls"] = Value::Array(kept_calls); + normalized.push(message); + normalized.extend( + messages[index + 1..result_end] + .iter() + .filter(|result| { + result + .get("tool_call_id") + .and_then(Value::as_str) + .is_some_and(|tool_call_id| { + call_id_counts.get(tool_call_id) == Some(&1) + && result_id_counts.get(tool_call_id) == Some(&1) + }) + }) + .cloned(), + ); + } + index = result_end; + } + normalized +} + +fn replay_content_text(message: &Value) -> String { + match message.get("content") { + None | Some(Value::Null) => String::new(), + Some(Value::String(text)) => text.clone(), + Some(content) => serde_json::to_string(content).unwrap_or_default(), + } +} + +pub(crate) fn raw_replay_message(message: &LcmRawMessage) -> Value { + if let Some(mut replay) = active_replay_message_from_metadata(message) { + replay["role"] = Value::String(message.role.clone()); + replay["store_id"] = Value::from(message.store_id); + return replay; + } + json!({ + "role": message.role, + "content": message.content, + "store_id": message.store_id, + }) +} + +fn active_replay_message_from_metadata(message: &LcmRawMessage) -> Option { + let metadata: Value = serde_json::from_str(message.metadata_json.as_deref()?).ok()?; + if metadata + .get(ACTIVE_REPLAY_METADATA_KEY) + .and_then(Value::as_bool) + != Some(true) + { + return None; + } + let mut replay = metadata + .get(ACTIVE_REPLAY_MESSAGE_KEY) + .and_then(Value::as_object) + .cloned() + .or_else(|| legacy_active_replay_message_from_metadata(&metadata))?; + if !replay.contains_key("content") { + replay.insert( + "content".to_string(), + Value::String(message.content.clone()), + ); + } + strip_disposable_assistant_replay_sidecars(&mut replay, &message.role); + Some(Value::Object(replay)) +} + +pub(crate) fn strip_disposable_assistant_replay_sidecars( + replay: &mut Map, + fallback_role: &str, +) { + if !replay + .get("role") + .and_then(Value::as_str) + .unwrap_or(fallback_role) + .eq_ignore_ascii_case("assistant") + { + return; + } + for key in [ + "codex_message_items", + "codex_reasoning_items", + "reasoning", + "reasoning_content", + "reasoning_details", + ] { + replay.remove(key); + } +} + +fn legacy_active_replay_message_from_metadata(metadata: &Value) -> Option> { + let mut replay = metadata.as_object()?.clone(); + replay.remove(ACTIVE_REPLAY_METADATA_KEY); + replay.remove(ACTIVE_REPLAY_MESSAGE_KEY); + replay.remove("ingest_protection"); + replay.remove("external_payload"); + replay.remove("payload_ref"); + replay.remove("byte_count"); + replay.remove("char_count"); + replay.remove("sha256"); + Some(replay) +} diff --git a/tests/agent_suite/agent_test.rs b/tests/agent_suite/agent_test.rs index 8f1cd8bc..3d559724 100644 --- a/tests/agent_suite/agent_test.rs +++ b/tests/agent_suite/agent_test.rs @@ -1276,7 +1276,7 @@ fn test_hermes_plugin_init_snapshot_matches_embedded_asset() { hasher.update(body.as_bytes()); assert_eq!( hex::encode(hasher.finalize()), - "954dfa34aad7753f3ad6ff5e536061066e8339f3ee167e65ea3f27ca5a879575", + "eaaa993f8b0742deca03d7a6f53ffca797695a0423637fd33206bf80bbe0f57d", "templates/plugin_init.py payload hash changed — verify the edit is intentional and update this snapshot" ); } diff --git a/tests/hermes_suite/lcm_bridge.rs b/tests/hermes_suite/lcm_bridge.rs index ce6ec8c9..9c7a99ef 100644 --- a/tests/hermes_suite/lcm_bridge.rs +++ b/tests/hermes_suite/lcm_bridge.rs @@ -2009,6 +2009,206 @@ assert engine.last_compress_result["reason"] == "noop_summarizer" ); } +#[test] +fn context_engine_rejects_orphan_heavy_replay_at_host_boundary() { + run_generated_plugin_script( + "check_orphan_heavy_replay_rejected.py", + r#" +messages = [ + {"role": "user", "content": "original valid transcript " * 80}, + {"role": "assistant", "content": "still valid"}, +] +replay = [ + {"role": "system", "content": "Useful compact summary"}, + {"role": "tool", "tool_call_id": "orphan-result", "content": "unmatched output"}, + { + "role": "assistant", + "content": "", + "tool_calls": [{ + "id": "orphan-call", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }], + }, +] + +engine = plugin.TraceDecayContextEngine() +engine.initialize(session_id="session-1", project_root="/tmp/project") +engine.last_real_prompt_tokens = 100 +engine._compress_to_result = lambda *args, **kwargs: { + "status": "ok", + "reason": "compressed_backlog", + "replay_messages": replay, + "replay_token_estimate": 40, +} + +compressed = engine.compress(list(messages), current_tokens=107) + +assert compressed == messages +assert engine._last_compress_aborted is True +assert engine._last_summary_error == "invalid_replay_tool_pairs" +assert engine.compression_count == 0 +assert engine.last_compress_result["status"] == "aborted" +assert engine.last_compress_result["replay_messages"] == [] +diagnostic = engine.last_compress_result["compression_diagnostic"] +assert diagnostic["type"] == "replay_boundary_rejection" +assert diagnostic["code"] == "invalid_replay_tool_pairs" +assert {issue["tool_call_id"] for issue in diagnostic["issues"]} == {"orphan-result"} +assert engine.should_defer_preflight_to_real_usage(rough_tokens=107) is True +"#, + "generated context engine should reject orphan-heavy replay before host adoption", + ); +} + +#[test] +fn context_engine_rejects_reported_107_to_201_token_expansion() { + run_generated_plugin_script( + "check_expanding_replay_rejected.py", + r#" +messages = [{"role": "user", "content": "small live input"}] +replay = [{"role": "system", "content": "expanded replay " * 100}] + +engine = plugin.TraceDecayContextEngine() +engine.initialize(session_id="session-1", project_root="/tmp/project") +engine.last_real_prompt_tokens = 100 +engine._compress_to_result = lambda *args, **kwargs: { + "status": "ok", + "reason": "compressed_backlog", + "replay_messages": replay, + "replay_token_estimate": 201, +} + +compressed = engine.compress(list(messages), current_tokens=107) + +assert compressed == messages +assert engine._last_compress_aborted is True +assert engine.compression_count == 0 +diagnostic = engine.last_compress_result["compression_diagnostic"] +assert diagnostic["type"] == "replay_boundary_rejection" +assert diagnostic["code"] in ("non_shrinking_replay", "non_shrinking_reported_replay") +assert diagnostic["reported_source_tokens"] == 107 +assert diagnostic["reported_replay_tokens"] == 201 +assert engine.last_compress_result["replay_messages"] == [] +assert engine.should_defer_preflight_to_real_usage(rough_tokens=107) is True +"#, + "generated context engine should reject the observed 107-to-201 token expansion", + ); +} + +#[test] +fn context_engine_prefers_reported_shrink_over_host_estimator_tie() { + run_generated_plugin_script( + "check_reported_shrink_wins.py", + r#" +messages = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "second"}, +] +replay = [ + {"role": "system", "content": "short"}, + {"role": "user", "content": "tail"}, +] + +engine = plugin.TraceDecayContextEngine() +engine.initialize(session_id="session-1", project_root="/tmp/project") +engine._compress_to_result = lambda *args, **kwargs: { + "status": "ok", + "reason": "compressed_backlog", + "replay_messages": replay, + "replay_token_estimate": 3, +} + +compressed = engine.compress(list(messages), current_tokens=50) + +assert compressed == replay +assert engine._last_compress_aborted is False +assert engine.compression_count == 1 +"#, + "backend shrink estimates should override a tied secondary host estimate", + ); +} + +#[test] +fn context_engine_preserves_valid_tool_pairs_and_summary() { + run_generated_plugin_script( + "check_valid_tool_pair_replay.py", + r#" +messages = [{"role": "user", "content": "old backlog " * 300}] +replay = [ + {"role": "system", "content": "Useful compact summary"}, + { + "role": "assistant", + "content": "", + "tool_calls": [{ + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }], + }, + {"role": "tool", "tool_call_id": "call-1", "content": "bounded result"}, +] + +engine = plugin.TraceDecayContextEngine() +engine.initialize(session_id="session-1", project_root="/tmp/project") +engine._compress_to_result = lambda *args, **kwargs: { + "status": "ok", + "reason": "compressed_backlog", + "replay_messages": replay, + "replay_token_estimate": 30, +} + +compressed = engine.compress(list(messages), current_tokens=500) + +assert compressed == replay +assert compressed[0]["content"] == "Useful compact summary" +assert compressed[1]["tool_calls"][0]["id"] == "call-1" +assert compressed[2]["tool_call_id"] == "call-1" +assert engine._last_compress_aborted is False +assert engine.compression_count == 1 +"#, + "generated context engine should preserve valid summaries and paired tool events", + ); +} + +#[test] +fn context_engine_preserves_trailing_open_tool_call() { + run_generated_plugin_script( + "check_trailing_open_tool_call.py", + r#" +messages = [{"role": "user", "content": "old backlog " * 200}] +replay = [ + {"role": "system", "content": "Useful compact summary"}, + { + "role": "assistant", + "content": "", + "tool_calls": [{ + "id": "call-open", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }], + }, +] + +engine = plugin.TraceDecayContextEngine() +engine.initialize(session_id="session-1", project_root="/tmp/project") +engine._compress_to_result = lambda *args, **kwargs: { + "status": "ok", + "reason": "compressed_backlog", + "replay_messages": replay, + "replay_token_estimate": 20, +} + +compressed = engine.compress(list(messages), current_tokens=500) + +assert compressed == replay +assert compressed[-1]["tool_calls"][0]["id"] == "call-open" +assert engine._last_compress_aborted is False +assert engine.compression_count == 1 +"#, + "generated context engine should preserve a trailing call awaiting its tool result", + ); +} + #[test] fn context_engine_treats_identical_replay_as_clean_noop() { run_generated_plugin_script( @@ -3819,7 +4019,7 @@ engine = plugin.TraceDecayContextEngine(config={"context_length": 100000, "conte engine.initialize(session_id="session-1", project_root="/tmp/project") compressed = engine.compress( - [{"role": "user", "content": "compress me"}], + [{"role": "user", "content": "compress this oversized turn now"}], current_tokens=90000, force=True, ) diff --git a/tests/session_suite/lcm_compression.rs b/tests/session_suite/lcm_compression.rs index 1b1511b1..9f01c22e 100644 --- a/tests/session_suite/lcm_compression.rs +++ b/tests/session_suite/lcm_compression.rs @@ -238,6 +238,54 @@ fn lcm_raw_message(store_id: i64, role: &str, content: &str) -> LcmRawMessage { } } +fn active_multi_tool_transaction() -> Vec { + vec![ + json!({ + "id": "assistant-tools", + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call-a", + "type": "function", + "function": {"name": "lookup_a", "arguments": "{\"query\":\"alpha\"}"}, + }, + { + "id": "call-b", + "type": "function", + "function": {"name": "lookup_b", "arguments": "{\"query\":\"beta\"}"}, + }, + ], + }), + json!({ + "id": "tool-a", + "role": "tool", + "tool_call_id": "call-a", + "content": "alpha result", + }), + json!({ + "id": "tool-b", + "role": "tool", + "tool_call_id": "call-b", + "content": "beta result", + }), + ] +} + +fn active_raw_message(store_id: i64, replay: Value) -> LcmRawMessage { + let role = replay["role"].as_str().unwrap(); + let content = replay["content"].as_str().unwrap_or(""); + let mut message = lcm_raw_message(store_id, role, content); + message.metadata_json = Some( + json!({ + "lcm_active_replay": true, + "active_replay": replay, + }) + .to_string(), + ); + message +} + fn lifecycle_state_with_debt(maintenance_debt: Vec) -> LcmLifecycleState { LcmLifecycleState { provider: "cursor".into(), @@ -1517,6 +1565,238 @@ async fn raw_replay_preserves_assistant_tool_calls_and_tool_result_linking() { assert_eq!(replay_from_raw.replay_messages[1]["name"], "lookup"); } +#[tokio::test] +async fn fresh_tail_boundary_keeps_multi_tool_transaction_atomic_and_shrinking() { + let tmp = TempDir::new().unwrap(); + let db = open_lcm_db(&tmp).await; + insert_session(&db, "cursor", "session-atomic-tail").await; + + let mut messages = vec![json!({ + "id": "old-user", + "role": "user", + "content": "old backlog words ".repeat(60), + })]; + messages.extend(active_multi_tool_transaction()); + messages.push(json!({"id": "fresh-user", "role": "user", "content": "fresh prompt"})); + let mut request = limited_compress_request( + "cursor", + "session-atomic-tail", + LcmSummarizerMode::Fake { + summary_text: "compact summary".into(), + }, + None, + None, + Some(400), + ); + request.messages = messages; + request.current_tokens = Some(107); + request.threshold_tokens = Some(80); + request.fresh_tail_count = Some(3); + + let response = db.lcm_compress(request).await.unwrap(); + let assistant = response + .replay_messages + .iter() + .find(|message| message["role"] == "assistant") + .unwrap_or_else(|| { + panic!( + "assistant tool call missing: {:?}", + response.replay_messages + ) + }); + assert_eq!(assistant["tool_calls"].as_array().unwrap().len(), 2); + assert_eq!( + response + .replay_messages + .iter() + .filter(|message| message["role"] == "tool") + .map(|message| message["tool_call_id"].as_str().unwrap()) + .collect::>(), + vec!["call-a", "call-b"] + ); + assert!(response.replay_token_estimate < 107); +} + +#[tokio::test] +async fn bounded_leaf_chunk_backs_off_before_multi_tool_transaction() { + let tmp = TempDir::new().unwrap(); + let db = open_lcm_db(&tmp).await; + insert_session(&db, "cursor", "session-atomic-leaf").await; + + let mut messages = vec![json!({ + "id": "old-user", + "role": "user", + "content": "old backlog words ".repeat(30), + })]; + messages.extend(active_multi_tool_transaction()); + messages.push(json!({"id": "fresh-user", "role": "user", "content": "fresh prompt"})); + let mut request = limited_compress_request( + "cursor", + "session-atomic-leaf", + LcmSummarizerMode::Fake { + summary_text: "transaction summary".into(), + }, + None, + Some(2), + None, + ); + request.messages = messages; + request.fresh_tail_count = Some(1); + + let response = db.lcm_compress(request).await.unwrap(); + assert_eq!(response.summary_nodes_created, 1); + assert_eq!(response.summary_nodes[0].source_refs.len(), 1); + assert_eq!( + response + .replay_messages + .iter() + .filter(|message| message["role"] == "tool") + .count(), + 2 + ); +} + +#[test] +fn bounded_leaf_caps_back_off_but_progress_takes_one_atomic_unit() { + let transaction = active_multi_tool_transaction() + .into_iter() + .enumerate() + .map(|(index, replay)| active_raw_message((index + 1) as i64, replay)) + .collect::>(); + + assert_eq!(bounded_leaf_chunk_len(&transaction, Some(1), None), 0); + assert_eq!(bounded_leaf_chunk_len(&transaction, None, Some(1)), 0); + + let request = limited_compress_request( + "cursor", + "session-1", + LcmSummarizerMode::Fake { + summary_text: "summary".into(), + }, + Some(1), + Some(1), + None, + ); + let plan = compression_plan(CompressionPlanInput { + request: &request, + backlog: &transaction, + }); + assert_eq!(plan.selected_backlog.len(), transaction.len()); +} + +#[tokio::test] +async fn budget_and_overflow_replay_never_split_tool_transaction() { + let tmp = TempDir::new().unwrap(); + let db = open_lcm_db(&tmp).await; + insert_session(&db, "cursor", "session-atomic-budget").await; + + let mut messages = vec![json!({ + "id": "old-user", + "role": "user", + "content": "old backlog words ".repeat(30), + })]; + messages.extend(active_multi_tool_transaction()); + let mut request = limited_compress_request( + "cursor", + "session-atomic-budget", + LcmSummarizerMode::Fake { + summary_text: "summary".into(), + }, + None, + None, + Some(8), + ); + request.messages = messages; + request.current_tokens = Some(107); + request.fresh_tail_count = Some(3); + + let response = db.lcm_compress(request).await.unwrap(); + let assistant_count = response + .replay_messages + .iter() + .filter(|message| message["role"] == "assistant") + .count(); + let tool_count = response + .replay_messages + .iter() + .filter(|message| message["role"] == "tool") + .count(); + assert!( + (assistant_count == 0 && tool_count == 0) || (assistant_count == 1 && tool_count == 2), + "budget assembly must keep or drop the transaction as a unit: {:?}", + response.replay_messages + ); +} + +#[tokio::test] +async fn legacy_partial_tool_transaction_keeps_matched_pair_and_repairs_orphans() { + let tmp = TempDir::new().unwrap(); + let db = open_lcm_db(&tmp).await; + insert_session(&db, "cursor", "session-legacy-tools").await; + + let mut transaction = active_multi_tool_transaction(); + transaction.remove(2); + transaction.push(json!({ + "id": "separator", + "role": "user", + "content": "new prompt", + })); + transaction.push(json!({ + "id": "orphan-b", + "role": "tool", + "tool_call_id": "call-b", + "content": "late orphan", + })); + transaction.push(json!({ + "id": "unmatched-assistant", + "role": "assistant", + "content": "visible assistant text", + "tool_calls": [{ + "id": "never-returned", + "type": "function", + "function": {"name": "missing", "arguments": "{}"}, + }], + })); + transaction.push(json!({ + "id": "after-unmatched-call", + "role": "user", + "content": "the unmatched call is now a closed legacy group", + })); + let mut request = compress_request( + "cursor", + "session-legacy-tools", + LcmSummarizerMode::Fake { + summary_text: "unused".into(), + }, + ); + request.messages = transaction; + request.fresh_tail_count = Some(10); + + let response = db.lcm_compress(request).await.unwrap(); + let paired_assistant = response + .replay_messages + .iter() + .find(|message| message["id"] == "assistant-tools") + .unwrap_or_else(|| panic!("partial assistant missing: {:?}", response.replay_messages)); + assert_eq!(paired_assistant["tool_calls"].as_array().unwrap().len(), 1); + assert_eq!(paired_assistant["tool_calls"][0]["id"], "call-a"); + assert_eq!( + response + .replay_messages + .iter() + .filter(|message| message["role"] == "tool") + .map(|message| message["tool_call_id"].as_str().unwrap()) + .collect::>(), + vec!["call-a"] + ); + let repaired_assistant = response + .replay_messages + .iter() + .find(|message| message["id"] == "unmatched-assistant") + .expect("visible unmatched assistant text should remain"); + assert!(repaired_assistant.get("tool_calls").is_none()); +} + #[tokio::test] async fn active_replay_tool_calls_apply_ingest_protection_and_externalize_media_spans() { let tmp = TempDir::new().unwrap();