Reduce repeated request context and diagnostic formatting work
Deduplicate exact shared context without dropping distinct instructions or tool results. Stop raw request diagnostic serialization at the configured prefix limit and mark truncated sizes as unknown.
This commit is contained in:
@@ -1059,7 +1059,7 @@ impl ResponseStream {
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn raw_model_request_payload(params: &api::RequestParams, max_chars: usize) -> serde_json::Value {
|
||||
use serde::ser::{SerializeMap, Serializer};
|
||||
use serde::ser::{SerializeMap, SerializeSeq, Serializer};
|
||||
|
||||
// collect_str writes debug output directly through serde_json's escaping
|
||||
// writer, without format! allocating another copy of potentially large input.
|
||||
@@ -1070,7 +1070,18 @@ fn raw_model_request_payload(params: &api::RequestParams, max_chars: usize) -> s
|
||||
}
|
||||
}
|
||||
|
||||
remote_logging::bounded_payload_context(max_chars, |writer| {
|
||||
struct NewestFirst<'a, T>(&'a [T]);
|
||||
impl<T: serde::Serialize> serde::Serialize for NewestFirst<'_, T> {
|
||||
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
let mut sequence = serializer.serialize_seq(Some(self.0.len()))?;
|
||||
for message in self.0.iter().rev() {
|
||||
sequence.serialize_element(message)?;
|
||||
}
|
||||
sequence.end()
|
||||
}
|
||||
}
|
||||
|
||||
remote_logging::prefix_bounded_payload_context(max_chars, |writer| {
|
||||
let mut serializer = serde_json::Serializer::pretty(writer);
|
||||
let mut map = serializer.serialize_map(None)?;
|
||||
map.serialize_entry("model_id", ¶ms.model.as_str())?;
|
||||
@@ -1079,10 +1090,14 @@ fn raw_model_request_payload(params: &api::RequestParams, max_chars: usize) -> s
|
||||
map.serialize_entry("computer_use_model_id", ¶ms.computer_use_model.as_str())?;
|
||||
map.serialize_entry("input", &DebugPayload(¶ms.input))?;
|
||||
map.serialize_entry("tool_results", &DebugPayload(¶ms.tool_results))?;
|
||||
map.serialize_entry("tasks", &DebugPayload(¶ms.tasks))?;
|
||||
map.serialize_entry("message_history", ¶ms.message_history)?;
|
||||
map.serialize_entry("progressive_summary", ¶ms.progressive_summary)?;
|
||||
map.serialize_entry("tool_result_archive", ¶ms.tool_result_archive)?;
|
||||
// Preserve the newest request context before old task/debug archives can
|
||||
// consume the diagnostic budget. This ordering affects diagnostics only.
|
||||
map.serialize_entry("message_history_order", "newest_first")?;
|
||||
map.serialize_entry(
|
||||
"message_history",
|
||||
&NewestFirst(params.message_history.as_slice()),
|
||||
)?;
|
||||
map.serialize_entry("global_rules", ¶ms.global_rules)?;
|
||||
map.serialize_entry("mcp_context", &DebugPayload(¶ms.mcp_context))?;
|
||||
map.serialize_entry("autonomy_level", &DebugPayload(¶ms.autonomy_level))?;
|
||||
@@ -1115,7 +1130,9 @@ fn raw_model_request_payload(params: &api::RequestParams, max_chars: usize) -> s
|
||||
"omitted_sensitive_fields",
|
||||
&["api_keys", "custom_model_providers", "custom_model_routers"],
|
||||
)?;
|
||||
map.end()
|
||||
map.serialize_entry("tasks", &DebugPayload(¶ms.tasks))?;
|
||||
map.serialize_entry("tool_result_archive", ¶ms.tool_result_archive)?;
|
||||
SerializeMap::end(map)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -311,6 +311,89 @@ fn truncate_chars(value: &str, max_chars: usize) -> String {
|
||||
truncated
|
||||
}
|
||||
|
||||
/// Stop serialization once the diagnostic prefix is full. Unlike the tail
|
||||
/// collector below, this does not walk the remaining history or tool archive.
|
||||
/// A truncated prefix is diagnostic text, not a complete serialized request.
|
||||
pub(crate) fn prefix_bounded_payload_context(
|
||||
max_chars: usize,
|
||||
serialize: impl FnOnce(&mut dyn std::io::Write) -> serde_json::Result<()>,
|
||||
) -> Value {
|
||||
struct PrefixWriter {
|
||||
bytes: Vec<u8>,
|
||||
max_chars: usize,
|
||||
max_bytes: usize,
|
||||
chars: usize,
|
||||
truncated: bool,
|
||||
}
|
||||
|
||||
impl std::io::Write for PrefixWriter {
|
||||
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
|
||||
if self.truncated {
|
||||
return Err(std::io::Error::other("diagnostic prefix limit reached"));
|
||||
}
|
||||
let mut included = 0;
|
||||
for byte in bytes {
|
||||
let starts_character = byte & 0xc0 != 0x80;
|
||||
if (starts_character && self.chars == self.max_chars)
|
||||
|| self.bytes.len().saturating_add(included) == self.max_bytes
|
||||
{
|
||||
self.bytes.extend_from_slice(&bytes[..included]);
|
||||
self.truncated = true;
|
||||
return if included == 0 {
|
||||
Err(std::io::Error::other("diagnostic prefix limit reached"))
|
||||
} else {
|
||||
// Respect Write's short-write contract. serde's write_all
|
||||
// immediately retries the remainder and receives the error.
|
||||
Ok(included)
|
||||
};
|
||||
}
|
||||
if starts_character {
|
||||
self.chars += 1;
|
||||
}
|
||||
included += 1;
|
||||
}
|
||||
self.bytes.extend_from_slice(bytes);
|
||||
Ok(bytes.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
let mut writer = PrefixWriter {
|
||||
bytes: Vec::new(),
|
||||
max_chars,
|
||||
max_bytes: max_chars.saturating_mul(4),
|
||||
chars: 0,
|
||||
truncated: false,
|
||||
};
|
||||
let result = serialize(&mut writer);
|
||||
if let Err(error) = result {
|
||||
if !writer.truncated {
|
||||
return json!({ "serialization_error": sanitize_error(error) });
|
||||
}
|
||||
}
|
||||
// Writers may split a UTF-8 character across calls. Keep only a complete
|
||||
// prefix if serialization stopped before the last character was finished.
|
||||
let valid_length = match std::str::from_utf8(&writer.bytes) {
|
||||
Ok(_) => writer.bytes.len(),
|
||||
Err(error) => error.valid_up_to(),
|
||||
};
|
||||
writer.bytes.truncate(valid_length);
|
||||
let payload = String::from_utf8(writer.bytes).expect("validated diagnostic UTF-8 prefix");
|
||||
let included_chars = payload.chars().count();
|
||||
json!({
|
||||
"payload": payload,
|
||||
"payload_total_chars": if writer.truncated { None } else { Some(included_chars) },
|
||||
"payload_included_chars": included_chars,
|
||||
"payload_max_chars": max_chars,
|
||||
"truncated": writer.truncated,
|
||||
"truncation_strategy": if writer.truncated { Some("prefix") } else { None },
|
||||
"serialization_stopped_early": writer.truncated,
|
||||
})
|
||||
}
|
||||
|
||||
/// Serialize directly into a bounded UTF-8 tail instead of materializing the full
|
||||
/// request and its archive as a JSON Value and then another large String.
|
||||
pub(crate) fn bounded_payload_context(
|
||||
|
||||
@@ -64,3 +64,79 @@ fn streamed_json_matches_serialized_payload_without_materializing_the_archive()
|
||||
});
|
||||
assert_eq!(actual, tail_limited_payload_context(&expected, 100));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefix_payload_preserves_utf8_and_reports_unknown_truncated_size() {
|
||||
let payload = "é日本🙂 ending";
|
||||
for limit in [0, 1, 3, 4, payload.chars().count(), payload.len()] {
|
||||
for chunk_size in [1, 2, 7, payload.len()] {
|
||||
let actual = super::prefix_bounded_payload_context(limit, |writer| {
|
||||
for chunk in payload.as_bytes().chunks(chunk_size) {
|
||||
writer.write_all(chunk).map_err(serde_json::Error::io)?;
|
||||
}
|
||||
Ok(())
|
||||
});
|
||||
let expected = payload.chars().take(limit).collect::<String>();
|
||||
let truncated = limit < payload.chars().count();
|
||||
assert_eq!(actual["payload"], expected);
|
||||
assert_eq!(actual["payload_included_chars"], expected.chars().count());
|
||||
assert_eq!(actual["truncated"], truncated);
|
||||
assert_eq!(actual["serialization_stopped_early"], truncated);
|
||||
assert_eq!(
|
||||
actual["payload_total_chars"],
|
||||
if truncated {
|
||||
serde_json::Value::Null
|
||||
} else {
|
||||
json!(payload.chars().count())
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefix_payload_stops_pretty_debug_serialization_at_the_budget() {
|
||||
use std::cell::Cell;
|
||||
|
||||
struct LargeDebug<'a>(&'a Cell<usize>);
|
||||
impl std::fmt::Debug for LargeDebug<'_> {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
for _ in 0..1_000_000 {
|
||||
self.0.set(self.0.get() + 1);
|
||||
formatter.write_str("large archive entry ")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
struct DebugPayload<'a>(LargeDebug<'a>);
|
||||
impl serde::Serialize for DebugPayload<'_> {
|
||||
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
serializer.collect_str(&format_args!("{:#?}", self.0))
|
||||
}
|
||||
}
|
||||
|
||||
let visited = Cell::new(0);
|
||||
let actual = super::prefix_bounded_payload_context(64, |writer| {
|
||||
serde_json::to_writer(writer, &DebugPayload(LargeDebug(&visited)))
|
||||
});
|
||||
|
||||
assert!(visited.get() < 10);
|
||||
assert_eq!(actual["payload_included_chars"], 64);
|
||||
assert_eq!(actual["truncated"], true);
|
||||
assert_eq!(actual["payload_total_chars"], serde_json::Value::Null);
|
||||
assert_eq!(actual["truncation_strategy"], "prefix");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefix_payload_keeps_real_serialization_errors_visible() {
|
||||
let actual = super::prefix_bounded_payload_context(100, |_writer| {
|
||||
Err(serde_json::Error::io(std::io::Error::other(
|
||||
"test serializer failed",
|
||||
)))
|
||||
});
|
||||
assert!(actual["serialization_error"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("test serializer failed"));
|
||||
assert!(actual.get("payload").is_none());
|
||||
}
|
||||
|
||||
@@ -844,6 +844,11 @@ fn provider_safe_mcp_tool_name(server_id: Option<&str>, tool_name: &str) -> Stri
|
||||
format!("{prefix}{component}{separator}{hash}")
|
||||
}
|
||||
|
||||
fn deduplicate_context_entries<T: Clone + Eq + std::hash::Hash>(entries: &mut Vec<T>) {
|
||||
let mut seen = HashSet::with_capacity(entries.len());
|
||||
entries.retain(|entry| seen.insert(entry.clone()));
|
||||
}
|
||||
|
||||
fn build_system_prompt(
|
||||
inputs: &[AIAgentInput],
|
||||
tools: &[ToolDefinition],
|
||||
@@ -988,6 +993,13 @@ fn build_system_prompt(
|
||||
AIAgentContext::Image(_) => {}
|
||||
}
|
||||
}
|
||||
// Grouped inputs, such as background-job completions, can share the same context.
|
||||
// Emit each exact context entry once, retaining its first position and all distinct
|
||||
// content. Conversational inputs and tool results remain separate and unchanged.
|
||||
deduplicate_context_entries(&mut environment);
|
||||
deduplicate_context_entries(&mut project_rules);
|
||||
deduplicate_context_entries(&mut attached_context);
|
||||
deduplicate_context_entries(&mut available_skills);
|
||||
if !environment.is_empty() {
|
||||
prompt.push_str("## Environment\n");
|
||||
for item in environment {
|
||||
|
||||
@@ -81,6 +81,177 @@ fn native_context_reaches_rig_without_a_proto_context_conversion() {
|
||||
assert!(prompt.contains("Indexed codebase: galaxy (/repo)"));
|
||||
}
|
||||
|
||||
fn shared_prompt_context(
|
||||
directory: &str,
|
||||
rule: &str,
|
||||
attachment: &str,
|
||||
skill_description: &str,
|
||||
) -> Vec<AIAgentContext> {
|
||||
vec![
|
||||
AIAgentContext::Directory {
|
||||
pwd: Some(directory.to_string()),
|
||||
home_dir: Some("/home/user".to_string()),
|
||||
are_file_symbols_indexed: false,
|
||||
},
|
||||
AIAgentContext::ProjectRules {
|
||||
root_path: "/repo".to_string(),
|
||||
active_rules: vec![FileContext::new(
|
||||
"/repo/AGENTS.md".to_string(),
|
||||
AnyFileContent::StringContent(rule.to_string()),
|
||||
None,
|
||||
None,
|
||||
)],
|
||||
additional_rule_paths: Vec::new(),
|
||||
},
|
||||
AIAgentContext::File(FileContext::new(
|
||||
"/repo/context.txt".to_string(),
|
||||
AnyFileContent::StringContent(attachment.to_string()),
|
||||
None,
|
||||
None,
|
||||
)),
|
||||
AIAgentContext::Skills {
|
||||
skills: vec![SkillDescriptor {
|
||||
reference: SkillReference::Path(LocalOrRemotePath::Local(PathBuf::from(
|
||||
"/repo/.agents/skills/context-probe/SKILL.md",
|
||||
))),
|
||||
name: "context-probe".to_string(),
|
||||
description: skill_description.to_string(),
|
||||
scope: SkillScope::Project,
|
||||
provider: SkillProvider::Agents,
|
||||
icon_override: None,
|
||||
}],
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grouped_completions_share_context_without_dropping_inputs_or_tool_results() {
|
||||
let context = shared_prompt_context(
|
||||
"/repo",
|
||||
"shared-rule-probe",
|
||||
"shared-attachment-probe",
|
||||
"shared-skill-description-probe",
|
||||
);
|
||||
let mut params = RequestParams::new_for_test();
|
||||
params.input = (0..4)
|
||||
.map(|index| AIAgentInput::CommandCompletionAssessment {
|
||||
prompt: format!("Completion evidence for job {index}"),
|
||||
context: Arc::from(context.clone()),
|
||||
completed_command: crate::ai::agent::RunningCommand {
|
||||
command: format!("command-{index}"),
|
||||
block_id: format!("background-job:{index}").into(),
|
||||
grid_contents: String::new(),
|
||||
cursor: String::new(),
|
||||
requested_command_id: None,
|
||||
is_alt_screen_active: false,
|
||||
},
|
||||
})
|
||||
.collect();
|
||||
// Repeated user restrictions remain distinct conversational inputs.
|
||||
params.input.extend([
|
||||
user_query_with_context("Do not deploy", context.clone()),
|
||||
user_query_with_context("Do not deploy", context),
|
||||
]);
|
||||
params.message_history = vec![galaxy_agent_core::ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: MessageContent::ToolUse {
|
||||
tool_use_id: "retained-call".to_string(),
|
||||
name: "read_files".to_string(),
|
||||
input: serde_json::json!({ "files": ["/repo/file.rs"] }),
|
||||
},
|
||||
}];
|
||||
params.tool_results = vec![ToolResult {
|
||||
call_id: "retained-call".to_string(),
|
||||
content: "retained-tool-result-probe".to_string(),
|
||||
status: ToolResultStatus::Success,
|
||||
images: Vec::new(),
|
||||
}];
|
||||
|
||||
let prepared = prepare_rig_turn(
|
||||
&config(),
|
||||
params,
|
||||
vec![ToolType::ReadFiles, ToolType::ReadSkill],
|
||||
Vec::new(),
|
||||
);
|
||||
let prompt = prepared.request.system_prompt.expect("system prompt");
|
||||
for entry in [
|
||||
"Working directory: /repo",
|
||||
"Home directory: /home/user",
|
||||
"shared-rule-probe",
|
||||
"shared-attachment-probe",
|
||||
"shared-skill-description-probe",
|
||||
] {
|
||||
assert_eq!(
|
||||
prompt.matches(entry).count(),
|
||||
1,
|
||||
"repeated context: {entry}"
|
||||
);
|
||||
}
|
||||
for index in 0..4 {
|
||||
assert!(prepared.request.messages.iter().any(|message| {
|
||||
matches!(&message.content, MessageContent::Text(text)
|
||||
if text == &format!("Completion evidence for job {index}"))
|
||||
}));
|
||||
}
|
||||
assert_eq!(
|
||||
prepared
|
||||
.request
|
||||
.messages
|
||||
.iter()
|
||||
.filter(|message| {
|
||||
matches!(&message.content, MessageContent::Text(text) if text == "Do not deploy")
|
||||
})
|
||||
.count(),
|
||||
2,
|
||||
);
|
||||
assert!(prepared.request.messages.iter().any(|message| {
|
||||
matches!(&message.content, MessageContent::ToolResult { tool_use_id, content, .. }
|
||||
if tool_use_id == "retained-call" && content == "retained-tool-result-probe")
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn distinct_context_content_keeps_its_first_order_even_with_shared_paths() {
|
||||
let first = shared_prompt_context(
|
||||
"/repo/first",
|
||||
"first-rule-probe",
|
||||
"first-attachment-probe",
|
||||
"first-skill-description-probe",
|
||||
);
|
||||
let second = shared_prompt_context(
|
||||
"/repo/second",
|
||||
"second-rule-probe",
|
||||
"second-attachment-probe",
|
||||
"second-skill-description-probe",
|
||||
);
|
||||
let mut params = RequestParams::new_for_test();
|
||||
params.input = vec![
|
||||
user_query_with_context("First instruction", first.clone()),
|
||||
user_query_with_context("Second instruction", second),
|
||||
user_query_with_context("Third instruction", first),
|
||||
];
|
||||
|
||||
let prepared = prepare_rig_turn(&config(), params, vec![ToolType::ReadSkill], Vec::new());
|
||||
let prompt = prepared.request.system_prompt.expect("system prompt");
|
||||
for (first, second) in [
|
||||
(
|
||||
"Working directory: /repo/first",
|
||||
"Working directory: /repo/second",
|
||||
),
|
||||
("first-rule-probe", "second-rule-probe"),
|
||||
("first-attachment-probe", "second-attachment-probe"),
|
||||
(
|
||||
"first-skill-description-probe",
|
||||
"second-skill-description-probe",
|
||||
),
|
||||
] {
|
||||
assert_eq!(prompt.matches(first).count(), 1);
|
||||
assert_eq!(prompt.matches(second).count(), 1);
|
||||
assert!(prompt.find(first).unwrap() < prompt.find(second).unwrap());
|
||||
}
|
||||
assert_eq!(prepared.request.messages.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn volatile_request_time_follows_the_cacheable_system_prompt_prefix() {
|
||||
let mut params = RequestParams::new_for_test();
|
||||
|
||||
Reference in New Issue
Block a user