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/mcp/tools/handlers/analytics.rs b/src/mcp/tools/handlers/analytics.rs index 08d57e7f..32725074 100644 --- a/src/mcp/tools/handlers/analytics.rs +++ b/src/mcp/tools/handlers/analytics.rs @@ -442,7 +442,6 @@ async fn facts_section( } }; let query = target - .db .conn() .query( "SELECT COUNT(*), diff --git a/src/mcp/tools/handlers/memory.rs b/src/mcp/tools/handlers/memory.rs index 62ac895d..e3b3d010 100644 --- a/src/mcp/tools/handlers/memory.rs +++ b/src/mcp/tools/handlers/memory.rs @@ -28,19 +28,33 @@ use super::support::{ const DEFAULT_FACT_LIMIT: usize = 20; const MAX_FACT_LIMIT: usize = 200; -pub(super) struct TargetMemoryDb { - pub(super) db: Database, +enum TargetMemoryDbHandle<'a> { + Active(&'a Database), + Owned(Box), +} + +pub(super) struct TargetMemoryDb<'a> { + db: TargetMemoryDbHandle<'a>, pub(super) project_root: PathBuf, pub(super) user_scope: bool, } +impl TargetMemoryDb<'_> { + pub(super) fn conn(&self) -> &libsql::Connection { + match &self.db { + TargetMemoryDbHandle::Active(db) => db.conn(), + TargetMemoryDbHandle::Owned(db) => db.conn(), + } + } +} + fn requests_user_memory(args: &Value) -> bool { args.get("memory_scope").and_then(Value::as_str) == Some("user") } -async fn open_user_memory_target(profile_root: &Path) -> Result { +async fn open_user_memory_target(profile_root: &Path) -> Result> { Ok(TargetMemoryDb { - db: open_user_memory_db(profile_root).await?, + db: TargetMemoryDbHandle::Owned(Box::new(open_user_memory_db(profile_root).await?)), project_root: profile_root.to_path_buf(), user_scope: true, }) @@ -65,12 +79,12 @@ fn rendered_fact_store(project_root: Option<&Path>, args: &Value, value: &Value) text_tool_result(&text) } -pub(super) async fn open_target_memory_db( - cg: &TraceDecay, +pub(super) async fn open_target_memory_db<'a>( + cg: &'a TraceDecay, args: &Value, global_db: Option<&GlobalDb>, allow_default_registry_fallback: bool, -) -> Result { +) -> Result> { if requests_user_memory(args) { if project_selector_present(args, &["project_path"]) { return Err(config_error( @@ -88,8 +102,13 @@ pub(super) async fn open_target_memory_db( ) .await? else { + let db = if cg.db_path() == cg.store_layout().graph_db_path { + TargetMemoryDbHandle::Active(cg.db()) + } else { + TargetMemoryDbHandle::Owned(Box::new(cg.open_project_store_db().await?)) + }; return Ok(TargetMemoryDb { - db: cg.open_project_store_db().await?, + db, project_root: cg.project_root().to_path_buf(), user_scope: false, }); @@ -116,7 +135,7 @@ pub(super) async fn open_target_memory_db( } let (db, _) = Database::open(&db_path).await?; Ok(TargetMemoryDb { - db, + db: TargetMemoryDbHandle::Owned(Box::new(db)), project_root: PathBuf::from(context.project.display_root), user_scope: false, }) @@ -317,10 +336,10 @@ pub(super) async fn handle_fact_store( async fn handle_fact_store_for_target( args: Value, cross_project_selector: bool, - target_memory: TargetMemoryDb, + target_memory: TargetMemoryDb<'_>, ) -> Result { let action = required_str(&args, "action")?; - let conn = target_memory.db.conn(); + let conn = target_memory.conn(); let store = MemoryStore::new(conn); let mut refresh_digest = false; let out = match action { @@ -548,7 +567,7 @@ pub(super) async fn handle_fact_feedback( .map(ToOwned::to_owned); let target_memory = open_target_memory_db(cg, &args, global_db, allow_default_registry_fallback).await?; - let result = MemoryStore::new(target_memory.db.conn()) + let result = MemoryStore::new(target_memory.conn()) .record_feedback_event(FeedbackRequest { fact_id: fact_id(&args)?, action: feedback_action(&args)?, @@ -561,7 +580,7 @@ pub(super) async fn handle_fact_feedback( .await?; if !target_memory.user_scope { refresh_memory_digest_after_memory_change( - target_memory.db.conn(), + target_memory.conn(), &target_memory.project_root, ) .await; @@ -582,7 +601,7 @@ pub(super) async fn handle_memory_status( ) -> Result { let target_memory = open_target_memory_db(cg, &args, global_db, allow_default_registry_fallback).await?; - let status = TraceDecay::memory_status_for_conn(target_memory.db.conn()).await?; + let status = TraceDecay::memory_status_for_conn(target_memory.conn()).await?; let value = json!({ "status": "ok", "memory": status }); Ok(rendered_tool_json( (!target_memory.user_scope).then_some(target_memory.project_root.as_path()), @@ -618,7 +637,7 @@ pub async fn handle_user_memory_tool( .or_else(|| args.get("reason")) .and_then(Value::as_str) .map(ToOwned::to_owned); - let result = MemoryStore::new(target_memory.db.conn()) + let result = MemoryStore::new(target_memory.conn()) .record_feedback_event(FeedbackRequest { fact_id: fact_id(&args)?, action: feedback_action(&args)?, @@ -636,7 +655,7 @@ pub async fn handle_user_memory_tool( )) } "tracedecay_memory_status" => { - let status = TraceDecay::memory_status_for_conn(target_memory.db.conn()).await?; + let status = TraceDecay::memory_status_for_conn(target_memory.conn()).await?; Ok(rendered_tool_json( None, &args, @@ -646,3 +665,32 @@ pub async fn handle_user_memory_tool( other => Err(config_error(format!("{other} is not a user-memory tool"))), } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn active_project_memory_uses_the_served_database_handle() { + let tmp = tempfile::tempdir().unwrap(); + let project_root = tmp.path().join("project"); + let profile_root = tmp.path().join("profile"); + std::fs::create_dir_all(&project_root).unwrap(); + let cg = TraceDecay::init_with_options( + &project_root, + crate::tracedecay::TraceDecayOpenOptions { + global_db_path: Some(profile_root.join("global.db")), + profile_root: Some(profile_root), + }, + ) + .await + .unwrap(); + + let target = open_target_memory_db(&cg, &json!({}), None, true) + .await + .unwrap(); + + assert!(matches!(target.db, TargetMemoryDbHandle::Active(_))); + assert!(std::ptr::eq(target.conn(), cg.db().conn())); + } +} diff --git a/src/migrate/consolidate/files.rs b/src/migrate/consolidate/files.rs index 8fb8f589..6a1b4f32 100644 --- a/src/migrate/consolidate/files.rs +++ b/src/migrate/consolidate/files.rs @@ -56,13 +56,6 @@ pub(super) fn tree_stats(root: &Path) -> Result<(usize, u64)> { Ok((files.len(), bytes)) } -pub(super) fn copy_tree_exact(source: &Path, target: &Path) -> Result<()> { - for (relative, path) in relative_file_map(source)? { - copy_file_exact(&path, &target.join(relative))?; - } - Ok(()) -} - pub(super) fn copy_file_exact(source: &Path, target: &Path) -> Result<()> { if target.exists() { if file_digest(source)? == file_digest(target)? { @@ -98,13 +91,33 @@ pub(super) fn copy_file_atomic(source: &Path, target: &Path) -> Result<()> { Err(error) if error.kind() == io::ErrorKind::NotFound => {} Err(error) => return Err(io_error(error)), } - let mut input = File::open(source).map_err(io_error)?; + let mut input = File::open(source).map_err(|error| { + config_error(format!( + "failed to open migration source '{}' for copy to '{}': {error}", + source.display(), + target.display() + )) + })?; let mut output = OpenOptions::new() .write(true) .create_new(true) .open(&temp) - .map_err(io_error)?; - io::copy(&mut input, &mut output).map_err(io_error)?; + .map_err(|error| { + config_error(format!( + "failed to create migration temp '{}' while copying '{}' to '{}': {error}", + temp.display(), + source.display(), + target.display() + )) + })?; + io::copy(&mut input, &mut output).map_err(|error| { + config_error(format!( + "failed to copy migration source '{}' to temp '{}' for '{}': {error}", + source.display(), + temp.display(), + target.display() + )) + })?; fs::set_permissions(&temp, fs::metadata(source).map_err(io_error)?.permissions()) .map_err(io_error)?; output.sync_all().map_err(io_error)?; @@ -183,9 +196,14 @@ pub(super) fn excluded_source_artifact(relative: &Path) -> bool { } pub(super) fn is_runtime_lock(relative: &Path) -> bool { + is_coordination_lock(relative) + || relative.file_name().and_then(|value| value.to_str()) == Some(".dirty") +} + +pub(super) fn is_coordination_lock(relative: &Path) -> bool { matches!( relative.file_name().and_then(|value| value.to_str()), - Some("sync.lock" | ".branch-add.lock" | ".dirty") + Some("sync.lock" | ".branch-add.lock") ) } diff --git a/src/migrate/consolidate/mod.rs b/src/migrate/consolidate/mod.rs index 47e798e8..a76fa004 100644 --- a/src/migrate/consolidate/mod.rs +++ b/src/migrate/consolidate/mod.rs @@ -27,9 +27,9 @@ use evidence::{GraphStoreEvidence, InputReadEvidence, capture_input_evidence}; #[cfg(test)] use files::sqlite_sidecar; use files::{ - copy_file_atomic, copy_sqlite_family_exact, copy_tree_exact, excluded_source_artifact, - file_digest, is_reference_artifact, is_runtime_lock, is_sqlite_database, is_sqlite_sidecar, - relative_file_map, tree_stats, + copy_file_atomic, copy_file_exact, copy_sqlite_family_exact, excluded_source_artifact, + file_digest, is_coordination_lock, is_reference_artifact, is_runtime_lock, is_sqlite_database, + is_sqlite_sidecar, relative_file_map, tree_stats, }; use finalize::{cut_over_markers, register_destination, verify_destination}; use preflight::{acquire_store_locks, ensure_profile_offline, preflight_disk_space}; @@ -166,8 +166,8 @@ struct ResolvedPlan { target_layout: StoreLayout, source_meta: BranchMeta, target_meta: BranchMeta, - scratch_root: MigrationScratchRoot, evidence: Arc, + scratch_root: MigrationScratchRoot, } static NEXT_MIGRATION_SCRATCH: AtomicU64 = AtomicU64::new(0); @@ -685,6 +685,7 @@ fn recover_untracked_branch_graphs(layout: &StoreLayout, meta: &mut BranchMeta) path.display() )) })?; + let db_file = db_file.replace('\\', "/"); let base = relative .file_stem() .and_then(|value| value.to_str()) @@ -703,7 +704,7 @@ fn recover_untracked_branch_graphs(layout: &StoreLayout, meta: &mut BranchMeta) meta.branches.insert( name, BranchEntry { - db_file: db_file.to_string(), + db_file, parent: Some(meta.default_branch.clone()), created_at: "0".to_string(), last_synced_at: "0".to_string(), @@ -1438,7 +1439,14 @@ fn read_optional_regular_file(path: &Path) -> Result>> { fn backup_store(layout: &StoreLayout, backup_root: &Path) -> Result<()> { let project_id = layout.identity.project_id.as_deref().unwrap_or("unknown"); - copy_tree_exact(&layout.data_root, &backup_root.join(project_id)) + let destination = backup_root.join(project_id); + for (relative, path) in relative_file_map(&layout.data_root)? { + if is_coordination_lock(&relative) { + continue; + } + copy_file_exact(&path, &destination.join(relative))?; + } + Ok(()) } async fn merge_databases(resolved: &ResolvedPlan, ledger: &mut ConsolidationLedger) -> Result<()> { diff --git a/src/migrate/consolidate/prepare.rs b/src/migrate/consolidate/prepare.rs index 0f4c0ade..95b45174 100644 --- a/src/migrate/consolidate/prepare.rs +++ b/src/migrate/consolidate/prepare.rs @@ -86,7 +86,7 @@ fn preserve_source_branch_graphs( merged.branches.insert( merged_name, BranchEntry { - db_file: relative.to_string_lossy().to_string(), + db_file: relative.to_string_lossy().replace('\\', "/"), parent: entry .parent .as_deref() @@ -171,7 +171,7 @@ fn expected_prepared_files( meta.branches.insert( merged_name, BranchEntry { - db_file: relative.to_string_lossy().to_string(), + db_file: relative.to_string_lossy().replace('\\', "/"), parent: entry .parent .as_deref() 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/src/sqlite_read_snapshot.rs b/src/sqlite_read_snapshot.rs index 1e4ed37d..05ab3ad6 100644 --- a/src/sqlite_read_snapshot.rs +++ b/src/sqlite_read_snapshot.rs @@ -164,11 +164,12 @@ enum SnapshotMode { struct ScratchDirectory { path: PathBuf, - _owner_lock: File, + owner_lock: Option, } impl Drop for ScratchDirectory { fn drop(&mut self) { + drop(self.owner_lock.take()); let _ = fs::remove_dir_all(&self.path); } } @@ -278,7 +279,7 @@ fn prepare_one( SnapshotMode::Copy } } else { - SnapshotMode::DirectImmutable + checkpointed_snapshot_mode() }; let mut copy_bytes = if matches!(mode, SnapshotMode::Copy) { main.bytes @@ -308,6 +309,20 @@ fn prepare_one( }) } +fn checkpointed_snapshot_mode() -> SnapshotMode { + // SQLite's immutable connection still holds a byte-range lock on Windows. + // Consolidation retains read snapshots while copying the frozen inputs, so + // opening a private copy keeps those handles off the source database. + #[cfg(windows)] + { + SnapshotMode::Copy + } + #[cfg(not(windows))] + { + SnapshotMode::DirectImmutable + } +} + async fn finish_one( prepared: PreparedSnapshot, scratch: Arc, @@ -393,7 +408,7 @@ fn create_scratch_directory( FileExt::unlock(&cleanup_lock)?; return Ok(ScratchDirectory { path, - _owner_lock: owner_lock, + owner_lock: Some(owner_lock), }); } Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {} @@ -631,6 +646,7 @@ mod tests { assert_eq!(family_state(&path).unwrap(), before); } + #[cfg(not(windows))] #[tokio::test] async fn checkpointed_database_reads_directly_without_copy_or_metadata_change() { let temp = TempDir::new().unwrap(); @@ -670,6 +686,27 @@ mod tests { assert_eq!(family_state(&path).unwrap(), before); } + #[cfg(windows)] + #[tokio::test] + async fn checkpointed_snapshot_does_not_lock_source_against_copying() { + let temp = TempDir::new().unwrap(); + let path = temp.path().join("source.db"); + let database = Builder::new_local(&path).build().await.unwrap(); + database + .connect() + .unwrap() + .execute_batch("CREATE TABLE durable(value TEXT NOT NULL);") + .await + .unwrap(); + drop(database); + + let snapshots = SnapshotSet::capture(std::slice::from_ref(&path)) + .await + .unwrap(); + assert_eq!(snapshots.copied_bytes(), fs::metadata(&path).unwrap().len()); + fs::copy(&path, temp.path().join("backup.db")).unwrap(); + } + #[test] fn empty_wal_does_not_change_the_content_fingerprint() { let temp = TempDir::new().unwrap(); 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..834fded1 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( @@ -2311,9 +2511,15 @@ def fake_call_tracedecay_tool(name, args, **kwargs): return '{"content":[{"type":"text","text":"{\\"status\\":\\"ok\\"}"}]}' plugin.tools.call_tracedecay_tool = fake_call_tracedecay_tool -plugin._project_scope_resolution = lambda path, *_args: ( - ("registered", str(path)) if str(path).startswith("/repos/") else ("unregistered", None) -) +def fake_project_scope_resolution(path, *_args): + normalized = str(path).replace("\\", "/") + marker = "/repos/" + marker_index = normalized.find(marker) + if marker_index < 0: + return ("unregistered", None) + return ("registered", normalized[marker_index:]) + +plugin._project_scope_resolution = fake_project_scope_resolution plugin._notify_turn_completed = lambda sid, root, watermark: completed.append((sid, root, watermark)) plugin._notify_turn_ingested = lambda sid, root, watermark: ingested.append((sid, root, watermark)) @@ -3819,7 +4025,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/hooks_lsp_suite/lsp_code_diagnostics_test.rs b/tests/hooks_lsp_suite/lsp_code_diagnostics_test.rs index 72f75462..882854ac 100644 --- a/tests/hooks_lsp_suite/lsp_code_diagnostics_test.rs +++ b/tests/hooks_lsp_suite/lsp_code_diagnostics_test.rs @@ -16,6 +16,10 @@ const FAKE_LSP_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(1 // recovery helper keeps the diagnostics quiet window small so success stays // cheap. const FAKE_LSP_RECOVERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3); +// macOS full-suite runs can briefly starve a newly spawned /usr/bin/python3 +// while other test binaries are starting. Keep that harness-only startup +// allowance independent from the deliberately short diagnostics quiet window. +const FAKE_LSP_START_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); const OUTER_ASYNC_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(6); const HANGING_WRITE_LINE_COUNT: usize = 64_000; @@ -37,6 +41,15 @@ fn recovery_fake_lsp_timeouts() -> lsp::client::LspRefreshTimeouts { ) } +fn loaded_runner_fake_lsp_timeouts() -> lsp::client::LspRefreshTimeouts { + lsp::client::LspRefreshTimeouts::new( + FAKE_LSP_TIMEOUT, + FAKE_LSP_START_TIMEOUT, + FAKE_LSP_START_TIMEOUT, + FAKE_LSP_TIMEOUT, + ) +} + #[test] fn builtin_registry_advertises_phase_one_setup_contract() { let adapters = lsp::adapters::builtin_adapters(); @@ -238,18 +251,18 @@ async fn broker_keeps_diagnostics_for_multiple_languages_in_one_snapshot() { ); broker - .refresh_documents( + .refresh_documents_with_timeouts( "alpha", vec![fake_document("alpha", "src/lib.alpha", "alpha nope")], - FAKE_LSP_TIMEOUT, + loaded_runner_fake_lsp_timeouts(), ) .await .unwrap(); broker - .refresh_documents( + .refresh_documents_with_timeouts( "beta", vec![fake_document("beta", "src/lib.beta", "beta nope")], - FAKE_LSP_TIMEOUT, + loaded_runner_fake_lsp_timeouts(), ) .await .unwrap(); @@ -743,11 +756,15 @@ async fn broker_ignores_stale_refresh_completion() { .unwrap() .expect("latest refresh should prepare"); - let latest_completed = latest.collect_diagnostics(FAKE_LSP_TIMEOUT).await; + let latest_completed = latest + .collect_diagnostics_with_timeouts(loaded_runner_fake_lsp_timeouts()) + .await; assert!(latest_completed.is_ok()); broker.finish_refresh(latest_completed).unwrap(); - let stale_completed = stale.collect_diagnostics(FAKE_LSP_TIMEOUT).await; + let stale_completed = stale + .collect_diagnostics_with_timeouts(loaded_runner_fake_lsp_timeouts()) + .await; assert!(stale_completed.is_ok()); broker.finish_refresh(stale_completed).unwrap(); @@ -825,11 +842,15 @@ async fn broker_keys_warm_lsp_clients_by_workspace_root() { ]; broker - .refresh_documents("fake", documents.clone(), FAKE_LSP_TIMEOUT) + .refresh_documents_with_timeouts( + "fake", + documents.clone(), + loaded_runner_fake_lsp_timeouts(), + ) .await .unwrap(); broker - .refresh_documents("fake", documents, FAKE_LSP_TIMEOUT) + .refresh_documents_with_timeouts("fake", documents, loaded_runner_fake_lsp_timeouts()) .await .unwrap(); 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();