fix: harden uploaded attachment retention and replay
This commit is contained in:
Generated
+1
@@ -4412,6 +4412,7 @@ dependencies = [
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"unicode-normalization",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ uuid = { workspace = true, features = ["v7", "serde"] }
|
||||
thiserror = { workspace = true }
|
||||
protocol = { workspace = true }
|
||||
tracing.workspace = true
|
||||
unicode-normalization = "0.1.25"
|
||||
|
||||
[dev-dependencies]
|
||||
async-trait = { workspace = true }
|
||||
|
||||
@@ -20,8 +20,8 @@ use crate::paste_artifact::{read_from_dir, write_to_dir};
|
||||
use crate::segment_log::LogEntry;
|
||||
use crate::store::{Store, StoreError};
|
||||
use crate::uploaded_file::{
|
||||
bind_uploaded_file, delete_uploaded_file, read_uploaded_file, read_uploaded_file_by_id,
|
||||
write_uploaded_file,
|
||||
bind_uploaded_file, delete_uncommitted_uploaded_files, delete_uploaded_file,
|
||||
read_uploaded_file, read_uploaded_file_by_id, write_uploaded_file,
|
||||
};
|
||||
use crate::{PasteArtifactLimits, SegmentId, SessionId, UploadedFileLimits};
|
||||
use protocol::{PasteArtifactRef, UploadedFileRef};
|
||||
@@ -459,6 +459,14 @@ impl Store for FsStore {
|
||||
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(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
@@ -678,6 +686,47 @@ mod tests {
|
||||
store.write_uploaded_file(session_id, "notes.txt", "not a type", b"x", limits),
|
||||
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
|
||||
.write_paste_artifact(
|
||||
session_id,
|
||||
|
||||
@@ -69,7 +69,7 @@ pub use system_item::{
|
||||
};
|
||||
pub use uploaded_file::{
|
||||
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::{
|
||||
CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata,
|
||||
|
||||
@@ -41,6 +41,24 @@ pub fn project_session_snapshot(session_id: SessionId, log: &[LogEntry]) -> Sess
|
||||
entries.clear();
|
||||
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 {
|
||||
ts,
|
||||
segments,
|
||||
@@ -357,6 +375,63 @@ mod tests {
|
||||
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]
|
||||
fn annotated_user_input_attaches_segments_to_first_user_role_entry_for_any_origin() {
|
||||
let session_id = crate::new_session_id();
|
||||
|
||||
@@ -63,6 +63,14 @@ pub enum LogEntry {
|
||||
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
|
||||
/// cycle (Invoke range). The range extends implicitly until 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))
|
||||
.collect();
|
||||
}
|
||||
LogEntry::InputSegmentsCheckpoint { user_segments, .. } => {
|
||||
state.user_segments = user_segments.clone();
|
||||
}
|
||||
LogEntry::Invoke { .. } => {
|
||||
// A terminal run record below clears or refines this. If the
|
||||
// log ends first, restore must treat the turn as interrupted.
|
||||
|
||||
@@ -221,6 +221,10 @@ pub trait Store: Send + Sync {
|
||||
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.
|
||||
fn append_trace(
|
||||
&self,
|
||||
|
||||
@@ -9,6 +9,7 @@ use fs4::fs_std::FileExt;
|
||||
use protocol::{UploadedFileAvailability, UploadedFileRef};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use unicode_normalization::UnicodeNormalization;
|
||||
use uuid::Uuid;
|
||||
|
||||
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_SESSION_ARTIFACT_BYTES: u64 = 32 * 1024 * 1024;
|
||||
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_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 == "."
|
||||
|| file_name == ".."
|
||||
|| file_name
|
||||
.chars()
|
||||
.any(|ch| ch.is_control() || matches!(ch, '/' | '\\'))
|
||||
|| file_name.chars().any(|ch| {
|
||||
ch.is_control()
|
||||
|| 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);
|
||||
}
|
||||
@@ -96,6 +104,35 @@ pub(crate) fn validate_media_type(media_type: &str) -> Result<()> {
|
||||
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> {
|
||||
let id = Uuid::parse_str(artifact_id).map_err(|_| StoreError::InvalidArtifactId)?;
|
||||
Ok(dir.join(format!("{id}.file.json")))
|
||||
@@ -133,7 +170,7 @@ pub(crate) fn stored_uploaded_file_usage(dir: &Path) -> Result<(u64, u64)> {
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(path)?)?;
|
||||
let stored: StoredUploadedFile = serde_json::from_slice(&fs::read(&path)?)?;
|
||||
bytes = bytes
|
||||
.checked_add(stored.byte_len)
|
||||
.ok_or(StoreError::ArtifactQuotaExceeded)?;
|
||||
@@ -153,7 +190,9 @@ pub(crate) fn write_uploaded_file(
|
||||
) -> Result<UploadedFileRef> {
|
||||
validate_file_name(file_name)?;
|
||||
validate_media_type(media_type)?;
|
||||
validate_content(media_type, content)?;
|
||||
let byte_len = u64::try_from(content.len()).map_err(|_| StoreError::ArtifactTooLarge)?;
|
||||
let sha256 = digest(content);
|
||||
if byte_len > limits.max_file_bytes {
|
||||
return Err(StoreError::ArtifactTooLarge);
|
||||
}
|
||||
@@ -166,7 +205,48 @@ pub(crate) fn write_uploaded_file(
|
||||
.open(dir.join(".aggregate.lock"))?;
|
||||
FileExt::lock_exclusive(&aggregate_lock)?;
|
||||
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
|
||||
.checked_add(file_bytes)
|
||||
.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 created_at_ms = now_ms()?;
|
||||
let sha256 = digest(content);
|
||||
let stored = StoredUploadedFile {
|
||||
file_name: file_name.to_owned(),
|
||||
media_type: media_type.to_owned(),
|
||||
@@ -274,6 +353,36 @@ pub(crate) fn bind_uploaded_file(
|
||||
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> {
|
||||
fs::create_dir_all(dir)?;
|
||||
let aggregate_lock = fs::OpenOptions::new()
|
||||
|
||||
@@ -928,14 +928,6 @@ impl App {
|
||||
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>) {
|
||||
self.blocks.push(Block::Alert {
|
||||
level: AlertLevel::Error,
|
||||
|
||||
@@ -792,17 +792,27 @@ async fn handle_terminal_event<T: Socket>(
|
||||
if let Some(method) = handle_key(app, key) {
|
||||
if let Some(path) = attachment_command_path(&method) {
|
||||
match client.upload_path(&path).await {
|
||||
Ok(reference) => app.push_notice(format!(
|
||||
Ok(reference) => app.flash_actionbar_notice(
|
||||
format!(
|
||||
"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) => {
|
||||
app.push_error(format!("Attachment upload failed: {error}"));
|
||||
}
|
||||
}
|
||||
} else if is_clear_attachments_command(&method) {
|
||||
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 {
|
||||
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() {
|
||||
Ok(mut workers) => workers.remove(handle.worker_ref()),
|
||||
Err(_) => {
|
||||
|
||||
@@ -88,6 +88,11 @@ impl WorkerHandle {
|
||||
.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> {
|
||||
self.working_event_tx.subscribe()
|
||||
}
|
||||
|
||||
@@ -340,7 +340,7 @@ mod tests {
|
||||
owner,
|
||||
"image.png",
|
||||
"image/png",
|
||||
&[0xff, 0xd8, 0x00],
|
||||
b"\x89PNG\r\n\x1a\nbody",
|
||||
session_store::UploadedFileLimits::default(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
@@ -3331,7 +3331,13 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
||||
}),
|
||||
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) =
|
||||
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<_>>();
|
||||
|
||||
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.
|
||||
// Inherits the source Segment's session_id so the compacted
|
||||
// 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,
|
||||
}),
|
||||
};
|
||||
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) =
|
||||
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
|
||||
// keep them. They are always the trailing K entries of
|
||||
// `self.user_segments` because submissions are appended in order.
|
||||
let drop_n = self.user_segments.len().saturating_sub(retained_user_msgs);
|
||||
if drop_n > 0 {
|
||||
self.user_segments.drain(..drop_n);
|
||||
}
|
||||
self.user_segments = retained_user_segments;
|
||||
|
||||
self.session.replace_history(compacted_history_entries);
|
||||
// Compaction-introduced system messages are part of the new
|
||||
|
||||
Reference in New Issue
Block a user