From 465269e5d69d5fa4d6924e86477e14ea18479bb6 Mon Sep 17 00:00:00 2001 From: Abhijay Jain Date: Wed, 29 Jul 2026 19:21:06 +0530 Subject: [PATCH] fix(moim): freeze turn-context timestamp at turn start to preserve prefix cache (#10734) --- crates/goose/src/agents/agent.rs | 12 +++++- crates/goose/src/agents/moim.rs | 73 ++++++++++++++++++++++---------- 2 files changed, 60 insertions(+), 25 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 9f1878012..304e268fc 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -2008,6 +2008,11 @@ impl Agent { let mut consecutive_stop_hook_blocks = 0u32; let stop_hook_block_cap = self.stop_hook_block_cap(); let mut can_drain_pending_steers = false; + let turn_start = chrono::Local::now(); + let turn_start_compaction_info = + super::moim::compute_compaction_info(&session_config.id, &self.extension_manager) + .await; + let turn_start_turns_taken = turns_taken; loop { if is_token_cancelled(&cancel_token) { @@ -2106,9 +2111,12 @@ impl Agent { &session_config.id, conversation.clone(), &self.extension_manager, - turns_taken, + turn_start_turns_taken, max_turns, - ).await; + turn_start, + turn_start_compaction_info.clone(), + ) + .await; let mut stream = Self::stream_response_from_provider( self.provider().await?, diff --git a/crates/goose/src/agents/moim.rs b/crates/goose/src/agents/moim.rs index eb398dd99..55f22be13 100644 --- a/crates/goose/src/agents/moim.rs +++ b/crates/goose/src/agents/moim.rs @@ -36,12 +36,49 @@ pub fn system_prompt_block() -> Option { } } +pub(super) async fn compute_compaction_info( + session_id: &str, + extension_manager: &ExtensionManager, +) -> Option { + let session = extension_manager + .get_context() + .session_manager + .get_session(session_id, false) + .await + .ok(); + let session_model_config = session + .as_ref() + .and_then(|session| session.model_config.clone()); + let context_limit = if let Some(model_config) = session_model_config.as_ref() { + let provider = extension_manager.get_provider().lock().await.clone(); + match provider { + Some(provider) => provider + .get_context_limit(model_config) + .await + .ok() + .or_else(|| Some(model_config.context_limit())), + None => Some(model_config.context_limit()), + } + } else { + None + }; + let total_tokens = session + .as_ref() + .and_then(|session| session.usage.total_tokens); + let compaction_threshold = crate::config::Config::global() + .get_param::("GOOSE_AUTO_COMPACT_THRESHOLD") + .unwrap_or(crate::context_mgmt::DEFAULT_COMPACTION_THRESHOLD); + compaction_remaining_line(total_tokens, context_limit, compaction_threshold) +} + pub async fn inject_moim( session_id: &str, conversation: Conversation, extension_manager: &ExtensionManager, turns_taken: u32, max_turns: u32, + turn_start: chrono::DateTime, + compaction_info: Option, ) -> Conversation { if SKIP.with(|f| f.get()) { return conversation; @@ -77,21 +114,14 @@ pub async fn inject_moim( .as_ref() .map(|session| session.working_dir.clone()) .unwrap_or_else(|| PathBuf::from(".")); - let total_tokens = session - .as_ref() - .and_then(|session| session.usage.total_tokens); - let compaction_threshold = crate::config::Config::global() - .get_param::("GOOSE_AUTO_COMPACT_THRESHOLD") - .unwrap_or(crate::context_mgmt::DEFAULT_COMPACTION_THRESHOLD); let extension_parts = extension_manager.collect_moim_parts(session_id).await; let moim = compose_moim( &working_dir, - total_tokens, - context_limit, - compaction_threshold, + compaction_info, turns_taken, max_turns, extension_parts, + turn_start, ); let mut messages = conversation.messages().clone(); @@ -135,23 +165,20 @@ fn should_skip_moim(context_limit: Option) -> bool { fn compose_moim( working_dir: &Path, - total_tokens: Option, - context_limit: Option, - compaction_threshold: f64, + compaction_info: Option, turns_taken: u32, max_turns: u32, extension_parts: Vec, + turn_start: chrono::DateTime, ) -> String { - let timestamp = chrono::Local::now().format("%Y-%m-%d %H:%M:00 %:z"); + let timestamp = turn_start.format("%Y-%m-%d %H:%M:00 %:z"); let mut lines = vec![ open_tag(TURN_CONTEXT_TAG), tag(CURRENT_TIME_TAG, ×tamp.to_string()), tag(WORKING_DIRECTORY_TAG, &working_dir.display().to_string()), ]; - if let Some(value) = - compaction_remaining_line(total_tokens, context_limit, compaction_threshold) - { + if let Some(value) = compaction_info { lines.push(tag("compaction", &value)); } if let Some(value) = turn_budget_line(turns_taken, max_turns) { @@ -256,7 +283,7 @@ mod tests { Message::assistant().with_text("Hi"), Message::user().with_text("Bye"), ]); - let result = inject_moim(&session.id, conv, &em, 0, 100).await; + let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await; let msgs = result.messages(); assert_eq!(msgs.len(), 3); @@ -283,7 +310,7 @@ mod tests { .unwrap(); let conv = Conversation::new_unvalidated(vec![Message::user().with_text("Hello")]); - let result = inject_moim(&session.id, conv, &em, 0, 100).await; + let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await; assert_eq!(result.messages().len(), 1); assert!(is_moim(&result.messages()[0].content[0])); @@ -311,7 +338,7 @@ mod tests { Message::assistant().with_text("reply"), Message::user().with_text("user only").user_only(), ]); - let result = inject_moim(&session.id, conv, &em, 0, 100).await; + let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await; let msgs = result.messages(); assert_eq!(msgs.len(), 2); @@ -346,7 +373,7 @@ mod tests { .with_tool_response("search_1", Ok(rmcp::model::CallToolResult::success(vec![]))), ]); - let result = inject_moim(&session.id, conv, &em, 0, 100).await; + let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await; let msgs = result.messages(); assert_eq!(msgs.len(), 3); @@ -378,14 +405,14 @@ mod tests { max_turns: u32, extension_parts: Vec, ) -> String { + let compaction_info = compaction_remaining_line(total_tokens, context_limit, 0.8); compose_moim( Path::new("/Users/me/code/goose"), - total_tokens, - context_limit, - 0.8, + compaction_info, turns_taken, max_turns, extension_parts, + chrono::Local::now(), ) }