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:
Ryan Ward
2026-11-03 16:20:18 +00:00
parent be5433177d
commit 62df16c605
5 changed files with 365 additions and 6 deletions
@@ -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", &params.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", &params.computer_use_model.as_str())?;
map.serialize_entry("input", &DebugPayload(&params.input))?;
map.serialize_entry("tool_results", &DebugPayload(&params.tool_results))?;
map.serialize_entry("tasks", &DebugPayload(&params.tasks))?;
map.serialize_entry("message_history", &params.message_history)?;
map.serialize_entry("progressive_summary", &params.progressive_summary)?;
map.serialize_entry("tool_result_archive", &params.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", &params.global_rules)?;
map.serialize_entry("mcp_context", &DebugPayload(&params.mcp_context))?;
map.serialize_entry("autonomy_level", &DebugPayload(&params.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(&params.tasks))?;
map.serialize_entry("tool_result_archive", &params.tool_result_archive)?;
SerializeMap::end(map)
})
}
+83
View File
@@ -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(
+76
View File
@@ -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());
}
+12
View File
@@ -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 {
+171
View File
@@ -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();