tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

fix(moim): freeze turn-context timestamp at turn start to preserve prefix cache (#10734)

Browse Source
This commit is contained in:
Abhijay Jain
2026-07-29 19:21:06 +05:30
committed by GitHub
parent 4a3e70c4dc
commit 465269e5d6
2 changed files with 60 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+10 -2
View File
@@ -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?,
crates/goose/src/agents/moim.rs
+50 -23
View File
@@ -36,12 +36,49 @@ pub fn system_prompt_block() -> Option<String> {
}
}
pub(super) async fn compute_compaction_info(
session_id: &str,
extension_manager: &ExtensionManager,
) -> Option<String> {
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::<f64>("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<chrono::Local>,
compaction_info: Option<String>,
) -> 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::<f64>("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<usize>) -> bool {
fn compose_moim(
working_dir: &Path,
total_tokens: Option<i32>,
context_limit: Option<usize>,
compaction_threshold: f64,
compaction_info: Option<String>,
turns_taken: u32,
max_turns: u32,
extension_parts: Vec<String>,
turn_start: chrono::DateTime<chrono::Local>,
) -> 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, &timestamp.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>,
) -> 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(),
)
}