use super::{SessionEvent, SessionEventKind};
use crate::output::sanitize_display_text;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::{HashMap, HashSet};
const PREVIEW_BYTES: usize = 8192;
const PREVIEW_LINES: usize = 128;
#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct AppliedPreview {
schema_version: u8,
diff: String,
truncated: bool,
}
impl AppliedPreview {
pub(crate) fn select(diff: &str, already_truncated: bool) -> Self {
let safe_diff = sanitize_display_text(diff);
let mut selected = Self {
schema_version: 1,
diff: String::new(),
truncated: already_truncated,
};
for (index, safe) in safe_diff.lines().enumerate() {
if index == PREVIEW_LINES {
selected.truncated = true;
break;
}
let mut end = safe.len().min(1024);
while !safe.is_char_boundary(end) {
end -= 1;
}
selected.truncated |= end != safe.len();
let previous = selected.diff.len();
selected.diff.push_str(&safe[..end]);
selected.diff.push('\n');
if serde_json::to_vec(&selected)
.expect("selected preview")
.len()
> PREVIEW_BYTES - 64
{
selected.diff.truncate(previous);
selected.truncated = true;
break;
}
}
selected
}
pub(crate) fn apply(&self, result: &mut crate::tools::ToolResult) {
result.display.edit_diff = Some(self.diff.clone());
result.metadata["diff_truncated"] =
Value::Bool(self.truncated || result.metadata["diff_truncated"] == true);
}
pub(crate) fn into_value(self) -> Value {
serde_json::to_value(self).expect("selected preview")
}
fn parse(value: &Value) -> Option<Self> {
if value.as_object()?.len() != 3
|| value["schema_version"].as_u64()? != 1
|| value["diff"].as_str()?.len() > PREVIEW_BYTES
|| value["truncated"].as_bool().is_none()
{
return None;
}
if serde_json::to_vec(value).ok()?.len() > PREVIEW_BYTES {
return None;
}
let preview: Self = serde_json::from_value(value.clone()).ok()?;
(preview.schema_version == 1
&& preview.diff.lines().count() <= PREVIEW_LINES
&& preview.diff.lines().all(|line| line.len() <= 1024)
&& sanitize_display_text(&preview.diff) == preview.diff)
.then_some(preview)
}
}
pub(crate) struct RecordedPreviews<'a> {
records: HashMap<(&'a str, &'a str), Option<(usize, &'a SessionEvent)>>,
result_ids: HashSet<(&'a str, &'a str)>,
duplicates: HashSet<(&'a str, &'a str)>,
}
impl<'a> RecordedPreviews<'a> {
pub(crate) fn new(events: &'a [SessionEvent]) -> Self {
let mut selected = Self {
records: HashMap::new(),
result_ids: HashSet::new(),
duplicates: HashSet::new(),
};
for (index, event) in events.iter().enumerate() {
let Some(id) = event.payload["display_id"]
.as_str()
.filter(|id| !id.is_empty() && id.len() <= 64)
else {
continue;
};
let key = (event.session_id.as_str(), id);
match event.kind() {
Some(
SessionEventKind::ToolResult
| SessionEventKind::ToolDisplayResult
| SessionEventKind::CodeModeToolResult,
) => {
if !selected.result_ids.insert(key) {
selected.duplicates.insert(key);
}
}
Some(SessionEventKind::ToolAppliedPreview) => {
selected
.records
.entry(key)
.and_modify(|record| *record = None)
.or_insert(Some((index, event)));
}
_ => {}
}
}
selected
}
pub(crate) fn apply(
&self,
index: usize,
event: &SessionEvent,
result: &mut crate::tools::ToolResult,
) {
let Some(id) = event.payload["display_id"].as_str() else {
return;
};
let key = (event.session_id.as_str(), id);
if self.duplicates.contains(&key) {
return;
}
let Some(Some((record_index, record))) = self.records.get(&key) else {
return;
};
if *record_index <= index
|| record.payload["parent_call_id"] != event.payload["parent_call_id"]
|| record.payload["sequence"] != event.payload["sequence"]
{
return;
}
if !matches!(result.tool_name.as_str(), "hash_edit" | "write") {
return;
}
if let Some(preview) = AppliedPreview::parse(&record.payload["preview"]) {
preview.apply(result);
}
}
}