fix: harden uploaded attachment retention and replay

This commit is contained in:
2026-09-03 05:50:47 +09:00
parent 6c5b8315a3
commit 8a70f3cb26
14 changed files with 321 additions and 29 deletions
Generated
+1
View File
@@ -4412,6 +4412,7 @@ dependencies = [
"thiserror 2.0.18", "thiserror 2.0.18",
"tokio", "tokio",
"tracing", "tracing",
"unicode-normalization",
"uuid", "uuid",
] ]
+1
View File
@@ -16,6 +16,7 @@ uuid = { workspace = true, features = ["v7", "serde"] }
thiserror = { workspace = true } thiserror = { workspace = true }
protocol = { workspace = true } protocol = { workspace = true }
tracing.workspace = true tracing.workspace = true
unicode-normalization = "0.1.25"
[dev-dependencies] [dev-dependencies]
async-trait = { workspace = true } async-trait = { workspace = true }
+51 -2
View File
@@ -20,8 +20,8 @@ use crate::paste_artifact::{read_from_dir, write_to_dir};
use crate::segment_log::LogEntry; use crate::segment_log::LogEntry;
use crate::store::{Store, StoreError}; use crate::store::{Store, StoreError};
use crate::uploaded_file::{ use crate::uploaded_file::{
bind_uploaded_file, delete_uploaded_file, read_uploaded_file, read_uploaded_file_by_id, bind_uploaded_file, delete_uncommitted_uploaded_files, delete_uploaded_file,
write_uploaded_file, read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file,
}; };
use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits}; use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits};
use protocol::{PasteArtifactRef, UploadedFileRef}; use protocol::{PasteArtifactRef, UploadedFileRef};
@@ -459,6 +459,14 @@ impl Store for FsStore {
delete_uploaded_file(&self.paste_artifact_dir(session_id), artifact_id) delete_uploaded_file(&self.paste_artifact_dir(session_id), artifact_id)
} }
fn delete_uncommitted_uploaded_files(&self, session_id: SessionId) -> Result<u64, StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
delete_uncommitted_uploaded_files(&self.paste_artifact_dir(session_id))
}
fn append_trace( fn append_trace(
&self, &self,
session_id: SessionId, session_id: SessionId,
@@ -678,6 +686,47 @@ mod tests {
store.write_uploaded_file(session_id, "notes.txt", "not a type", b"x", limits), store.write_uploaded_file(session_id, "notes.txt", "not a type", b"x", limits),
Err(StoreError::InvalidUploadedFileMediaType) Err(StoreError::InvalidUploadedFileMediaType)
)); ));
assert!(matches!(
store.write_uploaded_file(
session_id,
"safe\u{202e}txt.exe",
"text/plain",
b"x",
limits
),
Err(StoreError::InvalidUploadedFileName)
));
assert!(matches!(
store.write_uploaded_file(session_id, "image.png", "image/png", b"not a png", limits),
Err(StoreError::ArtifactIntegrityMismatch)
));
let pending = store
.write_uploaded_file(session_id, "Readme.txt", "text/plain", b"x", limits)
.unwrap();
let replay = store
.write_uploaded_file(session_id, "Readme.txt", "text/plain", b"x", limits)
.unwrap();
assert_eq!(replay.artifact_id, pending.artifact_id);
assert!(matches!(
store.write_uploaded_file(session_id, "README.txt", "text/plain", b"changed", limits),
Err(StoreError::InvalidUploadedFileName)
));
assert!(matches!(
store.write_uploaded_file(session_id, "README.txt", "text/plain", b"y", limits),
Err(StoreError::InvalidUploadedFileName)
));
let bound = store
.bind_uploaded_file(session_id, &pending, "entry-upload")
.unwrap();
let other = store
.write_uploaded_file(session_id, "other.txt", "text/plain", b"z", limits)
.unwrap();
assert_eq!(
store.delete_uncommitted_uploaded_files(session_id).unwrap(),
1
);
assert!(store.read_uploaded_file(session_id, &other).is_err());
assert_eq!(store.read_uploaded_file(session_id, &bound).unwrap(), b"x");
store store
.write_paste_artifact( .write_paste_artifact(
session_id, session_id,
+1 -1
View File
@@ -69,7 +69,7 @@ pub use system_item::{
}; };
pub use uploaded_file::{ pub use uploaded_file::{
DEFAULT_MAX_FILES_PER_SUBMISSION, DEFAULT_MAX_SESSION_ARTIFACT_BYTES, DEFAULT_MAX_FILES_PER_SUBMISSION, DEFAULT_MAX_SESSION_ARTIFACT_BYTES,
DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits, DEFAULT_MAX_SESSION_UPLOADED_FILES, DEFAULT_MAX_UPLOADED_FILE_BYTES, UploadedFileLimits,
}; };
pub use worker_metadata::{ pub use worker_metadata::{
CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata, CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata,
@@ -41,6 +41,24 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
entries.clear(); entries.clear();
extend_history(&mut entries, history, None, *ts); extend_history(&mut entries, history, None, *ts);
} }
LogEntry::InputSegmentsCheckpoint { user_segments, .. } => {
let mut segments = user_segments.iter();
for entry in &mut entries {
let is_user = matches!(
&entry.data,
SessionSnapshotEntryData::UserInput { .. }
| SessionSnapshotEntryData::Message {
role: SessionMessageRole::User,
..
}
);
if is_user && let Some(checkpoint) = segments.next() {
entry.data = SessionSnapshotEntryData::UserInput {
segments: checkpoint.clone(),
};
}
}
}
LogEntry::AnnotatedUserInput { LogEntry::AnnotatedUserInput {
ts, ts,
segments, segments,
@@ -357,6 +375,63 @@ mod tests {
assert!(json.contains("visible")); assert!(json.contains("visible"));
} }
#[test]
fn compacted_checkpoint_restores_uploaded_file_segments() {
let session_id = crate::new_session_id();
let user_entry_id = LoggedSessionHistoryEntryId::new();
let file = protocol::UploadedFileRef {
artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b3".into(),
file_name: "notes.md".into(),
media_type: "text/markdown".into(),
created_at_ms: 7,
availability: protocol::UploadedFileAvailability::Available,
byte_len: 12,
sha256: "a".repeat(64),
source_entry_id: Some(user_entry_id.0.clone()),
};
let segment = Segment::UploadedFile { file };
let log = vec![
LogEntry::AnnotatedSegmentStart {
ts: 10,
session_id,
system_prompt: None,
config: RequestConfig::default(),
history: vec![LoggedHistoryEntry {
item: LoggedItem::Message {
role: LoggedRole::User,
content: vec![LoggedContentPart::Text {
text: "[Attached file: notes.md]".into(),
}],
},
metadata: LoggedSessionHistoryMetadata {
entry_id: user_entry_id,
origin: LoggedSessionHistoryOrigin::HumanInput {
account_id: "account-1".into(),
},
derivation: None,
},
}],
forked_from: None,
compacted_from: Some(crate::SegmentOrigin {
segment_id: crate::new_segment_id(),
at_turn_index: 1,
}),
},
LogEntry::InputSegmentsCheckpoint {
ts: 10,
user_segments: vec![vec![segment.clone()]],
},
];
let snapshot = project_current_session_snapshot(&log);
assert_eq!(
snapshot.entries[0].data,
SessionSnapshotEntryData::UserInput {
segments: vec![segment]
}
);
}
#[test] #[test]
fn annotated_user_input_attaches_segments_to_first_user_role_entry_for_any_origin() { fn annotated_user_input_attaches_segments_to_first_user_role_entry_for_any_origin() {
let session_id = crate::new_session_id(); let session_id = crate::new_session_id();
+11
View File
@@ -63,6 +63,14 @@ pub enum LogEntry {
compacted_from: Option<SegmentOrigin>, compacted_from: Option<SegmentOrigin>,
}, },
/// Typed user-segment projection accompanying a compacted or forked
/// SegmentStart history snapshot. This keeps attachment identity and
/// metadata aligned with retained user entries without embedding bodies.
InputSegmentsCheckpoint {
ts: u64,
user_segments: Vec<Vec<Segment>>,
},
/// IDLE → active marker. Records the start of a new self-driving /// IDLE → active marker. Records the start of a new self-driving
/// cycle (Invoke range). The range extends implicitly until the /// cycle (Invoke range). The range extends implicitly until the
/// next `Invoke` entry; this entry carries the trigger only — the /// next `Invoke` entry; this entry carries the trigger only — the
@@ -273,6 +281,9 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
.map(|entry| Item::from(entry.item)) .map(|entry| Item::from(entry.item))
.collect(); .collect();
} }
LogEntry::InputSegmentsCheckpoint { user_segments, .. } => {
state.user_segments = user_segments.clone();
}
LogEntry::Invoke { .. } => { LogEntry::Invoke { .. } => {
// A terminal run record below clears or refines this. If the // A terminal run record below clears or refines this. If the
// log ends first, restore must treat the turn as interrupted. // log ends first, restore must treat the turn as interrupted.
+4
View File
@@ -221,6 +221,10 @@ pub trait Store: Send + Sync {
Err(StoreError::PasteArtifactUnsupported) Err(StoreError::PasteArtifactUnsupported)
} }
fn delete_uncommitted_uploaded_files(&self, _session_id: SessionId) -> Result<u64, StoreError> {
Ok(0)
}
/// Append a trace entry to the debug event trace file. /// Append a trace entry to the debug event trace file.
fn append_trace( fn append_trace(
&self, &self,
+115 -6
View File
@@ -9,6 +9,7 @@ use fs4::fs_std::FileExt;
use protocol::{UploadedFileAvailability, UploadedFileRef}; use protocol::{UploadedFileAvailability, UploadedFileRef};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
use unicode_normalization::UnicodeNormalization;
use uuid::Uuid; use uuid::Uuid;
use crate::StoreError; use crate::StoreError;
@@ -18,6 +19,7 @@ type Result<T> = std::result::Result<T, StoreError>;
pub const DEFAULT_MAX_UPLOADED_FILE_BYTES: u64 = 10 * 1024 * 1024; pub const DEFAULT_MAX_UPLOADED_FILE_BYTES: u64 = 10 * 1024 * 1024;
pub const DEFAULT_MAX_SESSION_ARTIFACT_BYTES: u64 = 32 * 1024 * 1024; pub const DEFAULT_MAX_SESSION_ARTIFACT_BYTES: u64 = 32 * 1024 * 1024;
pub const DEFAULT_MAX_FILES_PER_SUBMISSION: usize = 8; pub const DEFAULT_MAX_FILES_PER_SUBMISSION: usize = 8;
pub const DEFAULT_MAX_SESSION_UPLOADED_FILES: u64 = 256;
const MAX_FILE_NAME_CHARS: usize = 255; const MAX_FILE_NAME_CHARS: usize = 255;
const MAX_MEDIA_TYPE_BYTES: usize = 127; const MAX_MEDIA_TYPE_BYTES: usize = 127;
@@ -53,9 +55,15 @@ pub(crate) fn validate_file_name(file_name: &str) -> Result<()> {
|| file_name.chars().count() > MAX_FILE_NAME_CHARS || file_name.chars().count() > MAX_FILE_NAME_CHARS
|| file_name == "." || file_name == "."
|| file_name == ".." || file_name == ".."
|| file_name || file_name.chars().any(|ch| {
.chars() ch.is_control()
.any(|ch| ch.is_control() || matches!(ch, '/' | '\\')) || matches!(
ch,
'/' | '\\' | '\u{200b}' | '\u{200c}' | '\u{200d}' | '\u{2060}' | '\u{feff}'
)
|| ('\u{202a}'..='\u{202e}').contains(&ch)
|| ('\u{2066}'..='\u{2069}').contains(&ch)
})
{ {
return Err(StoreError::InvalidUploadedFileName); return Err(StoreError::InvalidUploadedFileName);
} }
@@ -96,6 +104,35 @@ pub(crate) fn validate_media_type(media_type: &str) -> Result<()> {
Ok(()) Ok(())
} }
fn normalized_file_name(file_name: &str) -> String {
file_name.nfkc().flat_map(char::to_lowercase).collect()
}
fn validate_content(media_type: &str, content: &[u8]) -> Result<()> {
if content.is_empty() {
return Err(StoreError::InvalidUploadedFileMediaType);
}
let matches_declared_type = if media_type.starts_with("text/") {
std::str::from_utf8(content).is_ok()
} else {
match media_type {
"application/json" => serde_json::from_slice::<serde_json::Value>(content).is_ok(),
"application/pdf" => content.starts_with(b"%PDF-"),
"image/png" => content.starts_with(b"\x89PNG\r\n\x1a\n"),
"image/jpeg" => content.starts_with(&[0xff, 0xd8, 0xff]),
"image/gif" => content.starts_with(b"GIF87a") || content.starts_with(b"GIF89a"),
"image/webp" => {
content.len() >= 12 && content.starts_with(b"RIFF") && &content[8..12] == b"WEBP"
}
_ => false,
}
};
if !matches_declared_type {
return Err(StoreError::ArtifactIntegrityMismatch);
}
Ok(())
}
fn record_path(dir: &Path, artifact_id: &str) -> Result<std::path::PathBuf> { fn record_path(dir: &Path, artifact_id: &str) -> Result<std::path::PathBuf> {
let id = Uuid::parse_str(artifact_id).map_err(|_| StoreError::InvalidArtifactId)?; let id = Uuid::parse_str(artifact_id).map_err(|_| StoreError::InvalidArtifactId)?;
Ok(dir.join(format!("{id}.file.json"))) Ok(dir.join(format!("{id}.file.json")))
@@ -133,7 +170,7 @@ pub(crate) fn stored_uploaded_file_usage(dir: &Path) -> Result<(u64, u64)> {
{ {
continue; continue;
} }
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(path)?)?; let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
bytes = bytes bytes = bytes
.checked_add(stored.byte_len) .checked_add(stored.byte_len)
.ok_or(StoreError::ArtifactQuotaExceeded)?; .ok_or(StoreError::ArtifactQuotaExceeded)?;
@@ -153,7 +190,9 @@ pub(crate) fn write_uploaded_file(
) -> Result<UploadedFileRef> { ) -> Result<UploadedFileRef> {
validate_file_name(file_name)?; validate_file_name(file_name)?;
validate_media_type(media_type)?; validate_media_type(media_type)?;
validate_content(media_type, content)?;
let byte_len = u64::try_from(content.len()).map_err(|_| StoreError::ArtifactTooLarge)?; let byte_len = u64::try_from(content.len()).map_err(|_| StoreError::ArtifactTooLarge)?;
let sha256 = digest(content);
if byte_len > limits.max_file_bytes { if byte_len > limits.max_file_bytes {
return Err(StoreError::ArtifactTooLarge); return Err(StoreError::ArtifactTooLarge);
} }
@@ -166,7 +205,48 @@ pub(crate) fn write_uploaded_file(
.open(dir.join(".aggregate.lock"))?; .open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?; FileExt::lock_exclusive(&aggregate_lock)?;
let (paste_bytes, _) = crate::paste_artifact::stored_paste_usage(dir)?; let (paste_bytes, _) = crate::paste_artifact::stored_paste_usage(dir)?;
let (file_bytes, _) = stored_uploaded_file_usage(dir)?; let (file_bytes, file_count) = stored_uploaded_file_usage(dir)?;
if file_count >= DEFAULT_MAX_SESSION_UPLOADED_FILES {
return Err(StoreError::ArtifactQuotaExceeded);
}
let normalized_name = normalized_file_name(file_name);
for entry in fs::read_dir(dir)? {
let path = entry?.path();
if !path
.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name.ends_with(".file.json"))
{
continue;
}
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.source_entry_id.is_none()
&& normalized_file_name(&stored.file_name) == normalized_name
{
if stored.media_type == media_type
&& stored.byte_len == byte_len
&& stored.sha256 == sha256
{
let artifact_id = path
.file_name()
.and_then(|name| name.to_str())
.and_then(|name| name.strip_suffix(".file.json"))
.ok_or(StoreError::InvalidArtifactId)?
.to_string();
return Ok(UploadedFileRef {
artifact_id,
file_name: stored.file_name,
media_type: stored.media_type,
created_at_ms: stored.created_at_ms,
availability: UploadedFileAvailability::Available,
byte_len: stored.byte_len,
sha256: stored.sha256,
source_entry_id: None,
});
}
return Err(StoreError::InvalidUploadedFileName);
}
}
if paste_bytes if paste_bytes
.checked_add(file_bytes) .checked_add(file_bytes)
.and_then(|total| total.checked_add(byte_len)) .and_then(|total| total.checked_add(byte_len))
@@ -177,7 +257,6 @@ pub(crate) fn write_uploaded_file(
let artifact_id = Uuid::now_v7().to_string(); let artifact_id = Uuid::now_v7().to_string();
let created_at_ms = now_ms()?; let created_at_ms = now_ms()?;
let sha256 = digest(content);
let stored = StoredUploadedFile { let stored = StoredUploadedFile {
file_name: file_name.to_owned(), file_name: file_name.to_owned(),
media_type: media_type.to_owned(), media_type: media_type.to_owned(),
@@ -274,6 +353,36 @@ pub(crate) fn bind_uploaded_file(
Ok(bound) Ok(bound)
} }
pub(crate) fn delete_uncommitted_uploaded_files(dir: &Path) -> Result<u64> {
fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(dir.join(".aggregate.lock"))?;
FileExt::lock_exclusive(&aggregate_lock)?;
let mut removed = 0_u64;
for entry in fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
if !path
.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name.ends_with(".file.json"))
{
continue;
}
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
if stored.source_entry_id.is_none() {
fs::remove_file(path)?;
removed = removed
.checked_add(1)
.ok_or(StoreError::ArtifactQuotaExceeded)?;
}
}
Ok(removed)
}
pub(crate) fn delete_uploaded_file(dir: &Path, artifact_id: &str) -> Result<bool> { pub(crate) fn delete_uploaded_file(dir: &Path, artifact_id: &str) -> Result<bool> {
fs::create_dir_all(dir)?; fs::create_dir_all(dir)?;
let aggregate_lock = fs::OpenOptions::new() let aggregate_lock = fs::OpenOptions::new()
-8
View File
@@ -928,14 +928,6 @@ impl App {
Some(self.method_for_run(queued.segments)) Some(self.method_for_run(queued.segments))
} }
pub fn push_notice(&mut self, message: impl Into<String>) {
self.blocks.push(Block::Alert {
level: AlertLevel::Warn,
source: AlertSource::Worker,
message: message.into(),
});
}
pub fn push_error(&mut self, message: impl Into<String>) { pub fn push_error(&mut self, message: impl Into<String>) {
self.blocks.push(Block::Alert { self.blocks.push(Block::Alert {
level: AlertLevel::Error, level: AlertLevel::Error,
+15 -5
View File
@@ -792,17 +792,27 @@ async fn handle_terminal_event<T: Socket>(
if let Some(method) = handle_key(app, key) { if let Some(method) = handle_key(app, key) {
if let Some(path) = attachment_command_path(&method) { if let Some(path) = attachment_command_path(&method) {
match client.upload_path(&path).await { match client.upload_path(&path).await {
Ok(reference) => app.push_notice(format!( Ok(reference) => app.flash_actionbar_notice(
"Attached {} ({} bytes); it will be sent with the next message.", format!(
reference.file_name, reference.byte_len "Attached {} ({} bytes); it will be sent with the next message.",
)), reference.file_name, reference.byte_len
),
ActionbarNoticeLevel::Info,
ActionbarNoticeSource::Tui,
Duration::from_secs(6),
),
Err(error) => { Err(error) => {
app.push_error(format!("Attachment upload failed: {error}")); app.push_error(format!("Attachment upload failed: {error}"));
} }
} }
} else if is_clear_attachments_command(&method) { } else if is_clear_attachments_command(&method) {
client.clear_pending_attachments().await; client.clear_pending_attachments().await;
app.push_notice("Removed pending attachments."); app.flash_actionbar_notice(
"Removed pending attachments.",
ActionbarNoticeLevel::Info,
ActionbarNoticeSource::Tui,
Duration::from_secs(4),
);
} else { } else {
client.send(&method).await?; client.send(&method).await?;
} }
@@ -2145,6 +2145,25 @@ where
), ),
); );
} }
let cleanup_handle = match self.workers.lock() {
Ok(workers) => workers
.get(handle.worker_ref())
.map(|execution| execution.handle.clone()),
Err(_) => {
return WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
"worker adapter registry lock is poisoned",
);
}
};
if let Some(worker) = cleanup_handle
&& let Err(error) = worker.delete_uncommitted_uploaded_files()
{
return WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
format!("uploaded_file_cleanup_failed: {error}"),
);
}
let execution = match self.workers.lock() { let execution = match self.workers.lock() {
Ok(mut workers) => workers.remove(handle.worker_ref()), Ok(mut workers) => workers.remove(handle.worker_ref()),
Err(_) => { Err(_) => {
+5
View File
@@ -88,6 +88,11 @@ impl WorkerHandle {
.delete_uploaded_file(self.session_id, artifact_id) .delete_uploaded_file(self.session_id, artifact_id)
} }
pub fn delete_uncommitted_uploaded_files(&self) -> Result<u64, session_store::StoreError> {
self.artifact_store
.delete_uncommitted_uploaded_files(self.session_id)
}
pub fn subscribe(&self) -> broadcast::Receiver<Event> { pub fn subscribe(&self) -> broadcast::Receiver<Event> {
self.working_event_tx.subscribe() self.working_event_tx.subscribe()
} }
+1 -1
View File
@@ -340,7 +340,7 @@ mod tests {
owner, owner,
"image.png", "image.png",
"image/png", "image/png",
&[0xff, 0xd8, 0x00], b"\x89PNG\r\n\x1a\nbody",
session_store::UploadedFileLimits::default(), session_store::UploadedFileLimits::default(),
) )
.unwrap(); .unwrap();
+22 -6
View File
@@ -3331,7 +3331,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}), }),
compacted_from: None, compacted_from: None,
}; };
let mut initial_entries = vec![entry.clone()]; let mut initial_entries = vec![
entry.clone(),
LogEntry::InputSegmentsCheckpoint {
ts: segment_log::now_millis(),
user_segments: self.user_segments.clone(),
},
];
if let Some(checkpoint) = if let Some(checkpoint) =
active_run_checkpoint_entry(w.active_run_turn_count(), w.turn_count()) active_run_checkpoint_entry(w.active_run_turn_count(), w.turn_count())
{ {
@@ -4295,6 +4301,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let retained_user_segments = self
.user_segments
.iter()
.skip(self.user_segments.len().saturating_sub(retained_user_msgs))
.cloned()
.collect::<Vec<_>>();
// Build the SegmentStart entry for the new compacted segment. // Build the SegmentStart entry for the new compacted segment.
// Inherits the source Segment's session_id so the compacted // Inherits the source Segment's session_id so the compacted
// lineage stays grouped under the same Session. Atomically // lineage stays grouped under the same Session. Atomically
@@ -4320,7 +4333,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
at_turn_index: source_turn_count, at_turn_index: source_turn_count,
}), }),
}; };
let mut initial_entries = vec![entry.clone()]; let mut initial_entries = vec![
entry.clone(),
LogEntry::InputSegmentsCheckpoint {
ts: segment_log::now_millis(),
user_segments: retained_user_segments.clone(),
},
];
if let Some(checkpoint) = if let Some(checkpoint) =
active_run_checkpoint_entry(w.active_run_turn_count(), source_turn_count) active_run_checkpoint_entry(w.active_run_turn_count(), source_turn_count)
{ {
@@ -4372,10 +4391,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
// segments; only the user_messages surviving in retained_items // segments; only the user_messages surviving in retained_items
// keep them. They are always the trailing K entries of // keep them. They are always the trailing K entries of
// `self.user_segments` because submissions are appended in order. // `self.user_segments` because submissions are appended in order.
let drop_n = self.user_segments.len().saturating_sub(retained_user_msgs); self.user_segments = retained_user_segments;
if drop_n > 0 {
self.user_segments.drain(..drop_n);
}
self.session.replace_history(compacted_history_entries); self.session.replace_history(compacted_history_entries);
// Compaction-introduced system messages are part of the new // Compaction-introduced system messages are part of the new