Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: make image tool results durably prunable

Browse Source
This commit is contained in:
Hare
2026-08-11 02:44:37 +09:00
parent 38b8f26a50
commit 8d2b8b690f
15 changed files with 469 additions and 268 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/llm-engine/src/engine.rs
+21 -72
View File
@@ -19,19 +19,15 @@ use crate::{
},
llm_client::{
ClientError, ConfigWarning, LlmClient, Request, RequestConfig, ResponseStream,
ToolDefinition,
error::is_retryable,
event::Event,
retry::RetryPolicy,
transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT,
types::{ContentPart, parse_tool_arguments},
ToolDefinition, error::is_retryable, event::Event, retry::RetryPolicy,
transport::DEFAULT_FIRST_STREAM_EVENT_TIMEOUT, types::parse_tool_arguments,
},
state::{EngineState, Locked, Mutable},
timeline::event::{ErrorEvent, StatusEvent, UsageEvent},
timeline::{TextBlockCollector, ThinkingBlockCollector, Timeline, ToolCallCollector},
tool::{
Attachment, ToolCall, ToolDefinition as EngineToolDefinition, ToolError,
ToolExecutionContext, ToolOutputLimits, ToolResult, truncate_content,
ToolCall, ToolDefinition as EngineToolDefinition, ToolError, ToolExecutionContext,
ToolOutputLimits, ToolResult, truncate_content,
},
tool_server::{ToolServer, ToolServerHandle},
};
@@ -159,33 +155,6 @@ enum StreamCompletion {
Interrupted { reason: String },
}
fn project_transient_attachments(items: &mut Vec<Item>) {
let mut projected = Vec::with_capacity(items.len());
let mut pending_parts = Vec::new();
for mut item in items.drain(..) {
let is_tool_result = matches!(item, Item::ToolResult { .. });
if !is_tool_result && !pending_parts.is_empty() {
projected.push(Item::user_message_parts(std::mem::take(&mut pending_parts)));
}
if let Item::ToolResult { attachments, .. } = &item {
for attachment in attachments {
let Attachment::Image(image) = attachment;
pending_parts.push(ContentPart::image(
image.mime_type(),
Arc::from(image.data()),
));
}
}
item.clear_transient_attachments();
projected.push(item);
}
if !pending_parts.is_empty() {
projected.push(Item::user_message_parts(pending_parts));
}
*items = projected;
}
pub struct Engine<C: LlmClient, S: EngineState = Mutable> {
/// LLM client
client: C,
@@ -1280,8 +1249,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
cb(current_llm_call);
}
project_transient_attachments(&mut request_context);
// Stream LLM response
self.emit_lifecycle_trace(
current_turn,
@@ -1290,12 +1257,6 @@ impl<C: LlmClient, S: EngineState> Engine<C, S> {
items_trace_payload(&request_context, tool_definitions.len(), None, false),
);
let request = self.build_request(&tool_definitions, &request_context);
// Structured attachments are single-use provider payloads. `request`
// owns its Arc-backed copy; history is cleared before any later turn
// can persist or resend the binary body.
for item in &mut self.history {
item.clear_transient_attachments();
}
self.emit_lifecycle_trace(
current_turn,
current_llm_call,
@@ -2199,48 +2160,36 @@ fn item_kind(item: &Item) -> &'static str {
#[cfg(test)]
mod tests {
use super::*;
use crate::llm_client::types::Role;
use crate::tool::{Attachment, ImageAttachment};
use std::time::Duration;
#[test]
fn transient_tool_attachment_becomes_non_persistent_user_image_part() {
let body: Arc<[u8]> = Arc::from(&b"secret-image-body"[..]);
let mut items = vec![Item::tool_result_item_with_attachments(
fn tool_attachment_round_trips_through_durable_history_json() {
let body: Arc<[u8]> = Arc::from(&b"image-body"[..]);
let items = vec![Item::tool_result_item_with_attachments(
"call_image",
"attached",
None,
false,
vec![Attachment::Image(crate::tool::ImageAttachment::new(
vec![Attachment::Image(ImageAttachment::new(
"image/png",
body.clone(),
))],
)];
project_transient_attachments(&mut items);
assert!(matches!(
&items[0],
Item::ToolResult { attachments, .. } if attachments.is_empty()
));
match &items[1] {
Item::Message {
role: Role::User,
content,
..
} => match &content[0] {
ContentPart::Image { media_type, source } => {
assert_eq!(media_type, "image/png");
assert_eq!(source.data(), body.as_ref());
assert_eq!(source.bytes(), body.len());
assert_eq!(content[0].as_text(), "[image attachment omitted]");
}
other => panic!("unexpected content part: {other:?}"),
},
other => panic!("unexpected projected item: {other:?}"),
}
let persisted = serde_json::to_string(&items).unwrap();
assert!(!persisted.contains("secret-image-body"));
assert!(!persisted.contains("base64"));
assert!(persisted.contains("aW1hZ2UtYm9keQ=="));
let restored: Vec<Item> = serde_json::from_str(&persisted).unwrap();
assert_eq!(restored, items);
assert!(matches!(
&restored[0],
Item::ToolResult { attachments, .. }
if matches!(
attachments.as_slice(),
[Attachment::Image(image)]
if image.mime_type() == "image/png" && image.data() == body.as_ref()
)
));
}
#[tokio::test]
crates/llm-engine/src/llm_client/scheme/anthropic/request.rs
-3
View File
@@ -268,9 +268,6 @@ impl AnthropicScheme {
.iter()
.map(|p| match p {
ContentPart::Text { text } => AnthropicContentPart::text(text.clone()),
ContentPart::Image { .. } => {
AnthropicContentPart::text(p.as_text().to_string())
}
ContentPart::Refusal { refusal } => {
AnthropicContentPart::text(refusal.clone())
}
crates/llm-engine/src/llm_client/scheme/openai_chat/request.rs
+71 -55
View File
@@ -5,10 +5,13 @@
use serde::Serialize;
use serde_json::Value;
use crate::llm_client::{
Request,
capability::{ModelCapability, ReasoningControl, ReasoningSupport},
types::{ContentPart, Item, Role, ToolDefinition, image_data_url, parse_tool_arguments},
use crate::{
llm_client::{
Request,
capability::{ModelCapability, ReasoningControl, ReasoningSupport},
types::{ContentPart, Item, Role, ToolDefinition, image_data_url, parse_tool_arguments},
},
tool::Attachment,
};
use super::OpenAIScheme;
@@ -185,6 +188,21 @@ impl OpenAIScheme {
/// - Assistant messages have role "assistant"
/// - Tool calls are within assistant messages as tool_calls array
/// - Tool results have role "tool" with tool_call_id
fn flush_pending_tool_result_images(
messages: &mut Vec<OpenAIMessage>,
pending_images: &mut Vec<OpenAIContentPart>,
) {
if !pending_images.is_empty() {
messages.push(OpenAIMessage {
role: "user".to_string(),
content: Some(OpenAIContent::Parts(std::mem::take(pending_images))),
tool_calls: vec![],
tool_call_id: None,
name: None,
});
}
}
fn convert_items_to_messages(
&self,
items: &[Item],
@@ -193,8 +211,15 @@ impl OpenAIScheme {
let mut messages = Vec::new();
let mut pending_tool_calls: Vec<OpenAIToolCall> = Vec::new();
let mut pending_assistant_text: Option<String> = None;
let mut pending_tool_result_images: Vec<OpenAIContentPart> = Vec::new();
for item in items {
if !matches!(item, Item::ToolResult { .. }) {
Self::flush_pending_tool_result_images(
&mut messages,
&mut pending_tool_result_images,
);
}
match item {
Item::Message { role, content, .. } => {
// Flush pending tool calls
@@ -209,41 +234,13 @@ impl OpenAIScheme {
Role::Assistant => "assistant",
Role::System => "system",
};
let has_image = matches!(role, Role::User)
&& supports_images
&& content
let message_content = OpenAIContent::Text(
content
.iter()
.any(|part| matches!(part, ContentPart::Image { .. }));
let message_content = if has_image {
OpenAIContent::Parts(
content
.iter()
.map(|part| match part {
ContentPart::Text { text } => {
OpenAIContentPart::Text { text: text.clone() }
}
ContentPart::Image { media_type, source } => {
OpenAIContentPart::ImageUrl {
image_url: ImageUrl {
url: image_data_url(media_type, source.data()),
},
}
}
ContentPart::Refusal { refusal } => OpenAIContentPart::Text {
text: refusal.clone(),
},
})
.collect(),
)
} else {
OpenAIContent::Text(
content
.iter()
.map(ContentPart::as_text)
.collect::<Vec<_>>()
.join(""),
)
};
.map(ContentPart::as_text)
.collect::<Vec<_>>()
.join(""),
);
messages.push(OpenAIMessage {
role: openai_role.to_string(),
@@ -277,19 +274,35 @@ impl OpenAIScheme {
call_id,
summary,
content,
attachments,
..
} => {
// Flush pending tool calls before tool result
// OpenAI requires every parallel tool result before a new user message.
self.flush_pending_assistant(
&mut messages,
&mut pending_tool_calls,
&mut pending_assistant_text,
);
let text = match content {
let mut text = match content {
Some(c) => format!("{summary}\n{c}"),
None => summary.clone(),
};
if supports_images {
pending_tool_result_images.extend(attachments.iter().map(|attachment| {
let Attachment::Image(image) = attachment;
OpenAIContentPart::ImageUrl {
image_url: ImageUrl {
url: image_data_url(image.mime_type(), image.data()),
},
}
}));
} else if !attachments.is_empty() {
text.push_str(&format!(
"\n[{} image attachment(s) omitted: model does not support images]",
attachments.len()
));
}
messages.push(OpenAIMessage {
role: "tool".to_string(),
content: Some(OpenAIContent::Text(text)),
@@ -317,6 +330,7 @@ impl OpenAIScheme {
&mut pending_tool_calls,
&mut pending_assistant_text,
);
Self::flush_pending_tool_result_images(&mut messages, &mut pending_tool_result_images);
messages
}
@@ -481,28 +495,27 @@ mod tests {
}
#[test]
fn parallel_tool_results_precede_synthetic_image_message() {
fn parallel_tool_results_precede_durable_image_projection() {
let scheme = OpenAIScheme::new();
let image = std::sync::Arc::<[u8]>::from(&b"\x89PNG\r\n\x1a\nbody"[..]);
let request = Request::new()
.item(Item::tool_call("call_image", "ViewImage", "{}"))
.item(Item::tool_call("call_text", "Read", "{}"))
.item(Item::tool_result_item(
.item(Item::tool_result_item_with_attachments(
"call_image",
"Attached image",
None,
false,
vec![crate::tool::Attachment::Image(
crate::tool::ImageAttachment::new("image/png", image),
)],
))
.item(Item::tool_result_item(
"call_text",
"Read text",
None,
false,
))
.item(Item::user_message_parts(vec![ContentPart::image(
"image/png",
image,
)]));
));
let json = serde_json::to_value(
&scheme
.build_request("gpt-4o", &request, &vision_cap())
@@ -518,7 +531,7 @@ mod tests {
}
#[test]
fn tool_image_is_structured_as_following_user_content_without_persisting_bytes() {
fn durable_tool_image_is_deterministically_lowered_to_following_user_content() {
let scheme = OpenAIScheme::new();
let image = std::sync::Arc::<[u8]>::from(&b"\x89PNG\r\n\x1a\nbody"[..]);
let attachment = crate::tool::Attachment::Image(crate::tool::ImageAttachment::new(
@@ -533,8 +546,8 @@ mod tests {
vec![attachment],
);
let persisted = serde_json::to_string(&item).unwrap();
assert!(!persisted.contains("base64"));
assert!(!persisted.contains("attachments"));
assert!(persisted.contains("attachments"));
let restored: Item = serde_json::from_str(&persisted).unwrap();
let request = Request::new()
.item(Item::tool_call(
@@ -542,13 +555,16 @@ mod tests {
"ViewImage",
r#"{"path":"a.png"}"#,
))
.item(item)
.item(Item::user_message_parts(vec![ContentPart::image(
"image/png",
image,
)]));
.item(restored);
let body = scheme.build_request("gpt-4o", &request, &vision_cap());
let json = serde_json::to_value(&body.messages).unwrap();
let rebuilt = serde_json::to_value(
&scheme
.build_request("gpt-4o", &request, &vision_cap())
.messages,
)
.unwrap();
assert_eq!(rebuilt, json);
assert_eq!(json[0]["role"], "assistant");
assert_eq!(json[1]["role"], "tool");
crates/llm-engine/src/llm_client/scheme/openai_responses/request.rs
+57 -29
View File
@@ -7,14 +7,31 @@
use serde::{Serialize, Serializer};
use serde_json::Value;
use crate::llm_client::{
Request,
capability::{ModelCapability, ReasoningControl, ReasoningSupport},
types::{ContentPart, Item, Role, ToolDefinition, image_data_url, parse_tool_arguments},
use crate::{
llm_client::{
Request,
capability::{ModelCapability, ReasoningControl, ReasoningSupport},
types::{ContentPart, Item, Role, ToolDefinition, image_data_url, parse_tool_arguments},
},
tool::Attachment,
};
use super::OpenAIResponsesScheme;
#[derive(Debug, Serialize)]
#[serde(untagged)]
pub(crate) enum FunctionCallOutputBody {
Text(String),
ContentItems(Vec<FunctionCallOutputContentItem>),
}
#[derive(Debug, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub(crate) enum FunctionCallOutputContentItem {
InputText { text: String },
InputImage { image_url: String },
}
/// `/v1/responses` のリクエスト body。
#[derive(Debug, Serialize)]
pub(crate) struct ResponsesRequest {
@@ -89,7 +106,10 @@ pub(crate) enum InputItem {
arguments: String,
},
/// function tool の結果（user 側）。
FunctionCallOutput { call_id: String, output: String },
FunctionCallOutput {
call_id: String,
output: FunctionCallOutputBody,
},
/// reasoning item。`encrypted_content` があれば必ず添える。
Reasoning {
#[serde(skip_serializing_if = "Option::is_none")]
@@ -114,10 +134,6 @@ pub(crate) enum InputContent {
/// user / developer 側のテキスト
InputText { text: String },
/// user 側の画像
InputImage {
image_url: String,
detail: &'static str,
},
/// assistant 側のテキスト
OutputText { text: String },
}
@@ -249,15 +265,6 @@ fn convert_items_to_input(items: &[Item], supports_images: bool) -> Vec<InputIte
.iter()
.map(|part| match part {
ContentPart::Text { text } => text_variant(text.clone()),
ContentPart::Image { media_type, source }
if matches!(role, Role::User) && supports_images =>
{
InputContent::InputImage {
image_url: image_data_url(media_type, source.data()),
detail: "auto",
}
}
ContentPart::Image { .. } => text_variant(part.as_text().to_string()),
ContentPart::Refusal { refusal } => text_variant(refusal.clone()),
})
.collect();
@@ -284,12 +291,30 @@ fn convert_items_to_input(items: &[Item], supports_images: bool) -> Vec<InputIte
call_id,
summary,
content,
attachments,
..
} => {
let output = match content {
let text = match content {
Some(c) => format!("{summary}\n{c}"),
None => summary.clone(),
};
let output = if attachments.is_empty() {
FunctionCallOutputBody::Text(text)
} else if supports_images {
let mut parts = vec![FunctionCallOutputContentItem::InputText { text }];
parts.extend(attachments.iter().map(|attachment| {
let Attachment::Image(image) = attachment;
FunctionCallOutputContentItem::InputImage {
image_url: image_data_url(image.mime_type(), image.data()),
}
}));
FunctionCallOutputBody::ContentItems(parts)
} else {
FunctionCallOutputBody::Text(format!(
"{text}\n[{} image attachment(s) omitted: model does not support images]",
attachments.len()
))
};
out.push(InputItem::FunctionCallOutput {
call_id: call_id.clone(),
output,
@@ -701,7 +726,7 @@ mod tests {
}
#[test]
fn synthetic_user_image_uses_responses_input_image_content() {
fn durable_tool_image_uses_function_call_output_content_items() {
let scheme = OpenAIResponsesScheme::new();
let image = std::sync::Arc::<[u8]>::from(&b"\x89PNG\r\n\x1a\nbody"[..]);
let item = Item::tool_result_item_with_attachments(
@@ -710,32 +735,35 @@ mod tests {
None,
false,
vec![crate::tool::Attachment::Image(
crate::tool::ImageAttachment::new("image/png", image.clone()),
crate::tool::ImageAttachment::new("image/png", image),
)],
);
let persisted = serde_json::to_string(&item).unwrap();
let restored: Item = serde_json::from_str(&persisted).unwrap();
let req = Request::new()
.item(Item::tool_call(
"call_image",
"ViewImage",
r#"{"path":"a.png"}"#,
))
.item(item)
.item(Item::user_message_parts(vec![ContentPart::image(
"image/png",
image,
)]));
.item(restored);
let body = scheme.build_request("gpt-5", &req, &cap_with_reasoning());
let json = serde_json::to_value(&body).unwrap();
assert_eq!(json["input"][1]["type"], "function_call_output");
assert_eq!(json["input"][2]["type"], "message");
assert_eq!(json["input"][2]["content"][0]["type"], "input_image");
assert_eq!(json["input"].as_array().unwrap().len(), 2);
assert_eq!(json["input"][1]["output"][0]["type"], "input_text");
assert_eq!(json["input"][1]["output"][1]["type"], "input_image");
assert!(
json["input"][2]["content"][0]["image_url"]
json["input"][1]["output"][1]["image_url"]
.as_str()
.unwrap()
.starts_with("data:image/png;base64,")
);
let rebuilt =
serde_json::to_value(scheme.build_request("gpt-5", &req, &cap_with_reasoning()))
.unwrap();
assert_eq!(rebuilt["input"], json["input"]);
let mut no_vision = cap_with_reasoning();
no_vision.vision = false;
crates/llm-engine/src/llm_client/types.rs
+4 -59
View File
@@ -124,8 +124,8 @@ pub enum Item {
/// Whether the tool result represents an execution error.
#[serde(default, skip_serializing_if = "is_false")]
is_error: bool,
/// Request-local structured payloads. Never serialized or persisted.
#[serde(skip, default)]
/// Durable binary details (removed with `content` by normal pruning).
#[serde(default, skip_serializing_if = "Vec::is_empty")]
attachments: Vec<Attachment>,
},
@@ -264,7 +264,7 @@ impl Item {
Self::tool_result_item_with_attachments(call_id, summary, content, is_error, Vec::new())
}
/// Create a tool result item with request-local structured attachments.
/// Create a tool result item with durable, prunable structured attachments.
pub fn tool_result_item_with_attachments(
call_id: impl Into<String>,
summary: impl Into<String>,
@@ -282,13 +282,6 @@ impl Item {
}
}
/// Drop request-local attachments after constructing the provider request.
pub fn clear_transient_attachments(&mut self) {
if let Self::ToolResult { attachments, .. } = self {
attachments.clear();
}
}
/// Create a tool result item with summary and content.
pub fn tool_result_with_content(
call_id: impl Into<String>,
@@ -457,38 +450,6 @@ pub fn parse_tool_arguments(arguments: &str) -> serde_json::Value {
// Content Parts - Components within message items
// ============================================================================
#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ImageSource {
bytes: usize,
#[serde(skip, default)]
data: Arc<[u8]>,
}
impl ImageSource {
pub fn new(data: Arc<[u8]>) -> Self {
Self {
bytes: data.len(),
data,
}
}
pub fn data(&self) -> &[u8] {
&self.data
}
pub fn bytes(&self) -> usize {
self.bytes
}
}
impl fmt::Debug for ImageSource {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ImageSource")
.field("bytes", &self.bytes)
.finish()
}
}
/// Content part within a message item
///
/// Text content is role-agnostic; the containing Item's Role determines
@@ -502,12 +463,6 @@ pub enum ContentPart {
text: String,
},
/// Request-local image content. The source bytes are never serialized.
Image {
media_type: String,
source: ImageSource,
},
/// Refusal content (for assistant messages)
Refusal {
/// The refusal message
@@ -528,20 +483,10 @@ impl ContentPart {
}
}
pub fn image(media_type: impl Into<String>, data: Arc<[u8]>) -> Self {
Self::Image {
media_type: media_type.into(),
source: ImageSource::new(data),
}
}
/// Get a bounded textual projection. Image content is represented by an
/// explicit placeholder rather than silently becoming an empty string,
/// filesystem path, data URL, or base64 body.
/// Get a textual projection of the content part.
pub fn as_text(&self) -> &str {
match self {
Self::Text { text } => text,
Self::Image { .. } => "[image attachment omitted]",
Self::Refusal { refusal } => refusal,
}
}
crates/llm-engine/src/prune.rs
+48 -9
View File
@@ -110,25 +110,30 @@ impl Default for PruneConfig {
}
}
/// Set `content = None` on each `Item::ToolResult` at the given indices.
/// Remove detailed text and attachments from each `Item::ToolResult` at the given indices.
///
/// Returns the number of items that were actually modified — items that
/// are already content-less are counted as 0. Intended for use on a
/// request-context clone (never on a persistent history).
/// The mandatory summary remains. Returns the number of items that were actually
/// modified — results that already contain no detail are counted as 0. Intended
/// for use on a request-context clone (never on a persistent history).
pub fn project(items: &mut [Item], indices: &[usize]) -> usize {
let mut count = 0;
for &i in indices {
if let Item::ToolResult { content, .. } = &mut items[i]
&& content.is_some()
if let Item::ToolResult {
content,
attachments,
..
} = &mut items[i]
&& (content.is_some() || !attachments.is_empty())
{
*content = None;
attachments.clear();
count += 1;
}
}
count
}
/// Indices of `Item::ToolResult { content: Some(_), .. }` that lie before
/// Indices of detailed `Item::ToolResult` values that lie before
/// the suffix protected by `protected_tokens`. Pure: does not mutate `items`.
///
/// Returns an empty vector when token estimates are unavailable (`NoData`) or
@@ -159,8 +164,10 @@ pub fn evaluate_candidates(
.enumerate()
.filter_map(|(i, item)| match item {
Item::ToolResult {
content: Some(_), ..
} => Some(i),
content,
attachments,
..
} if content.is_some() || !attachments.is_empty() => Some(i),
_ => None,
})
.collect();
@@ -373,6 +380,38 @@ mod tests {
}
}
#[test]
fn project_drops_image_detail_but_keeps_summary_and_persistent_source() {
let original = vec![Item::tool_result_item_with_attachments(
"call_image",
"Attached image/png image (12 bytes)",
None,
false,
vec![crate::tool::Attachment::Image(
crate::tool::ImageAttachment::new("image/png", b"image-body".to_vec()),
)],
)];
let mut request_context = original.clone();
let estimates = uniform_estimates(&original, 100);
assert_eq!(prunable_indices(&original, 0, &estimates), vec![0]);
assert_eq!(project(&mut request_context, &[0]), 1);
assert!(matches!(
&request_context[0],
Item::ToolResult {
summary,
content: None,
attachments,
..
} if summary == "Attached image/png image (12 bytes)" && attachments.is_empty()
));
assert!(matches!(
&original[0],
Item::ToolResult { attachments, .. } if attachments.len() == 1
));
}
#[test]
fn project_skips_already_pruned_items() {
// indices points at an item whose content is already None.
crates/llm-engine/src/tool.rs
+41 -13
View File
@@ -6,7 +6,8 @@
use std::{collections::HashMap, fmt, sync::Arc};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _};
use serde_json::Value;
use thiserror::Error;
@@ -120,12 +121,39 @@ impl fmt::Debug for ImageAttachment {
}
}
/// Request-local binary payload emitted by a tool.
///
/// Attachments are deliberately excluded from serde. They may be projected into
/// the immediately following provider request, but never into persisted history,
/// protocol events, logs, or telemetry.
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Serialize, Deserialize)]
struct ImageAttachmentWire {
mime_type: String,
data: String,
}
impl Serialize for ImageAttachment {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
ImageAttachmentWire {
mime_type: self.mime_type.clone(),
data: STANDARD.encode(self.data.as_ref()),
}
.serialize(serializer)
}
}
impl<'de> Deserialize<'de> for ImageAttachment {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let wire = ImageAttachmentWire::deserialize(deserializer)?;
let data = STANDARD.decode(wire.data).map_err(D::Error::custom)?;
Ok(Self::new(wire.mime_type, data))
}
}
/// Durable binary detail emitted by a tool and handled by normal ToolResult pruning.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload", rename_all = "snake_case")]
pub enum Attachment {
Image(ImageAttachment),
}
@@ -133,8 +161,8 @@ pub enum Attachment {
/// Tool execution result.
///
/// Every output has a mandatory `summary` (1-2 lines) that persists in
/// conversation history even after pruning. The optional `content` carries
/// full text details. `attachments` are request-local and are never serialized.
/// conversation history even after pruning. Optional text and binary details are
/// committed to history and may later be omitted only by normal ToolResult pruning.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolOutput {
/// Short summary (1-2 lines). Always remains in history.
@@ -142,8 +170,8 @@ pub struct ToolOutput {
/// Detailed text output. Removed by Prune when old enough.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
/// Structured binary payloads for the immediately following model request.
#[serde(skip, default)]
/// Durable binary details handled by the same pruning lifecycle as `content`.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub attachments: Vec<Attachment>,
}
@@ -409,8 +437,8 @@ pub struct ToolResult {
/// Whether this is an error
#[serde(default)]
pub is_error: bool,
/// Request-local structured payloads. Never serialized or persisted.
#[serde(skip, default)]
/// Durable binary details (prunable with `content`).
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub attachments: Vec<Attachment>,
}
crates/session-store/Cargo.toml
+1
View File
@@ -6,6 +6,7 @@ edition.workspace = true
license.workspace = true
[dependencies]
base64.workspace = true
llm-engine = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
crates/session-store/src/logged_item.rs
+73 -11
View File
@@ -12,13 +12,36 @@
//! `Reasoning::encrypted_content` is preserved because OpenAI Responses ZDR
//! requires it on stateless re-send.
use llm_engine::llm_client::types::{ContentPart, Item, Role};
use serde::{Deserialize, Serialize};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use llm_engine::{
llm_client::types::{ContentPart, Item, Role},
tool::{Attachment, ImageAttachment},
};
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _};
fn is_false(value: &bool) -> bool {
!*value
}
mod base64_bytes {
use super::*;
pub fn serialize<S>(data: &[u8], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&STANDARD.encode(data))
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
let encoded = String::deserialize(deserializer)?;
STANDARD.decode(encoded).map_err(D::Error::custom)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum LoggedItem {
@@ -36,6 +59,8 @@ pub enum LoggedItem {
summary: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
content: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
attachments: Vec<LoggedAttachment>,
#[serde(default, skip_serializing_if = "is_false")]
is_error: bool,
},
@@ -67,6 +92,16 @@ pub enum LoggedContentPart {
Refusal { refusal: String },
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum LoggedAttachment {
Image {
mime_type: String,
#[serde(with = "base64_bytes")]
data: Vec<u8>,
},
}
// ---------------------------------------------------------------------------
// Item ↔ LoggedItem
// ---------------------------------------------------------------------------
@@ -92,12 +127,14 @@ impl From<&Item> for LoggedItem {
call_id,
summary,
content,
attachments,
is_error,
..
} => Self::ToolResult {
call_id: call_id.clone(),
summary: summary.clone(),
content: content.clone(),
attachments: attachments.iter().map(LoggedAttachment::from).collect(),
is_error: *is_error,
},
Item::Reasoning {
@@ -146,6 +183,7 @@ impl From<LoggedItem> for Item {
call_id,
summary,
content,
attachments,
is_error,
} => Item::ToolResult {
id: None,
@@ -153,7 +191,7 @@ impl From<LoggedItem> for Item {
summary,
content,
is_error,
attachments: Vec::new(),
attachments: attachments.into_iter().map(Attachment::from).collect(),
},
LoggedItem::Reasoning {
text,
@@ -214,9 +252,6 @@ impl From<&ContentPart> for LoggedContentPart {
fn from(part: &ContentPart) -> Self {
match part {
ContentPart::Text { text } => Self::Text { text: text.clone() },
ContentPart::Image { .. } => Self::Text {
text: part.as_text().to_string(),
},
ContentPart::Refusal { refusal } => Self::Refusal {
refusal: refusal.clone(),
},
@@ -233,6 +268,27 @@ impl From<LoggedContentPart> for ContentPart {
}
}
impl From<&Attachment> for LoggedAttachment {
fn from(attachment: &Attachment) -> Self {
match attachment {
Attachment::Image(image) => Self::Image {
mime_type: image.mime_type().to_string(),
data: image.data().to_vec(),
},
}
}
}
impl From<LoggedAttachment> for Attachment {
fn from(attachment: LoggedAttachment) -> Self {
match attachment {
LoggedAttachment::Image { mime_type, data } => {
Self::Image(ImageAttachment::new(mime_type, data))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -375,7 +431,7 @@ mod tests {
}
#[test]
fn tool_result_persistence_drops_binary_attachments() {
fn tool_result_persistence_round_trips_binary_attachments() {
let original = Item::tool_result_item_with_attachments(
"call_image",
"attached",
@@ -390,11 +446,17 @@ mod tests {
);
let logged: LoggedItem = (&original).into();
let json = serde_json::to_string(&logged).unwrap();
assert!(!json.contains("secret-image-body"));
assert!(!json.contains("attachments"));
assert!(json.contains("attachments"));
assert!(json.contains("c2VjcmV0LWltYWdlLWJvZHk="));
match Item::from(logged) {
Item::ToolResult { attachments, .. } => assert!(attachments.is_empty()),
let restored: LoggedItem = serde_json::from_str(&json).unwrap();
match Item::from(restored) {
Item::ToolResult { attachments, .. } => assert!(matches!(
attachments.as_slice(),
[Attachment::Image(image)]
if image.mime_type() == "image/png"
&& image.data() == b"secret-image-body"
)),
other => panic!("unexpected variant: {other:?}"),
}
}
crates/session-store/src/segment_log.rs
+30
View File
@@ -473,6 +473,36 @@ mod tests {
assert!(state.history[2].is_tool_result());
}
#[test]
fn replay_restores_durable_tool_image_detail() {
let entry = LogEntry::ToolResult {
ts: 3500,
item: Item::tool_result_item_with_attachments(
"call_image",
"attached",
None,
false,
vec![llm_engine::tool::Attachment::Image(
llm_engine::tool::ImageAttachment::new("image/png", b"durable-image".to_vec()),
)],
)
.into(),
};
let persisted = serde_json::to_string(&entry).unwrap();
let restored_entry: LogEntry = serde_json::from_str(&persisted).unwrap();
let state = collect_state(&[restored_entry]);
assert!(matches!(
&state.history[0],
Item::ToolResult { attachments, .. }
if matches!(
attachments.as_slice(),
[llm_engine::tool::Attachment::Image(image)]
if image.data() == b"durable-image"
)
));
}
#[test]
fn replay_config_changed() {
let state = collect_state(&[
crates/tools/tests/integration.rs
+5 -2
View File
@@ -101,7 +101,7 @@ fn meta_has_description_and_schema() {
}
#[tokio::test]
async fn view_image_reads_scoped_bytes_without_serializing_them_as_text() {
async fn view_image_reads_scoped_bytes_into_durable_tool_detail() {
let dir = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let scope = scope_with_spill(dir.path(), spill.path());
@@ -119,7 +119,10 @@ async fn view_image_reads_scoped_bytes_without_serializing_them_as_text() {
assert_eq!(image.data(), png);
let serialized = serde_json::to_string(&output).unwrap();
assert!(!serialized.contains("private-image-body"));
assert!(!serialized.contains("attachments"));
assert!(serialized.contains("attachments"));
let restored: llm_engine::tool::ToolOutput = serde_json::from_str(&serialized).unwrap();
let llm_engine::tool::Attachment::Image(restored_image) = &restored.attachments[0];
assert_eq!(restored_image.data(), png);
let escaped = call_err(&tool, json!({ "path": "../outside.png" })).await;
assert!(escaped.to_string().contains("scope") || escaped.to_string().contains("path"));
crates/worker/src/session_capture.rs
+66 -15
View File
@@ -271,9 +271,20 @@ impl SessionCapture {
});
}
Item::ToolResult {
summary, content, ..
summary,
content,
attachments,
..
} => {
let text = format!("{summary}\n{}", content.as_deref().unwrap_or_default());
let attachment_marker = if attachments.is_empty() {
String::new()
} else {
format!("\n[{} image attachment(s)]", attachments.len())
};
let text = format!(
"{summary}\n{}{attachment_marker}",
content.as_deref().unwrap_or_default(),
);
index.push(ReferenceEntry {
id: SessionEntryRef::new(idx),
entry_range,
@@ -506,21 +517,29 @@ fn render_item(
Item::ToolResult {
summary,
content,
attachments,
is_error,
..
} => match detail {
ReadDetail::Compact => format!(
"[{} ToolOutput{}]\nsummary: {summary}\ncontent: (omitted)",
entry.id,
if *is_error { " error" } else { "" }
),
ReadDetail::Full => format!(
"[{} ToolOutput{}]\nsummary: {summary}\ncontent: {}",
entry.id,
if *is_error { " error" } else { "" },
content.as_deref().unwrap_or_default()
),
},
} => {
let attachment_line = if attachments.is_empty() {
String::new()
} else {
format!("\nattachments: {} image(s)", attachments.len())
};
match detail {
ReadDetail::Compact => format!(
"[{} ToolOutput{}]\nsummary: {summary}\ncontent: (omitted){attachment_line}",
entry.id,
if *is_error { " error" } else { "" },
),
ReadDetail::Full => format!(
"[{} ToolOutput{}]\nsummary: {summary}\ncontent: {}{attachment_line}",
entry.id,
if *is_error { " error" } else { "" },
content.as_deref().unwrap_or_default(),
),
}
}
Item::Reasoning { .. } => format!("[{} Reasoning omitted]", entry.id),
};
truncate_chars(&text, max_bytes)
@@ -633,6 +652,38 @@ mod tests {
);
}
#[test]
fn tool_image_projection_exposes_only_bounded_metadata() {
let view = SessionCapture::new(
"segment-1",
vec![Item::tool_result_item_with_attachments(
"c1",
"attached",
None,
false,
vec![llm_engine::tool::Attachment::Image(
llm_engine::tool::ImageAttachment::new(
"image/png",
b"private-image-body".to_vec(),
),
)],
)],
);
let result = view.read(
ReadSelector::Id("E00000000"),
ReadOptions {
include_tools: true,
detail: ReadDetail::Full,
..ReadOptions::default()
},
);
assert_eq!(result.entries.len(), 1);
assert!(result.entries[0].text.contains("attachments: 1 image(s)"));
assert!(!result.entries[0].text.contains("private-image-body"));
assert!(!result.entries[0].text.contains("cHJpdmF0ZS"));
}
#[test]
fn system_prompt_and_reasoning_are_excluded_from_every_projection() {
let view = SessionCapture::new(
crates/worker/src/worker.rs
+2
View File
@@ -6507,6 +6507,7 @@ mod build_summary_prompt_tests {
call_id: "call-1".into(),
summary: "wrote a file".into(),
content: None,
attachments: Vec::new(),
is_error: false,
},
},
@@ -6548,6 +6549,7 @@ mod build_summary_prompt_tests {
call_id: "call-1".into(),
summary: "wrote a file".into(),
content: None,
attachments: Vec::new(),
is_error: false,
},
},