fix: harden uploaded attachment retention and replay
This commit is contained in:
Generated
+1
@@ -4412,6 +4412,7 @@ dependencies = [
|
|||||||
"thiserror 2.0.18",
|
"thiserror 2.0.18",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tracing",
|
"tracing",
|
||||||
|
"unicode-normalization",
|
||||||
"uuid",
|
"uuid",
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|||||||
@@ -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 }
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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();
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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(_) => {
|
||||||
|
|||||||
@@ -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()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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();
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user