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

Cache-safe request assembly: append-only turn context and declared cache semantics (#11022)

Browse Source
This commit is contained in:
filip
2026-08-06 20:08:49 -07:00
committed by GitHub
parent ca52cce628
commit 66051ec7d2
39 changed files with 1589 additions and 1231 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/acp/provider.rs
+34 -2
View File
@@ -1442,14 +1442,14 @@ fn messages_to_prompt(messages: &[Message], include_handoff_context: bool) -> Ve
fn last_user_message_index(messages: &[Message]) -> Option<usize> {
messages
.iter()
.rposition(|m| m.role == Role::User && m.is_agent_visible())
.rposition(|m| m.role == Role::User && m.is_agent_visible() && !m.is_turn_context())
}
fn has_handoff_context(messages: &[Message]) -> bool {
last_user_message_index(messages).is_some_and(|last_user_index| {
messages[..last_user_index]
.iter()
.any(Message::is_agent_visible)
.any(|m| m.is_agent_visible() && !m.is_turn_context())
})
}
@@ -1458,6 +1458,7 @@ fn build_handoff_context_memo(prior_messages: &[Message]) -> Option<String> {
Conversation::new_unvalidated(prior_messages.iter().cloned())
.agent_visible_messages()
.iter()
.filter(|message| !message.is_turn_context())
.map(|message| format_message_for_compacting(&message.agent_visible_content()))
.collect();
@@ -1810,6 +1811,37 @@ mod tests {
assert_eq!(prompt_text(&blocks[1]), "continue from there");
}
#[test]
fn messages_to_prompt_skips_turn_context_events() {
use crate::conversation::message::MessageMetadata;
let turn_context = |text: &str| {
Message::user()
.with_text(text)
.with_metadata(MessageMetadata::agent_only().with_turn_context())
};
let messages = vec![
Message::user().with_text("inspect src/lib.rs"),
turn_context("<turn-context>old cwd /repo</turn-context>"),
Message::assistant().with_text("I found the file"),
Message::user().with_text("continue from there"),
turn_context("<turn-context>new cwd /repo</turn-context>"),
];
let blocks = messages_to_prompt(&messages, true);
assert_eq!(blocks.len(), 2);
assert!(
!prompt_text(&blocks[0]).contains("turn-context"),
"handoff memo must not include turn-context events"
);
assert_eq!(
prompt_text(&blocks[1]),
"continue from there",
"the current prompt must be the user's request, not a trailing turn-context event"
);
}
#[test]
fn messages_to_prompt_drops_user_only_acp_rows_from_handoff() {
let user_only = TextContent::new("SECRET_USER_ONLY")
crates/goose/src/agents/agent.rs
+57 -20
View File
@@ -46,7 +46,9 @@ use crate::conversation::message::{
ActionRequiredData, InferenceMetadata, Message, MessageContent, MessageUsage, ProviderMetadata,
SystemNotificationType, ToolRequest,
};
use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation};
use crate::conversation::{
debug_conversation_fix, fix_conversation, merge_consecutive_messages_for_request, Conversation,
};
use crate::mcp_utils::ToolResult;
use crate::permission::permission_inspector::PermissionInspector;
use crate::permission::permission_judge::PermissionCheckResult;
@@ -150,7 +152,6 @@ pub struct ReplyContext {
pub system_prompt: String,
pub goose_mode: GooseMode,
pub tool_call_cut_off: usize,
pub initial_messages: Vec<Message>,
pub model_config: goose_providers::model::ModelConfig,
}
@@ -787,8 +788,6 @@ impl Agent {
)
);
}
let initial_messages = conversation.messages().clone();
let (tools, toolshim_tools, system_prompt, model_config) = self
.prepare_tools_and_prompt(session_id, working_dir)
.await?;
@@ -820,7 +819,6 @@ impl Agent {
system_prompt,
goose_mode,
tool_call_cut_off,
initial_messages,
model_config,
})
}
@@ -2154,7 +2152,6 @@ impl Agent {
mut system_prompt,
tool_call_cut_off,
goose_mode,
initial_messages,
model_config,
} = context;
@@ -2264,7 +2261,27 @@ impl Agent {
let turn_start_compaction_info =
super::moim::compute_compaction_info(&session_config.id, &self.extension_manager)
.await;
let turn_start_turns_taken = turns_taken;
if let Some(turn_context) = super::moim::turn_context_message(
&session_config.id,
&self.extension_manager,
turns_taken,
max_turns,
turn_start,
turn_start_compaction_info,
)
.await
{
persist_and_push_message_with_id(
&session_manager,
&session_config.id,
&mut conversation,
turn_context,
)
.await?;
}
// Snapshot after the turn-context append so a retry keeps the sent prefix.
let initial_messages = conversation.messages().clone();
loop {
if is_token_cancelled(&cancel_token) {
@@ -2359,23 +2376,12 @@ impl Agent {
break;
}
let conversation_with_moim = super::moim::inject_moim(
&session_config.id,
conversation.clone(),
&self.extension_manager,
turn_start_turns_taken,
max_turns,
turn_start,
turn_start_compaction_info.clone(),
)
.await;
let mut stream = crate::agents::reply_parts::stream_response_from_provider(
self.provider().await?,
model_config.clone(),
&session_config.id,
&system_prompt,
conversation_with_moim.messages(),
conversation.messages(),
&tools,
&toolshim_tools,
).await?;
@@ -3734,7 +3740,7 @@ impl Agent {
.filter(super::reply_parts::is_tool_visible_to_model)
.collect();
messages = Conversation::new_unvalidated(messages.agent_visible_messages());
messages = Conversation::new_unvalidated(recipe_conversation_history(&messages));
messages.push(Message::user().with_text(recipe_prompt));
let (messages, issues) = fix_conversation(messages);
@@ -3743,6 +3749,9 @@ impl Agent {
.iter()
.for_each(|issue| tracing::warn!(recipe.conversation.issue = issue));
}
let messages = Conversation::new_unvalidated(merge_consecutive_messages_for_request(
messages.messages().clone(),
));
tracing::debug!(
"Added recipe prompt to messages, total messages: {}",
@@ -3908,6 +3917,15 @@ impl Agent {
}
}
fn recipe_conversation_history(messages: &Conversation) -> Vec<Message> {
// The recipe prompt has no turn-context instructions; drop the blocks.
messages
.agent_visible_messages()
.into_iter()
.filter(|message| !message.is_turn_context())
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
@@ -3923,6 +3941,25 @@ mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use tempfile::TempDir;
#[test]
fn recipe_history_excludes_turn_context_events() {
use crate::conversation::message::MessageMetadata;
let history = Conversation::new_unvalidated([
Message::user().with_text("build me a recipe"),
Message::user()
.with_text("<turn-context>cwd /repo</turn-context>")
.with_metadata(MessageMetadata::agent_only().with_turn_context()),
Message::assistant().with_text("on it"),
]);
let texts: Vec<String> = recipe_conversation_history(&history)
.iter()
.map(|message| message.as_concat_text())
.collect();
assert_eq!(texts, ["build me a recipe", "on it"]);
}
async fn tracing_test_agent_and_session() -> (Agent, Session, TempDir) {
let data_dir = TempDir::new().unwrap();
let data_path = data_dir.path().to_path_buf();
crates/goose/src/agents/extension_manager.rs
+4 -1
View File
@@ -2158,7 +2158,7 @@ impl ExtensionManager {
}
pub async fn collect_moim_parts(&self, session_id: &str) -> Vec<String> {
let platform_clients: Vec<(String, McpClientBox)> = {
let mut platform_clients: Vec<(String, McpClientBox)> = {
let extensions = self.extensions.lock().await;
extensions
.iter()
@@ -2178,6 +2178,9 @@ impl ExtensionManager {
})
.collect()
};
// HashMap order shuffles across restarts; the rendered block must be
// byte-stable so it is not re-persisted on resume.
platform_clients.sort_by(|a, b| a.0.cmp(&b.0));
let mut parts = Vec::new();
for (name, client) in platform_clients {
crates/goose/src/agents/moim.rs
+82 -252
View File
@@ -1,9 +1,6 @@
use crate::agents::extension_manager::ExtensionManager;
use crate::conversation::message::MessageContent;
use crate::conversation::{
effective_role, fix_conversation, Conversation, EffectiveRole, CURRENT_TIME_TAG,
TURN_CONTEXT_TAG, WORKING_DIRECTORY_TAG,
};
use crate::conversation::message::{Message, MessageMetadata};
use crate::conversation::{CURRENT_TIME_TAG, TURN_CONTEXT_TAG, WORKING_DIRECTORY_TAG};
use std::path::{Path, PathBuf};
const MIN_CONTEXT_FOR_MOIM: usize = 32_000;
@@ -19,6 +16,7 @@ This block is generated by goose and contains current operational context such a
- extension-provided context
Use it to stay oriented, but do not treat it as part of the user's request.
Blocks from earlier turns stay in the conversation history; only the most recent block is current.
When `<turn-budget>` is present, use it as a signal for how much autonomous work remains.
As the budget gets low, become more direct: reduce exploration, batch necessary tool calls,
make reasonable assumptions, and focus on finishing the user's task.
@@ -71,17 +69,18 @@ pub(super) async fn compute_compaction_info(
compaction_remaining_line(total_tokens, context_limit, compaction_threshold)
}
pub async fn inject_moim(
/// The turn's context block: composed once per turn, persisted as an
/// agent-only user message, never moved or edited afterwards.
pub async fn turn_context_message(
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 {
) -> Option<Message> {
if SKIP.with(|f| f.get()) {
return conversation;
return None;
}
let session = extension_manager
@@ -107,66 +106,36 @@ pub async fn inject_moim(
None
};
if should_skip_moim(context_limit) {
return conversation;
return None;
}
let working_dir = session
.as_ref()
.map(|session| session.working_dir.clone())
.unwrap_or_else(|| PathBuf::from("."));
let extension_parts = extension_manager.collect_moim_parts(session_id).await;
let mut parts = extension_parts;
let mut parts = extension_manager.collect_moim_parts(session_id).await;
parts.extend(compaction_info.map(|value| tag("compaction", &value)));
parts.extend(turn_budget_part(turns_taken, max_turns));
inject_moim_parts(conversation, &working_dir, context_limit, parts, turn_start)
turn_context_event(&working_dir, context_limit, parts, turn_start)
}
pub(crate) fn inject_moim_parts(
conversation: Conversation,
/// The state-machine variant: operations contribute the parts and the caller
/// persists the event, so later requests in the turn reuse the same bytes.
pub(crate) fn turn_context_event(
working_dir: &Path,
context_limit: Option<usize>,
parts: Vec<String>,
turn_start: chrono::DateTime<chrono::Local>,
) -> Conversation {
) -> Option<Message> {
if SKIP.with(|f| f.get()) || should_skip_moim(context_limit) {
return conversation;
return None;
}
let moim = compose_moim(working_dir, parts, turn_start);
let mut messages = conversation.messages().clone();
let Some(idx) = messages
.iter()
.rposition(|m| m.is_agent_visible() && effective_role(m) == EffectiveRole::User)
else {
return conversation;
};
let insert_idx = messages[idx]
.content
.iter()
.take_while(|content| matches!(content, MessageContent::ToolResponse(_)))
.count();
messages[idx]
.content
.insert(insert_idx, MessageContent::text(moim));
let (fixed, issues) = fix_conversation(Conversation::new_unvalidated(messages));
let has_unexpected_issues = issues.iter().any(|issue| {
!issue.contains("Merged consecutive user messages")
&& !issue.contains("Merged consecutive assistant messages")
&& !issue.contains("Added placeholder to empty tool result")
&& !issue.contains("Trimmed trailing whitespace from assistant message")
&& !issue.contains("Removed trailing assistant message")
&& !issue.contains("Merged text content")
});
if has_unexpected_issues {
tracing::warn!("MOIM injection caused unexpected issues: {:?}", issues);
return conversation;
}
fixed
Some(
Message::user()
.with_text(compose_moim(working_dir, parts, turn_start))
.with_metadata(MessageMetadata::agent_only().with_turn_context()),
)
}
fn should_skip_moim(context_limit: Option<usize>) -> bool {
@@ -252,21 +221,8 @@ fn turn_budget_part(turns_taken: u32, max_turns: u32) -> Option<String> {
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::message::Message;
use rmcp::model::CallToolRequestParams;
fn text_at(message: &crate::conversation::message::Message, index: usize) -> &str {
message.content[index].as_text().unwrap()
}
fn is_moim(content: &MessageContent) -> bool {
content
.as_text()
.is_some_and(|text| text.starts_with(&format!("<{}>\n", TURN_CONTEXT_TAG)))
}
#[tokio::test]
async fn test_moim_prepended_to_latest_user_message() {
async fn session_and_manager() -> (String, ExtensionManager, tempfile::TempDir) {
let temp_dir = tempfile::tempdir().unwrap();
let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
let session = em
@@ -280,209 +236,83 @@ mod tests {
)
.await
.unwrap();
let conv = Conversation::new_unvalidated(vec![
Message::user().with_text("Hello"),
Message::assistant().with_text("Hi"),
Message::user().with_text("Bye"),
]);
let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await;
let msgs = result.messages();
assert_eq!(msgs.len(), 3);
assert_eq!(text_at(&msgs[0], 0), "Hello");
assert_eq!(text_at(&msgs[1], 0), "Hi");
assert!(is_moim(&msgs[2].content[0]));
assert_eq!(text_at(&msgs[2], 1), "Bye");
(session.id, em, temp_dir)
}
#[tokio::test]
async fn test_moim_injection_no_assistant() {
let temp_dir = tempfile::tempdir().unwrap();
let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
let session = em
.get_context()
.session_manager
.create_session(
PathBuf::from("/test/dir"),
"test".to_string(),
crate::session::SessionType::User,
crate::config::GooseMode::Auto,
)
async fn turn_context_message_is_an_agent_only_user_message() {
let (session_id, em, _tmp) = session_and_manager().await;
let message = turn_context_message(&session_id, &em, 0, 100, chrono::Local::now(), None)
.await
.unwrap();
.expect("turn context should be produced");
let conv = Conversation::new_unvalidated(vec![Message::user().with_text("Hello")]);
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]));
assert_eq!(text_at(&result.messages()[0], 1), "Hello");
assert_eq!(message.role, rmcp::model::Role::User);
assert!(message.is_agent_visible());
assert!(!message.is_user_visible());
assert!(message.is_turn_context());
let text = message.content[0].as_text().unwrap();
assert!(text.starts_with(&format!("<{TURN_CONTEXT_TAG}>\n")));
assert!(text.trim_end().ends_with(&format!("</{TURN_CONTEXT_TAG}>")));
}
#[tokio::test]
async fn test_moim_skips_user_messages_not_visible_to_agent() {
let temp_dir = tempfile::tempdir().unwrap();
let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
let session = em
.get_context()
.session_manager
.create_session(
PathBuf::from("/test/dir"),
"test".to_string(),
crate::session::SessionType::User,
crate::config::GooseMode::Auto,
)
async fn turn_context_bytes_are_stable_for_a_turn() {
let (session_id, em, _tmp) = session_and_manager().await;
let turn_start = chrono::Local::now();
let first = turn_context_message(&session_id, &em, 0, 100, turn_start, None)
.await
.unwrap();
let second = turn_context_message(&session_id, &em, 0, 100, turn_start, None)
.await
.unwrap();
let conv = Conversation::new_unvalidated(vec![
Message::user().with_text("agent visible"),
Message::assistant().with_text("reply"),
Message::user().with_text("user only").user_only(),
]);
let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await;
let msgs = result.messages();
assert_eq!(msgs.len(), 2);
assert!(is_moim(&msgs[0].content[0]));
assert_eq!(text_at(&msgs[0], 1), "agent visible");
assert_eq!(text_at(&msgs[1], 0), "user only");
assert!(!msgs[1].is_agent_visible());
assert_eq!(
first.content[0].as_text(),
second.content[0].as_text(),
"the same turn inputs must render byte-identical blocks"
);
}
#[tokio::test]
async fn test_moim_with_tool_calls() {
let temp_dir = tempfile::tempdir().unwrap();
let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
let session = em
.get_context()
.session_manager
.create_session(
PathBuf::from("/test/dir"),
"test".to_string(),
crate::session::SessionType::User,
crate::config::GooseMode::Auto,
)
.await
.unwrap();
#[test]
fn turn_context_event_is_marked_and_skips_small_contexts() {
let event = turn_context_event(
Path::new("/test/dir"),
Some(200_000),
vec![],
chrono::Local::now(),
)
.expect("large-context models should get a turn-context event");
assert!(event.is_turn_context());
assert!(!event.is_user_visible());
let conv = Conversation::new_unvalidated(vec![
Message::user().with_text("Search for something"),
Message::assistant()
.with_text("I'll search for you")
.with_tool_request("search_1", Ok(CallToolRequestParams::new("search"))),
Message::user()
.with_tool_response("search_1", Ok(rmcp::model::CallToolResult::success(vec![]))),
]);
let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await;
let msgs = result.messages();
assert_eq!(msgs.len(), 3);
assert!(is_moim(&msgs[0].content[0]));
assert_eq!(text_at(&msgs[0], 1), "Search for something");
assert!(matches!(
&msgs[2].content[0],
MessageContent::ToolResponse(_)
));
assert_eq!(msgs[2].content.len(), 1);
assert!(turn_context_event(
Path::new("/test/dir"),
Some(16_000),
vec![],
chrono::Local::now(),
)
.is_none());
}
/// The turn-context block is produced here (`compose_moim`, in goose) but
/// recognized in a different crate (`is_turn_context_text`, in
/// goose-providers) by matching its textual shape. That recognition is the
/// only thing that keeps the volatile block out of the Anthropic cache
/// prefix. The two are coupled by string layout alone, with no shared type
/// or compiler check, so a reorder or rename on either side would silently
/// collapse prompt caching. These tests fail loudly if they drift apart.
mod turn_context_detector_coupling {
use super::*;
use crate::conversation::is_turn_context_text;
use std::path::Path;
#[test]
fn current_time_is_timezone_aware() {
let block = compose_moim(
Path::new("/Users/me/code/goose"),
vec![],
chrono::Local::now(),
);
let time_line = block
.lines()
.find(|line| line.starts_with(&format!("<{CURRENT_TIME_TAG}>")))
.expect("turn-context block should contain a current-time line");
fn moim(
total_tokens: Option<i32>,
context_limit: Option<usize>,
turns_taken: u32,
max_turns: u32,
mut parts: Vec<String>,
) -> String {
parts.extend(
compaction_remaining_line(total_tokens, context_limit, 0.8)
.map(|value| tag("compaction", &value)),
);
parts.extend(turn_budget_part(turns_taken, max_turns));
compose_moim(
Path::new("/Users/me/code/goose"),
parts,
chrono::Local::now(),
)
}
let value = time_line
.trim_start_matches(&format!("<{CURRENT_TIME_TAG}>"))
.trim_end_matches(&format!("</{CURRENT_TIME_TAG}>"));
#[test]
fn compose_moim_output_is_recognized_by_the_anthropic_detector() {
let guardrail = || vec!["<guardrail>stay on task</guardrail>".to_string()];
let cases = [
("minimal", moim(None, None, 0, 0, vec![])),
(
"with compaction line",
moim(Some(100_000), Some(200_000), 0, 0, vec![]),
),
("with turn-budget line", moim(None, None, 20, 40, vec![])),
(
"with extension context",
moim(None, None, 0, 0, guardrail()),
),
(
"fully populated",
moim(Some(100_000), Some(200_000), 20, 40, guardrail()),
),
];
for (label, block) in &cases {
assert!(
is_turn_context_text(block),
"is_turn_context_text rejected real compose_moim output ({label}); the \
producer in goose and the detector in goose-providers have drifted, so the \
volatile block will no longer be kept out of the Anthropic cache prefix:\n\
{block}"
);
}
}
#[test]
fn detector_rejects_text_that_is_not_a_turn_context_block() {
assert!(
!is_turn_context_text("Please refactor the argument parser."),
"ordinary user text must not be treated as a turn-context block"
);
assert!(
!is_turn_context_text(
"<turn-context>\n\
<current-time>2026-06-25 12:00:00</current-time>\n\
<session-id>abc</session-id>\n\
</turn-context>"
),
"a block whose first field after current-time is not working-directory must not \
match: the detector's shape checks are what discriminate the injected block"
);
}
#[test]
fn current_time_is_timezone_aware() {
let block = moim(None, None, 0, 0, vec![]);
let time_line = block
.lines()
.find(|line| line.starts_with(&format!("<{CURRENT_TIME_TAG}>")))
.expect("turn-context block should contain a current-time line");
let value = time_line
.trim_start_matches(&format!("<{CURRENT_TIME_TAG}>"))
.trim_end_matches(&format!("</{CURRENT_TIME_TAG}>"));
chrono::DateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S %:z")
.expect("current-time should include a numeric UTC offset");
}
chrono::DateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S %:z")
.expect("current-time should include a numeric UTC offset");
}
}
crates/goose/src/agents/platform_extensions/chatrecall.rs
+5 -1
View File
@@ -41,7 +41,11 @@ pub struct ChatRecallClient {
}
fn format_agent_visible_excerpt(conversation: &Conversation) -> Option<(usize, String)> {
let messages = conversation.agent_visible_messages();
let messages: Vec<_> = conversation
.agent_visible_messages()
.into_iter()
.filter(|message| !message.is_turn_context())
.collect();
let total = messages.len();
if total == 0 {
return None;
crates/goose/src/agents/platform_extensions/orchestrator.rs
+12 -7
View File
@@ -336,12 +336,13 @@ impl OrchestratorClient {
) -> Result<String, String> {
let provider = self.get_provider().await?;
let conversation_text = Conversation::new_unvalidated(messages.iter().cloned())
.agent_visible_messages()
.iter()
.map(format_message_for_compacting)
.collect::<Vec<_>>()
.join("\n");
let conversation_text = agent_visible_session_messages(&Conversation::new_unvalidated(
messages.iter().cloned(),
))
.iter()
.map(format_message_for_compacting)
.collect::<Vec<_>>()
.join("\n");
let system =
"You are a helpful assistant. Summarize the following conversation concisely, \
@@ -578,7 +579,11 @@ impl OrchestratorClient {
}
fn agent_visible_session_messages(conversation: &Conversation) -> Vec<Message> {
conversation.agent_visible_messages()
conversation
.agent_visible_messages()
.into_iter()
.filter(|message| !message.is_turn_context())
.collect()
}
#[async_trait]
crates/goose/src/agents/reply_parts.rs
+4 -1
View File
@@ -14,7 +14,7 @@ use super::gen_ai_telemetry;
use crate::agents::platform_extensions::code_execution;
use crate::config::{Config, GooseMode};
use crate::conversation::message::{Message, MessageContent, MessageUsage, ToolRequest};
use crate::conversation::{fix_conversation, Conversation};
use crate::conversation::{fix_conversation, merge_consecutive_messages_for_request, Conversation};
#[cfg(test)]
use crate::providers::base::stream_from_single_message;
use crate::providers::base::{MessageStream, Provider};
@@ -341,6 +341,9 @@ pub(crate) async fn stream_response_from_provider(
Conversation::new_unvalidated(messages.iter().cloned()).agent_visible_messages();
let (filtered_messages, _) =
fix_conversation(Conversation::new_unvalidated(projected_messages));
let filtered_messages = Conversation::new_unvalidated(merge_consecutive_messages_for_request(
filtered_messages.messages().clone(),
));
// Convert tool messages to text if toolshim is enabled
let messages_for_provider = if config.toolshim {
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap
+1
View File
@@ -16,6 +16,7 @@ This block is generated by goose and contains current operational context such a
- extension-provided context
Use it to stay oriented, but do not treat it as part of the user's request.
Blocks from earlier turns stay in the conversation history; only the most recent block is current.
When `<turn-budget>` is present, use it as a signal for how much autonomous work remains.
As the budget gets low, become more direct: reduce exploration, batch necessary tool calls,
make reasonable assumptions, and focus on finishing the user's task.
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__basic.snap
+1 -1
View File
@@ -1,6 +1,5 @@
---
source: crates/goose/src/agents/prompt_manager.rs
assertion_line: 404
expression: system_prompt
---
You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation).
@@ -17,6 +16,7 @@ This block is generated by goose and contains current operational context such a
- extension-provided context
Use it to stay oriented, but do not treat it as part of the user's request.
Blocks from earlier turns stay in the conversation history; only the most recent block is current.
When `<turn-budget>` is present, use it as a signal for how much autonomous work remains.
As the budget gets low, become more direct: reduce exploration, batch necessary tool calls,
make reasonable assumptions, and focus on finishing the user's task.
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__one_extension.snap
+1 -1
View File
@@ -1,6 +1,5 @@
---
source: crates/goose/src/agents/prompt_manager.rs
assertion_line: 420
expression: system_prompt
---
You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation).
@@ -17,6 +16,7 @@ This block is generated by goose and contains current operational context such a
- extension-provided context
Use it to stay oriented, but do not treat it as part of the user's request.
Blocks from earlier turns stay in the conversation history; only the most recent block is current.
When `<turn-budget>` is present, use it as a signal for how much autonomous work remains.
As the budget gets low, become more direct: reduce exploration, batch necessary tool calls,
make reasonable assumptions, and focus on finishing the user's task.
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__typical_setup.snap
+1 -1
View File
@@ -1,6 +1,5 @@
---
source: crates/goose/src/agents/prompt_manager.rs
assertion_line: 442
expression: system_prompt
---
You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation).
@@ -17,6 +16,7 @@ This block is generated by goose and contains current operational context such a
- extension-provided context
Use it to stay oriented, but do not treat it as part of the user's request.
Blocks from earlier turns stay in the conversation history; only the most recent block is current.
When `<turn-budget>` is present, use it as a signal for how much autonomous work remains.
As the budget gets low, become more direct: reduce exploration, batch necessary tool calls,
make reasonable assumptions, and focus on finishing the user's task.
crates/goose/src/agents/state_machine/ops_llm.rs
+26 -7
View File
@@ -434,18 +434,32 @@ impl Inference for InferenceRunner<'_> {
.get_context_limit(&self.model_config)
.await
.unwrap_or_else(|_| self.model_config.context_limit());
let turn_start = messages_since_kickoff(conversation)?
let turn = messages_since_kickoff(conversation)?;
let turn_start = turn
.first()
.and_then(|message| chrono::DateTime::from_timestamp(message.created, 0))
.map(|timestamp| timestamp.with_timezone(&chrono::Local))
.unwrap_or_else(chrono::Local::now);
let conversation_for_provider = crate::agents::moim::inject_moim_parts(
Conversation::new_unvalidated(messages_for_provider),
// Persist a fresh turn-context event only when its bytes changed
// (e.g. the turn budget ticked); earlier events stay in place.
let last_turn_context = turn
.iter()
.rev()
.find(|message| message.is_turn_context())
.map(Message::as_concat_text);
let turn_context = crate::agents::moim::turn_context_event(
&session.working_dir,
Some(context_limit),
input.moim_parts,
turn_start,
);
)
.filter(|event| Some(event.as_concat_text()) != last_turn_context);
if let Some(event) = &turn_context {
messages_for_provider.push(event.clone());
}
let conversation_for_provider = Conversation::new_unvalidated(messages_for_provider);
let mut usage_effects: Vec<StateEffect> =
turn_context.into_iter().map(StateEffect::from).collect();
let stream = crate::agents::reply_parts::stream_response_from_provider(
self.provider.clone(),
@@ -460,7 +474,10 @@ impl Inference for InferenceRunner<'_> {
let mut stream = match stream {
Ok(stream) => stream,
Err(err) => return applied(self.error_outcome(&err, emit).await),
Err(err) => {
usage_effects.extend(self.error_outcome(&err, emit).await);
return applied(usage_effects);
}
};
let requested_model = self.model_config.model_name.clone();
@@ -477,7 +494,6 @@ impl Inference for InferenceRunner<'_> {
});
let mut accumulator = Conversation::empty();
let mut usage_effects = Vec::new();
let mut tool_request_ids = std::collections::HashSet::new();
loop {
tokio::select! {
@@ -540,7 +556,10 @@ impl Inference for InferenceRunner<'_> {
return yielded_with(usage_effects);
}
if usage_effects.is_empty() {
let has_recorded_usage = usage_effects
.iter()
.any(|effect| matches!(effect, StateEffect::RecordUsage(_)));
if !has_recorded_usage {
let mut usage = ProviderUsage::new(
self.model_config.model_name.clone(),
goose_providers::conversation::token_usage::Usage::default(),
crates/goose/src/agents/state_machine/ops_toolcalling.rs
+2
View File
@@ -648,6 +648,8 @@ impl Operation for ToolExecutionOperation<'_> {
if extensions.is_empty() {
return Ok(prompt_parts);
}
// HashMap order shuffles across restarts and would bust the prompt cache.
extensions.sort_by(|a, b| a.name.cmp(&b.name));
let mut lines = vec![
"# Extensions".to_string(),
crates/goose/src/agents/state_machine/tests/mod.rs
+29
View File
@@ -49,6 +49,35 @@ async fn max_turns_counts_inference_calls_and_injects_budget() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn turn_context_is_persisted_once_per_turn_and_reused_across_inferences() -> Result<()> {
let (pipeline, api) = test_pipeline().await?;
api.on("add one").call(ADD, value(1));
api.on("result: 1").reply("The total is 1");
api.on("hello").reply("hi there!");
let result = pipeline.run(["add one", "hello"]).await?;
let conversation = result.conversation();
let events: Vec<_> = conversation
.messages()
.iter()
.filter(|message| message.is_turn_context())
.collect();
assert_eq!(
events.len(),
2,
"one turn-context event per turn; the turn's second inference reuses it"
);
assert!(events.iter().all(|event| !event.is_user_visible()));
assert!(api
.calls()
.iter()
.all(|call| call.input_contains("<turn-context>")));
Ok(())
}
#[tokio::test]
async fn goal_starts_nudges_and_clears_when_met() -> Result<()> {
let (pipeline, api) = test_pipeline().await?;
crates/goose/src/agents/state_machine/tests/pipeline.rs
+1
View File
@@ -970,6 +970,7 @@ impl TestRun {
.conversation()
.messages()
.iter()
.filter(|message| !message.is_turn_context())
.flat_map(|message| {
message
.content
crates/goose/src/agents/state_machine/tests/recipe_scheduling_lifecycle.rs
+6 -1
View File
@@ -266,7 +266,12 @@ async fn scheduler_is_advertised_only_when_configured_and_manages_jobs() -> Resu
assert_eq!(scheduled_sessions.len(), 1);
assert_eq!(
scheduled_sessions[0].1.message_count,
listed.conversation().messages().len()
listed
.conversation()
.messages()
.iter()
.filter(|message| message.is_user_visible())
.count()
);
api.on("list scheduled runs").call(
MANAGE_SCHEDULE_TOOL_NAME_COMPLETE,
crates/goose/src/context_mgmt/mod.rs
+233 -9
View File
@@ -100,9 +100,13 @@ pub async fn compact_messages(
has_text && !has_tool_content
};
let (preserved_user_message, is_most_recent) = if !manual_compact {
// Turn-context events are agent-appended, never the message to preserve.
let (preserved_user_message, preserved_idx, is_most_recent) = if !manual_compact {
let found_msg = messages.iter().enumerate().rev().find_map(|(idx, msg)| {
if !msg.is_agent_visible() || !matches!(msg.role, rmcp::model::Role::User) {
if !msg.is_agent_visible()
|| msg.is_turn_context()
|| !matches!(msg.role, rmcp::model::Role::User)
{
return None;
}
@@ -123,13 +127,13 @@ pub async fn compact_messages(
});
if let Some((idx, msg)) = found_msg {
let is_last = idx == messages.len() - 1;
(Some(msg), is_last)
let is_last = messages[idx + 1..].iter().all(Message::is_turn_context);
(Some(msg), Some(idx), is_last)
} else {
(None, false)
(None, None, false)
}
} else {
(None, false)
(None, None, false)
};
let messages_to_compact = messages.as_slice();
@@ -175,6 +179,25 @@ pub async fn compact_messages(
final_messages.push(user_msg);
}
// Carry the turn's own context event (it follows the preserved prompt) so
// a mid-turn retry keeps it; anything earlier belongs to a previous turn.
if let Some(carry_from) = preserved_idx.map(|idx| idx + 1) {
if let Some(turn_context) = messages_to_compact[carry_from..]
.iter()
.rev()
.find(|msg| msg.is_turn_context() && msg.is_agent_visible())
{
let mut carried = turn_context.clone();
carried.id = None;
// Storage reloads order by created_timestamp; the copy must keep
// its appended position, not resurface at the original event's time.
if let Some(latest) = final_messages.iter().map(|msg| msg.created).max() {
carried.created = carried.created.max(latest);
}
final_messages.push(carried);
}
}
let conversation = Conversation::new_unvalidated(final_messages);
let retained_context_tokens = match count_context_tokens(&conversation).await {
Ok(tokens) => tokens,
@@ -317,8 +340,14 @@ async fn do_compact(
session_id: &str,
messages: &[Message],
) -> Result<(Message, ProviderUsage), anyhow::Error> {
let agent_visible_messages =
Conversation::new_unvalidated(messages.iter().cloned()).agent_visible_messages();
// Keep stale per-turn state out of the summary.
let agent_visible_messages = Conversation::new_unvalidated(
messages
.iter()
.filter(|msg| !msg.is_turn_context())
.cloned(),
)
.agent_visible_messages();
// Try progressively removing more tool response messages from the middle to reduce context length
let removal_percentages = [0, 10, 20, 50, 100];
@@ -697,6 +726,7 @@ mod tests {
message: Message,
config: ModelConfig,
max_tool_responses: Option<usize>,
captured_system: std::sync::Mutex<Option<String>>,
}
impl MockProvider {
@@ -715,6 +745,7 @@ mod tests {
request_headers: None,
},
max_tool_responses: None,
captured_system: std::sync::Mutex::new(None),
}
}
@@ -733,10 +764,11 @@ mod tests {
async fn stream(
&self,
_model_config: &ModelConfig,
_system: &str,
system: &str,
messages: &[Message],
_tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
*self.captured_system.lock().unwrap() = Some(system.to_string());
// If max_tool_responses is set, fail if we have too many
if let Some(max) = self.max_tool_responses {
let tool_response_count = messages
@@ -956,6 +988,198 @@ mod tests {
assert!(!user_text.contains("assistant-only preprompt"));
}
#[tokio::test]
async fn preserved_user_message_skips_turn_context_events() {
let conversation = Conversation::new_unvalidated([
Message::user().with_text("earlier request"),
Message::assistant().with_text("earlier response"),
Message::user().with_text("the real current request"),
Message::user()
.with_text("<turn-context>frozen block</turn-context>")
.with_metadata(MessageMetadata::agent_only().with_turn_context()),
]);
let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000);
let compacted = compact_messages(
&provider,
&provider.config,
"test-session-id",
&conversation,
false,
)
.await
.unwrap()
.conversation;
let preserved: Vec<_> = compacted
.messages()
.iter()
.filter(|message| message.is_agent_visible() && message.role == Role::User)
.collect();
let turn_context_events: Vec<_> = preserved
.iter()
.filter(|message| message.is_turn_context())
.collect();
let [carried_block] = turn_context_events.as_slice() else {
panic!("expected exactly one carried turn-context event");
};
assert!(
carried_block.as_concat_text().contains("frozen block"),
"the turn's context event must be carried forward for the mid-turn retry"
);
let last = compacted.messages().last().unwrap();
assert!(
last.is_turn_context() && last.is_agent_visible(),
"the carried event must trail the preserved user message"
);
let user_prompt = preserved[preserved.len() - 2];
assert!(
user_prompt
.as_concat_text()
.contains("the real current request"),
"the user's prompt must survive compaction verbatim"
);
assert!(!user_prompt.is_turn_context());
let continuation = compacted
.messages()
.iter()
.find(|message| message.role == Role::Assistant && message.is_agent_visible())
.unwrap()
.as_concat_text();
assert!(
continuation.contains(CONVERSATION_CONTINUATION_TEXT),
"a trailing turn-context event must not demote the compaction to a tool-loop continuation"
);
}
#[tokio::test]
async fn stale_turn_context_from_an_earlier_turn_is_not_carried() {
let conversation = Conversation::new_unvalidated([
Message::user().with_text("earlier request"),
Message::user()
.with_text("<turn-context>stale block</turn-context>")
.with_metadata(MessageMetadata::agent_only().with_turn_context()),
Message::assistant().with_text("earlier response"),
Message::user().with_text("the new prompt"),
]);
let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000);
let compacted = compact_messages(
&provider,
&provider.config,
"test-session-id",
&conversation,
false,
)
.await
.unwrap()
.conversation;
assert!(
!compacted
.messages()
.iter()
.any(|message| message.is_agent_visible() && message.is_turn_context()),
"pre-turn compaction must not resurrect a previous turn's context event"
);
let last = compacted.messages().last().unwrap();
assert_eq!(last.as_concat_text(), "the new prompt");
}
#[tokio::test]
async fn carried_turn_context_stays_last_after_persist_and_reload() {
let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000);
let mut prompt = Message::user().with_text("the real current request");
prompt.created -= 3600;
let mut block = Message::user()
.with_text("<turn-context>frozen block</turn-context>")
.with_metadata(MessageMetadata::agent_only().with_turn_context());
block.created -= 3600;
let conversation = Conversation::new_unvalidated([prompt, block]);
let compacted = compact_messages(
&provider,
&provider.config,
"test-session-id",
&conversation,
false,
)
.await
.unwrap()
.conversation;
assert!(compacted.messages().last().unwrap().is_turn_context());
let temp_dir = tempfile::tempdir().unwrap();
let manager = crate::session::SessionManager::new(temp_dir.path().to_path_buf());
let session = manager
.create_session(
std::path::PathBuf::from("/tmp/test"),
"carry order".to_string(),
crate::session::session_manager::SessionType::User,
crate::config::GooseMode::default(),
)
.await
.unwrap();
manager
.replace_conversation(&session.id, &compacted)
.await
.unwrap();
let reloaded = manager
.get_session(&session.id, true)
.await
.unwrap()
.conversation
.unwrap();
assert_eq!(reloaded.messages().len(), compacted.messages().len());
let last = reloaded.messages().last().unwrap();
assert!(
last.is_turn_context(),
"the carried event must not resurface at its original timestamp on reload"
);
}
#[tokio::test]
async fn summarizer_input_excludes_turn_context_events() {
let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000);
let turn_context = |text: &str| {
Message::user()
.with_text(text)
.with_metadata(MessageMetadata::agent_only().with_turn_context())
};
let conversation = Conversation::new_unvalidated([
Message::user().with_text("please refactor the parser"),
turn_context("<turn-context>cwd /old/dir</turn-context>"),
Message::assistant().with_text("working on it"),
turn_context("<turn-context>cwd /new/dir</turn-context>"),
]);
let compacted = compact_messages(
&provider,
&provider.config,
"test-session-id",
&conversation,
false,
)
.await
.unwrap()
.conversation;
let system = provider.captured_system.lock().unwrap().clone().unwrap();
assert!(system.contains("please refactor the parser"));
assert!(
!system.contains("/old/dir") && !system.contains("/new/dir"),
"turn-context events must not reach the summarizer as dialogue"
);
let carried = compacted.messages().last().unwrap();
assert!(
carried.is_turn_context() && carried.as_concat_text().contains("/new/dir"),
"the newest turn-context event must still be carried forward"
);
}
#[tokio::test]
async fn tool_pair_summary_projects_nested_audiences_before_provider_input() {
let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000);
crates/goose/src/providers/litellm.rs
+3 -67
View File
@@ -1,10 +1,11 @@
use anyhow::Result;
use async_trait::async_trait;
use futures::future::BoxFuture;
use goose_providers::cache_semantics::apply_chat_payload_breakpoints;
use goose_providers::conversation::token_usage::ProviderUsage;
use goose_providers::errors::ProviderError;
use goose_providers::images::ImageFormat;
use serde_json::{json, Value};
use serde_json::Value;
use std::collections::HashMap;
use super::api_client::{ApiClient, AuthMethod};
@@ -248,7 +249,7 @@ impl Provider for LiteLLMProvider {
)?;
if self.supports_cache_control(model_config).await {
payload = update_request_for_cache_control(&payload);
apply_chat_payload_breakpoints(&mut payload);
}
let response = self
@@ -280,71 +281,6 @@ impl Provider for LiteLLMProvider {
}
}
/// Updates the request payload to include cache control headers for automatic prompt caching
/// Adds ephemeral cache control to the last 2 user messages, system message, and last tool
pub fn update_request_for_cache_control(original_payload: &Value) -> Value {
let mut payload = original_payload.clone();
if let Some(messages_spec) = payload
.as_object_mut()
.and_then(|obj| obj.get_mut("messages"))
.and_then(|messages| messages.as_array_mut())
{
let mut user_count = 0;
for message in messages_spec.iter_mut().rev() {
if message.get("role") == Some(&json!("user")) {
if let Some(content) = message.get_mut("content") {
if let Some(content_str) = content.as_str() {
*content = json!([{
"type": "text",
"text": content_str,
"cache_control": { "type": "ephemeral" }
}]);
}
}
user_count += 1;
if user_count >= 2 {
break;
}
}
}
if let Some(system_message) = messages_spec
.iter_mut()
.find(|msg| msg.get("role") == Some(&json!("system")))
{
if let Some(content) = system_message.get_mut("content") {
if let Some(content_str) = content.as_str() {
*system_message = json!({
"role": "system",
"content": [{
"type": "text",
"text": content_str,
"cache_control": { "type": "ephemeral" }
}]
});
}
}
}
}
if let Some(tools_spec) = payload
.as_object_mut()
.and_then(|obj| obj.get_mut("tools"))
.and_then(|tools| tools.as_array_mut())
{
if let Some(last_tool) = tools_spec.last_mut() {
if let Some(function) = last_tool.get_mut("function") {
function
.as_object_mut()
.unwrap()
.insert("cache_control".to_string(), json!({ "type": "ephemeral" }));
}
}
}
payload
}
fn parse_custom_headers(headers_str: String) -> HashMap<String, String> {
let mut headers = HashMap::new();
for line in headers_str.lines() {
crates/goose/src/providers/openrouter.rs
+5 -82
View File
@@ -11,6 +11,7 @@ use super::openai_compatible::{handle_status, stream_openai_compat};
use super::retry::ProviderRetry;
use crate::conversation::message::Message;
use crate::providers::formats::openrouter as openrouter_format;
use goose_providers::cache_semantics::{apply_chat_payload_breakpoints, CacheSemantics};
use goose_providers::errors::ProviderError;
use goose_providers::formats::openai::create_request;
use goose_providers::model::ModelConfig;
@@ -21,7 +22,6 @@ pub const OPENROUTER_PROVIDER_NAME: &str = "openrouter";
const OPENROUTER_PARAMETERS_CONFIG_KEY: &str = "OPENROUTER_PARAMETERS";
pub const OPENROUTER_DEFAULT_MODEL: &str = "anthropic/claude-sonnet-4";
pub const OPENROUTER_DEFAULT_FAST_MODEL: &str = "google/gemini-2.5-flash";
pub const OPENROUTER_MODEL_PREFIX_ANTHROPIC: &str = "anthropic";
// OpenRouter can run many models, we suggest the default
pub const OPENROUTER_KNOWN_MODELS: &[&str] = &[
@@ -98,79 +98,6 @@ fn is_mandatory_reasoning_error(error: &ProviderError) -> bool {
matches!(error, ProviderError::RequestFailed(message) if message.contains("Reasoning is mandatory"))
}
/// Update the request when using anthropic model.
/// For anthropic model, we can enable prompt caching to save cost. Since openrouter is the OpenAI compatible
/// endpoint, we need to modify the open ai request to have anthropic cache control field.
fn update_request_for_anthropic(original_payload: &Value) -> Value {
let mut payload = original_payload.clone();
if let Some(messages_spec) = payload
.as_object_mut()
.and_then(|obj| obj.get_mut("messages"))
.and_then(|messages| messages.as_array_mut())
{
// Add "cache_control" to the last and second-to-last "user" messages.
// During each turn, we mark the final message with cache_control so the conversation can be
// incrementally cached. The second-to-last user message is also marked for caching with the
// cache_control parameter, so that this checkpoint can read from the previous cache.
let mut user_count = 0;
for message in messages_spec.iter_mut().rev() {
if message.get("role") == Some(&json!("user")) {
if let Some(content) = message.get_mut("content") {
if let Some(content_str) = content.as_str() {
*content = json!([{
"type": "text",
"text": content_str,
"cache_control": { "type": "ephemeral" }
}]);
}
}
user_count += 1;
if user_count >= 2 {
break;
}
}
}
// Update the system message to have cache_control field.
if let Some(system_message) = messages_spec
.iter_mut()
.find(|msg| msg.get("role") == Some(&json!("system")))
{
if let Some(content) = system_message.get_mut("content") {
if let Some(content_str) = content.as_str() {
*system_message = json!({
"role": "system",
"content": [{
"type": "text",
"text": content_str,
"cache_control": { "type": "ephemeral" }
}]
});
}
}
}
}
if let Some(tools_spec) = payload
.as_object_mut()
.and_then(|obj| obj.get_mut("tools"))
.and_then(|tools| tools.as_array_mut())
{
// Add "cache_control" to the last tool spec, if any. This means that all tool definitions,
// will be cached as a single prefix.
if let Some(last_tool) = tools_spec.last_mut() {
if let Some(function) = last_tool.get_mut("function") {
function
.as_object_mut()
.unwrap()
.insert("cache_control".to_string(), json!({ "type": "ephemeral" }));
}
}
}
payload
}
fn is_gemini_model(model_name: &str) -> bool {
model_name.starts_with("google/")
}
@@ -350,8 +277,10 @@ impl Provider for OpenRouterProvider {
}
}
if supports_cache_control(model_config) {
payload = update_request_for_anthropic(&payload);
if CacheSemantics::for_model(OPENROUTER_PROVIDER_NAME, &model_config.model_name)
.uses_explicit_breakpoints()
{
apply_chat_payload_breakpoints(&mut payload);
}
if is_gemini_model(&model_config.model_name) {
@@ -386,12 +315,6 @@ impl Provider for OpenRouterProvider {
}
}
fn supports_cache_control(model: &ModelConfig) -> bool {
model
.model_name
.starts_with(OPENROUTER_MODEL_PREFIX_ANTHROPIC)
}
#[cfg(test)]
mod tests {
use super::*;
crates/goose/src/scheduler.rs
+7 -1
View File
@@ -1493,7 +1493,13 @@ mod tests {
.unwrap();
assert_eq!(
sessions[0].message_count,
session.conversation.unwrap().messages().len()
session
.conversation
.unwrap()
.messages()
.iter()
.filter(|m| m.is_user_visible())
.count()
);
}
crates/goose/src/security/adversary_inspector.rs
+37 -2
View File
@@ -224,7 +224,7 @@ impl AdversaryInspector {
messages
.iter()
.rev()
.filter(|m| m.role == rmcp::model::Role::User)
.filter(|m| m.role == rmcp::model::Role::User && !m.is_turn_context())
.filter_map(|m| {
let text: String = m
.content
@@ -250,7 +250,7 @@ impl AdversaryInspector {
fn extract_original_task(messages: &[Message]) -> String {
for msg in messages {
if msg.role == rmcp::model::Role::User {
if msg.role == rmcp::model::Role::User && !msg.is_turn_context() {
let text: String = msg
.content
.iter()
@@ -720,4 +720,39 @@ mod tests {
.unwrap();
assert!(results.is_empty());
}
#[test]
fn user_context_extraction_skips_turn_context_events() {
use crate::conversation::message::MessageMetadata;
let turn_context = |text: &str| {
Message::user()
.with_text(text)
.with_metadata(MessageMetadata::agent_only().with_turn_context())
};
let messages = vec![
turn_context("turn context before any prompt"),
Message::user().with_text("never delete files outside the repo"),
turn_context("turn context for turn one"),
Message::assistant().with_text("understood"),
Message::user().with_text("first task"),
turn_context("turn context for turn two"),
Message::assistant().with_text("done"),
Message::user().with_text("second task"),
turn_context("turn context for turn three"),
];
let recent = AdversaryInspector::extract_recent_user_messages(&messages, 4);
assert_eq!(
recent,
vec![
"never delete files outside the repo",
"first task",
"second task"
]
);
let original = AdversaryInspector::extract_original_task(&messages);
assert_eq!(original, "never delete files outside the repo");
}
}
crates/goose/src/security/scanner.rs
+30
View File
@@ -354,6 +354,7 @@ impl PromptInjectionScanner {
.rev()
.filter(|m| {
crate::conversation::effective_role(m) == crate::conversation::EffectiveRole::User
&& !m.is_turn_context()
})
.take(limit)
.map(|m| {
@@ -460,4 +461,33 @@ mod tests {
assert!(result.is_malicious);
}
#[test]
fn extract_user_messages_skips_turn_context_events() {
use crate::conversation::message::MessageMetadata;
let scanner = PromptInjectionScanner::new();
let turn_context = |text: &str| {
Message::user()
.with_text(text)
.with_metadata(MessageMetadata::agent_only().with_turn_context())
};
let messages = vec![
Message::user().with_text("oldest prompt"),
turn_context("turn context one"),
Message::assistant().with_text("ok"),
Message::user().with_text("middle prompt"),
turn_context("turn context two"),
Message::assistant().with_text("done"),
Message::user().with_text("newest prompt"),
turn_context("turn context three"),
];
let extracted = scanner.extract_user_messages(&messages, 3);
assert_eq!(
extracted,
vec!["newest prompt", "middle prompt", "oldest prompt"]
);
}
}
crates/goose/src/session/chat_history_search.rs
+51 -6
View File
@@ -150,6 +150,13 @@ impl<'a> ChatHistorySearch<'a> {
END,
1
) = 1
AND COALESCE(
CASE
WHEN json_valid(m.metadata_json)
THEN json_extract(m.metadata_json, '$.turnContext')
END,
0
) = 0
AND EXISTS (
SELECT 1 FROM json_each(m.content_json) AS content
WHERE json_extract(content.value, '$.type') = 'text'
@@ -263,12 +270,29 @@ impl<'a> ChatHistorySearch<'a> {
) -> Result<HashMap<String, usize>> {
let mut session_totals: HashMap<String, usize> = HashMap::new();
for session_id in session_messages.keys() {
let count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM messages WHERE session_id = ?")
.bind(session_id)
.fetch_one(self.pool)
.await
.unwrap_or(0);
let count: i64 = sqlx::query_scalar(
r#"
SELECT COUNT(*) FROM messages m WHERE m.session_id = ?
AND COALESCE(
CASE
WHEN json_valid(m.metadata_json)
THEN json_extract(m.metadata_json, '$.agentVisible')
END,
1
) = 1
AND COALESCE(
CASE
WHEN json_valid(m.metadata_json)
THEN json_extract(m.metadata_json, '$.turnContext')
END,
0
) = 0
"#,
)
.bind(session_id)
.fetch_one(self.pool)
.await
.unwrap_or(0);
session_totals.insert(session_id.clone(), count as usize);
}
Ok(session_totals)
@@ -443,5 +467,26 @@ mod tests {
.await
.unwrap();
assert_eq!(hidden_only.total_matches, 0);
insert_message(
&pool,
&Message::user()
.with_text("<turn-context>needle in operational context</turn-context>")
.with_metadata(MessageMetadata::agent_only().with_turn_context()),
now,
)
.await;
let needle = ChatHistorySearch::new(&pool, "needle", Some(10), None, None, None, vec![])
.execute()
.await
.unwrap();
assert_eq!(
needle.total_matches, 1,
"turn-context events are operational context, not searchable conversation"
);
assert_eq!(
needle.results[0].total_messages_in_session, 2,
"session totals must not count turn-context or user-only rows"
);
}
}
crates/goose/src/session/session_manager.rs
+54 -4
View File
@@ -365,12 +365,14 @@ fn message_keyword_clause(keyword_count: usize) -> String {
.collect::<Vec<_>>()
.join(" OR ");
let visible = user_visible_message_sql("mq.metadata_json");
format!(
r#"
EXISTS (
SELECT 1
FROM messages mq
WHERE mq.session_id = s.id
AND {visible}
AND EXISTS (
SELECT 1
FROM json_each(mq.content_json)
@@ -619,7 +621,7 @@ impl SessionManager {
let user_message_count = conversation
.messages()
.iter()
.filter(|m| matches!(m.role, Role::User))
.filter(|m| matches!(m.role, Role::User) && m.is_user_visible())
.count();
let should_generate_name = if provider.manages_own_context() {
@@ -714,6 +716,10 @@ fn normalized_message_timestamp_sql(column: &str) -> String {
)
}
fn user_visible_message_sql(column: &str) -> String {
format!("COALESCE(json_extract({column}, '$.userVisible'), 1) != 0")
}
fn session_sort_at(session: &Session) -> DateTime<Utc> {
session.last_message_at.unwrap_or(session.updated_at)
}
@@ -1629,7 +1635,11 @@ impl SessionStorage {
if include_messages {
let conv = self.get_conversation(&session.id).await?;
session.message_count = conv.messages().len();
session.message_count = conv
.messages()
.iter()
.filter(|m| m.is_user_visible())
.count();
session.last_message_at = conv
.messages()
.iter()
@@ -1638,7 +1648,8 @@ impl SessionStorage {
session.conversation = Some(conv);
} else {
let sql = format!(
"SELECT COUNT(*), MAX({}) FROM messages WHERE session_id = ?",
"SELECT COUNT(*) FILTER (WHERE {}), MAX({}) FROM messages WHERE session_id = ?",
user_visible_message_sql("metadata_json"),
normalized_message_timestamp_sql("created_timestamp")
);
let (count, last_message_timestamp): (i64, Option<i64>) = sqlx::query_as(&sql)
@@ -1979,7 +1990,7 @@ impl SessionStorage {
s.schedule_id, s.recipe_json, s.user_recipe_values_json,
s.provider_name, s.model_config_json, s.goose_mode,
s.archived_at, s.project_id, s.parent_session_id,
COUNT(m.id) as message_count,
COUNT(m.id) FILTER (WHERE {}) as message_count,
MAX({}) as last_message_timestamp,
{} as sort_timestamp
FROM sessions s
@@ -1990,6 +2001,7 @@ impl SessionStorage {
{}
{}
"#,
user_visible_message_sql("m.metadata_json"),
normalized_message_timestamp,
sort_timestamp_sql,
message_join,
@@ -2985,6 +2997,44 @@ mod tests {
assert_eq!(with_messages.last_message_at, Some(expected));
}
#[tokio::test]
async fn test_message_count_excludes_agent_only_messages() {
let temp_dir = TempDir::new().unwrap();
let sm = SessionManager::new(temp_dir.path().to_path_buf());
let session = sm
.create_session(
PathBuf::from("/tmp/test"),
"Hidden events".to_string(),
SessionType::User,
GooseMode::default(),
)
.await
.unwrap();
add_user_message(&sm, &session.id).await;
sm.add_message(
&session.id,
&Message::user()
.with_text("<turn-context>frozen</turn-context>")
.with_visibility(false, true),
)
.await
.unwrap();
sm.add_message(&session.id, &Message::assistant().with_text("hi"))
.await
.unwrap();
let counted = sm.get_session(&session.id, false).await.unwrap();
assert_eq!(counted.message_count, 2);
let with_messages = sm.get_session(&session.id, true).await.unwrap();
assert_eq!(with_messages.message_count, 2);
let listed = sm.list_sessions().await.unwrap();
let listed_session = listed.iter().find(|s| s.id == session.id).unwrap();
assert_eq!(listed_session.message_count, 2);
}
#[tokio::test]
async fn test_truncate_conversation_from_message_keeps_same_second_previous_rows() {
let temp_dir = TempDir::new().unwrap();
crates/goose/src/session/session_naming.rs
+1 -1
View File
@@ -75,7 +75,7 @@ fn extract_short_title(text: &str) -> String {
fn get_initial_user_messages(messages: &Conversation) -> Vec<String> {
messages
.iter()
.filter(|m| m.role == rmcp::model::Role::User)
.filter(|m| m.role == rmcp::model::Role::User && m.is_user_visible())
.take(MSG_COUNT_FOR_SESSION_NAME_GENERATION)
.map(|m| {
m.content
crates/goose/tests/agent.rs
+9 -2
View File
@@ -1404,7 +1404,11 @@ mod tests {
.messages()
.to_vec();
let user_count = messages.iter().filter(|m| m.role == Role::User).count();
// Turn-context events are excluded; this test is about streaming deltas.
let user_count = messages
.iter()
.filter(|m| m.role == Role::User && m.is_user_visible())
.count();
let asst_count = messages
.iter()
.filter(|m| m.role == Role::Assistant)
@@ -1453,7 +1457,10 @@ mod tests {
.messages()
.to_vec();
let user_count2 = messages2.iter().filter(|m| m.role == Role::User).count();
let user_count2 = messages2
.iter()
.filter(|m| m.role == Role::User && m.is_user_visible())
.count();
let asst_count2 = messages2
.iter()
.filter(|m| m.role == Role::Assistant)
crates/goose/tests/compaction.rs
+6 -1
View File
@@ -320,7 +320,12 @@ fn assert_conversation_compacted(conversation: &Conversation) {
"Message after compaction at index {} should be agent visible",
idx
);
if idx == continuation_end && matches!(msg.role, rmcp::model::Role::User) {
if msg.is_turn_context() {
assert!(
!msg.is_user_visible(),
"Carried turn-context event should be user-invisible"
);
} else if idx == continuation_end && matches!(msg.role, rmcp::model::Role::User) {
assert!(
!msg.is_user_visible(),
"Projected preserved user message should be user-invisible"