feat: store large paste inputs as artifacts
This commit is contained in:
Generated
+1
@@ -4406,6 +4406,7 @@ dependencies = [
|
|||||||
"protocol",
|
"protocol",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
|
"sha2 0.11.0",
|
||||||
"tempfile",
|
"tempfile",
|
||||||
"thiserror 2.0.18",
|
"thiserror 2.0.18",
|
||||||
"tokio",
|
"tokio",
|
||||||
|
|||||||
@@ -193,6 +193,23 @@ impl WorkerEvent {
|
|||||||
/// variants — emits an alert and inserts a `[unknown input segment]`
|
/// variants — emits an alert and inserts a `[unknown input segment]`
|
||||||
/// placeholder into the LLM context so neither user nor LLM is blind to
|
/// placeholder into the LLM context so neither user nor LLM is blind to
|
||||||
/// the dropped intent.
|
/// the dropped intent.
|
||||||
|
/// Session-owned reference to a large pasted-input artifact.
|
||||||
|
///
|
||||||
|
/// The reference contains only bounded integrity and provenance metadata. The
|
||||||
|
/// artifact body remains in session storage and is available to the model only
|
||||||
|
/// through the scoped paste-artifact tools installed by Worker.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||||
|
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
|
||||||
|
pub struct PasteArtifactRef {
|
||||||
|
pub artifact_id: String,
|
||||||
|
pub byte_len: u64,
|
||||||
|
pub char_count: u64,
|
||||||
|
pub line_count: u64,
|
||||||
|
pub sha256: String,
|
||||||
|
pub source_entry_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
|
||||||
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
|
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
|
||||||
@@ -210,6 +227,10 @@ pub enum Segment {
|
|||||||
lines: u32,
|
lines: u32,
|
||||||
content: String,
|
content: String,
|
||||||
},
|
},
|
||||||
|
/// Internal reference produced when Worker stores a large `Paste` before
|
||||||
|
/// committing input. Clients may receive this in history/event projections;
|
||||||
|
/// the body is intentionally absent.
|
||||||
|
PasteArtifact { artifact: PasteArtifactRef },
|
||||||
/// `@<path>` file-system reference. Worker resolves readable files to
|
/// `@<path>` file-system reference. Worker resolves readable files to
|
||||||
/// `[File: <path>]` attachments and readable normal directories to shallow
|
/// `[File: <path>]` attachments and readable normal directories to shallow
|
||||||
/// `[Dir: <path>]` listings; the flattened user text keeps the literal
|
/// `[Dir: <path>]` listings; the flattened user text keeps the literal
|
||||||
@@ -250,6 +271,18 @@ impl Segment {
|
|||||||
match seg {
|
match seg {
|
||||||
Segment::Text { content } => out.push_str(content),
|
Segment::Text { content } => out.push_str(content),
|
||||||
Segment::Paste { content, .. } => out.push_str(content),
|
Segment::Paste { content, .. } => out.push_str(content),
|
||||||
|
Segment::PasteArtifact { artifact } => {
|
||||||
|
use std::fmt::Write as _;
|
||||||
|
let _ = write!(
|
||||||
|
out,
|
||||||
|
"[Large paste stored as artifact {}: {} bytes, {} chars, {} lines, sha256 {}; use SearchInputArtifact and ReadInputArtifact to inspect it]",
|
||||||
|
artifact.artifact_id,
|
||||||
|
artifact.byte_len,
|
||||||
|
artifact.char_count,
|
||||||
|
artifact.line_count,
|
||||||
|
artifact.sha256
|
||||||
|
);
|
||||||
|
}
|
||||||
Segment::FileRef { path } => {
|
Segment::FileRef { path } => {
|
||||||
out.push('@');
|
out.push('@');
|
||||||
out.push_str(path);
|
out.push_str(path);
|
||||||
@@ -1202,6 +1235,29 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn paste_artifact_segment_roundtrips_without_body() {
|
||||||
|
let artifact = PasteArtifactRef {
|
||||||
|
artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b2".to_string(),
|
||||||
|
byte_len: 65_536,
|
||||||
|
char_count: 65_530,
|
||||||
|
line_count: 200,
|
||||||
|
sha256: "a".repeat(64),
|
||||||
|
source_entry_id: "entry-1".to_string(),
|
||||||
|
};
|
||||||
|
let segment = Segment::PasteArtifact {
|
||||||
|
artifact: artifact.clone(),
|
||||||
|
};
|
||||||
|
let json = serde_json::to_string(&segment).unwrap();
|
||||||
|
assert!(!json.contains("pasted body"));
|
||||||
|
assert_eq!(serde_json::from_str::<Segment>(&json).unwrap(), segment);
|
||||||
|
let projected = Segment::flatten_to_text(&[segment]);
|
||||||
|
assert!(projected.contains(&artifact.artifact_id));
|
||||||
|
assert!(projected.contains("SearchInputArtifact"));
|
||||||
|
assert!(projected.contains("ReadInputArtifact"));
|
||||||
|
assert!(!projected.contains("pasted body"));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn method_run_flow_segment_roundtrip() {
|
fn method_run_flow_segment_roundtrip() {
|
||||||
let method = Method::Run {
|
let method = Method::Run {
|
||||||
|
|||||||
@@ -7,10 +7,11 @@ use crate::{
|
|||||||
CommandStreamSlice, CompactionLifecycle, CompactionLifecycleState, CompletionEntry,
|
CommandStreamSlice, CompactionLifecycle, CompactionLifecycleState, CompletionEntry,
|
||||||
CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot,
|
CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot,
|
||||||
InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot,
|
InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot,
|
||||||
InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary, RewindTarget, RewindTargetId,
|
InvokeKind, MemoryWorkerEvent, Method, PasteArtifactRef, Permission, RewindSummary,
|
||||||
RunResult, ScopeRule, Segment, SessionContentPart, SessionEntryProvenance, SessionMessageRole,
|
RewindTarget, RewindTargetId, RunResult, ScopeRule, Segment, SessionContentPart,
|
||||||
SessionSnapshot, SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment,
|
SessionEntryProvenance, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry,
|
||||||
ToolResultDisposition, TurnResult, WorkerEvent, WorkerStatus,
|
SessionSnapshotEntryData, SessionToolAttachment, ToolResultDisposition, TurnResult,
|
||||||
|
WorkerEvent, WorkerStatus,
|
||||||
subscription::{
|
subscription::{
|
||||||
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
|
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
|
||||||
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
|
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
|
||||||
@@ -78,6 +79,7 @@ pub fn generated_protocol_types() -> String {
|
|||||||
push_decl::<Greeting>(&cfg, &mut output);
|
push_decl::<Greeting>(&cfg, &mut output);
|
||||||
push_decl::<Alert>(&cfg, &mut output);
|
push_decl::<Alert>(&cfg, &mut output);
|
||||||
push_decl::<MemoryWorkerEvent>(&cfg, &mut output);
|
push_decl::<MemoryWorkerEvent>(&cfg, &mut output);
|
||||||
|
push_decl::<PasteArtifactRef>(&cfg, &mut output);
|
||||||
push_decl::<Segment>(&cfg, &mut output);
|
push_decl::<Segment>(&cfg, &mut output);
|
||||||
push_decl::<WorkerEvent>(&cfg, &mut output);
|
push_decl::<WorkerEvent>(&cfg, &mut output);
|
||||||
push_decl::<SubscriptionRequestId>(&cfg, &mut output);
|
push_decl::<SubscriptionRequestId>(&cfg, &mut output);
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ base64.workspace = true
|
|||||||
agen = { workspace = true }
|
agen = { workspace = true }
|
||||||
serde = { workspace = true, features = ["derive"] }
|
serde = { workspace = true, features = ["derive"] }
|
||||||
serde_json = { workspace = true }
|
serde_json = { workspace = true }
|
||||||
|
sha2.workspace = true
|
||||||
uuid = { workspace = true, features = ["v7", "serde"] }
|
uuid = { workspace = true, features = ["v7", "serde"] }
|
||||||
thiserror = { workspace = true }
|
thiserror = { workspace = true }
|
||||||
protocol = { workspace = true }
|
protocol = { workspace = true }
|
||||||
|
|||||||
@@ -16,9 +16,11 @@
|
|||||||
//! enumerable by the picker.
|
//! enumerable by the picker.
|
||||||
|
|
||||||
use crate::event_trace::TraceEntry;
|
use crate::event_trace::TraceEntry;
|
||||||
|
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::{SegmentId, SessionId};
|
use crate::{PasteArtifactLimits, SegmentId, SessionId};
|
||||||
|
use protocol::PasteArtifactRef;
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::io::{Read, Seek, SeekFrom, Write};
|
use std::io::{Read, Seek, SeekFrom, Write};
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
@@ -109,6 +111,16 @@ impl FsStore {
|
|||||||
.join(format!("{segment_id}.trace.jsonl"))
|
.join(format!("{segment_id}.trace.jsonl"))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn paste_artifact_dir(&self, session_id: SessionId) -> PathBuf {
|
||||||
|
self.session_dir(session_id).join("artifacts").join("paste")
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
fn paste_artifact_path(&self, session_id: SessionId, artifact_id: &str) -> PathBuf {
|
||||||
|
self.paste_artifact_dir(session_id)
|
||||||
|
.join(format!("{artifact_id}.json"))
|
||||||
|
}
|
||||||
|
|
||||||
fn append_line(&self, path: &Path, line: &str) -> Result<(), StoreError> {
|
fn append_line(&self, path: &Path, line: &str) -> Result<(), StoreError> {
|
||||||
let _guard = self
|
let _guard = self
|
||||||
.append_lock
|
.append_lock
|
||||||
@@ -350,6 +362,33 @@ impl Store for FsStore {
|
|||||||
Ok(complete.lines().filter(|l| !l.trim().is_empty()).count())
|
Ok(complete.lines().filter(|l| !l.trim().is_empty()).count())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn write_paste_artifact(
|
||||||
|
&self,
|
||||||
|
session_id: SessionId,
|
||||||
|
source_entry_id: &str,
|
||||||
|
content: &str,
|
||||||
|
limits: PasteArtifactLimits,
|
||||||
|
) -> Result<PasteArtifactRef, StoreError> {
|
||||||
|
let _guard = self
|
||||||
|
.append_lock
|
||||||
|
.lock()
|
||||||
|
.map_err(|_| std::io::Error::other("session store append lock was poisoned"))?;
|
||||||
|
write_to_dir(
|
||||||
|
&self.paste_artifact_dir(session_id),
|
||||||
|
source_entry_id,
|
||||||
|
content,
|
||||||
|
limits,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_paste_artifact(
|
||||||
|
&self,
|
||||||
|
session_id: SessionId,
|
||||||
|
artifact_id: &str,
|
||||||
|
) -> Result<(PasteArtifactRef, String), StoreError> {
|
||||||
|
read_from_dir(&self.paste_artifact_dir(session_id), artifact_id)
|
||||||
|
}
|
||||||
|
|
||||||
fn append_trace(
|
fn append_trace(
|
||||||
&self,
|
&self,
|
||||||
session_id: SessionId,
|
session_id: SessionId,
|
||||||
@@ -398,4 +437,87 @@ mod tests {
|
|||||||
store.create_segment(session_id, segment_id, &[]).unwrap();
|
store.create_segment(session_id, segment_id, &[]).unwrap();
|
||||||
assert!(store.session_modified_at(session_id).unwrap().is_some());
|
assert!(store.session_modified_at(session_id).unwrap().is_some());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn paste_artifacts_are_atomic_integrity_checked_and_session_scoped() {
|
||||||
|
let tmp = tempfile::TempDir::new().unwrap();
|
||||||
|
let store = FsStore::new(tmp.path()).unwrap();
|
||||||
|
let owner = new_session_id();
|
||||||
|
let other = new_session_id();
|
||||||
|
let content = "αβγ\nsecond line\n";
|
||||||
|
let reference = store
|
||||||
|
.write_paste_artifact(owner, "entry-1", content, PasteArtifactLimits::default())
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert_eq!(reference.byte_len, content.len() as u64);
|
||||||
|
assert_eq!(reference.char_count, content.chars().count() as u64);
|
||||||
|
assert_eq!(reference.source_entry_id, "entry-1");
|
||||||
|
assert_eq!(
|
||||||
|
store
|
||||||
|
.read_paste_artifact(owner, &reference.artifact_id)
|
||||||
|
.unwrap()
|
||||||
|
.1,
|
||||||
|
content
|
||||||
|
);
|
||||||
|
assert!(matches!(
|
||||||
|
store.read_paste_artifact(other, &reference.artifact_id),
|
||||||
|
Err(StoreError::PasteArtifactNotFound(_))
|
||||||
|
));
|
||||||
|
assert!(
|
||||||
|
self::fs::read_dir(store.paste_artifact_dir(owner))
|
||||||
|
.unwrap()
|
||||||
|
.all(|entry| !entry
|
||||||
|
.unwrap()
|
||||||
|
.file_name()
|
||||||
|
.to_string_lossy()
|
||||||
|
.ends_with(".tmp"))
|
||||||
|
);
|
||||||
|
let very_large = "z".repeat(1024 * 1024);
|
||||||
|
let very_large_ref = store
|
||||||
|
.write_paste_artifact(
|
||||||
|
owner,
|
||||||
|
"entry-2",
|
||||||
|
&very_large,
|
||||||
|
PasteArtifactLimits::default(),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
store
|
||||||
|
.read_paste_artifact(owner, &very_large_ref.artifact_id)
|
||||||
|
.unwrap()
|
||||||
|
.1,
|
||||||
|
very_large
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn paste_artifact_limits_and_corruption_fail_closed() {
|
||||||
|
let tmp = tempfile::TempDir::new().unwrap();
|
||||||
|
let store = FsStore::new(tmp.path()).unwrap();
|
||||||
|
let session_id = new_session_id();
|
||||||
|
let limits = PasteArtifactLimits {
|
||||||
|
max_artifact_bytes: 5,
|
||||||
|
max_session_bytes: 8,
|
||||||
|
};
|
||||||
|
let first = store
|
||||||
|
.write_paste_artifact(session_id, "entry-1", "1234", limits)
|
||||||
|
.unwrap();
|
||||||
|
assert!(matches!(
|
||||||
|
store.write_paste_artifact(session_id, "entry-2", "56789", limits),
|
||||||
|
Err(StoreError::PasteArtifactLimit(_))
|
||||||
|
));
|
||||||
|
assert!(matches!(
|
||||||
|
store.write_paste_artifact(session_id, "entry-2", "5678", limits),
|
||||||
|
Ok(_)
|
||||||
|
));
|
||||||
|
std::fs::write(
|
||||||
|
store.paste_artifact_path(session_id, &first.artifact_id),
|
||||||
|
b"{}",
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
assert!(matches!(
|
||||||
|
store.read_paste_artifact(session_id, &first.artifact_id),
|
||||||
|
Err(StoreError::Serde(_)) | Err(StoreError::PasteArtifactIntegrity(_))
|
||||||
|
));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -35,6 +35,7 @@ pub mod fs_store;
|
|||||||
pub mod history;
|
pub mod history;
|
||||||
mod legacy_session_log;
|
mod legacy_session_log;
|
||||||
pub mod logged_item;
|
pub mod logged_item;
|
||||||
|
mod paste_artifact;
|
||||||
pub mod public_snapshot;
|
pub mod public_snapshot;
|
||||||
pub mod segment;
|
pub mod segment;
|
||||||
pub mod segment_log;
|
pub mod segment_log;
|
||||||
@@ -53,6 +54,7 @@ pub use history::{
|
|||||||
LoggedWorkerSubject,
|
LoggedWorkerSubject,
|
||||||
};
|
};
|
||||||
pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged};
|
pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged};
|
||||||
|
pub use paste_artifact::PasteArtifactLimits;
|
||||||
pub use segment::{
|
pub use segment::{
|
||||||
SegmentStartState, append_entry, append_system_item, classify_logged_history_entry,
|
SegmentStartState, append_entry, append_system_item, classify_logged_history_entry,
|
||||||
create_compacted_segment, create_segment, create_segment_with_ids, ensure_head_or_fork, fork,
|
create_compacted_segment, create_segment, create_segment_with_ids, ensure_head_or_fork, fork,
|
||||||
|
|||||||
@@ -0,0 +1,156 @@
|
|||||||
|
//! Session-owned storage for large pasted-input artifacts.
|
||||||
|
|
||||||
|
use std::fs;
|
||||||
|
use std::io::Write as _;
|
||||||
|
use std::path::Path;
|
||||||
|
|
||||||
|
use protocol::PasteArtifactRef;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
|
||||||
|
use crate::StoreError;
|
||||||
|
|
||||||
|
/// Bounded storage policy applied before a large paste becomes durable input.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub struct PasteArtifactLimits {
|
||||||
|
pub max_artifact_bytes: u64,
|
||||||
|
pub max_session_bytes: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for PasteArtifactLimits {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
max_artifact_bytes: 8 * 1024 * 1024,
|
||||||
|
max_session_bytes: 64 * 1024 * 1024,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Integrity-bearing on-disk record. The body and metadata are committed in one
|
||||||
|
/// atomic file replacement so readers never observe a half-written artifact.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub(crate) struct StoredPasteArtifact {
|
||||||
|
pub reference: PasteArtifactRef,
|
||||||
|
pub content: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn write_to_dir(
|
||||||
|
artifact_dir: &Path,
|
||||||
|
source_entry_id: &str,
|
||||||
|
content: &str,
|
||||||
|
limits: PasteArtifactLimits,
|
||||||
|
) -> Result<PasteArtifactRef, StoreError> {
|
||||||
|
let byte_len = content.len() as u64;
|
||||||
|
if byte_len > limits.max_artifact_bytes {
|
||||||
|
return Err(StoreError::PasteArtifactLimit(format!(
|
||||||
|
"artifact has {byte_len} bytes; maximum is {}",
|
||||||
|
limits.max_artifact_bytes
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
fs::create_dir_all(artifact_dir)?;
|
||||||
|
let mut aggregate = 0_u64;
|
||||||
|
for entry in fs::read_dir(artifact_dir)? {
|
||||||
|
let path = entry?.path();
|
||||||
|
if path.extension().and_then(|value| value.to_str()) != Some("json") {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let stored: StoredPasteArtifact = serde_json::from_slice(&fs::read(&path)?)?;
|
||||||
|
verify(&stored, &stored.reference.artifact_id)?;
|
||||||
|
aggregate = aggregate
|
||||||
|
.checked_add(stored.reference.byte_len)
|
||||||
|
.ok_or_else(|| {
|
||||||
|
StoreError::PasteArtifactLimit("session aggregate size overflow".to_string())
|
||||||
|
})?;
|
||||||
|
}
|
||||||
|
let projected = aggregate.checked_add(byte_len).ok_or_else(|| {
|
||||||
|
StoreError::PasteArtifactLimit("session aggregate size overflow".to_string())
|
||||||
|
})?;
|
||||||
|
if projected > limits.max_session_bytes {
|
||||||
|
return Err(StoreError::PasteArtifactLimit(format!(
|
||||||
|
"session artifacts would use {projected} bytes; maximum is {}",
|
||||||
|
limits.max_session_bytes
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let artifact_id = uuid::Uuid::now_v7().to_string();
|
||||||
|
let reference = PasteArtifactRef {
|
||||||
|
artifact_id: artifact_id.clone(),
|
||||||
|
byte_len,
|
||||||
|
char_count: content.chars().count() as u64,
|
||||||
|
line_count: line_count(content),
|
||||||
|
sha256: sha256_hex(content),
|
||||||
|
source_entry_id: source_entry_id.to_string(),
|
||||||
|
};
|
||||||
|
let bytes = serde_json::to_vec(&StoredPasteArtifact {
|
||||||
|
reference: reference.clone(),
|
||||||
|
content: content.to_string(),
|
||||||
|
})?;
|
||||||
|
let target = artifact_dir.join(format!("{artifact_id}.json"));
|
||||||
|
let temporary = artifact_dir.join(format!(".{artifact_id}.tmp"));
|
||||||
|
let mut file = fs::OpenOptions::new()
|
||||||
|
.create_new(true)
|
||||||
|
.write(true)
|
||||||
|
.open(&temporary)?;
|
||||||
|
if let Err(error) = file.write_all(&bytes).and_then(|_| file.sync_all()) {
|
||||||
|
let _ = fs::remove_file(&temporary);
|
||||||
|
return Err(error.into());
|
||||||
|
}
|
||||||
|
if let Err(error) = fs::rename(&temporary, &target) {
|
||||||
|
let _ = fs::remove_file(&temporary);
|
||||||
|
return Err(error.into());
|
||||||
|
}
|
||||||
|
if let Ok(directory) = fs::File::open(artifact_dir) {
|
||||||
|
directory.sync_all()?;
|
||||||
|
}
|
||||||
|
Ok(reference)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn read_from_dir(
|
||||||
|
artifact_dir: &Path,
|
||||||
|
artifact_id: &str,
|
||||||
|
) -> Result<(PasteArtifactRef, String), StoreError> {
|
||||||
|
let parsed = uuid::Uuid::parse_str(artifact_id)
|
||||||
|
.map_err(|_| StoreError::PasteArtifactNotFound(artifact_id.to_string()))?;
|
||||||
|
if parsed.to_string() != artifact_id {
|
||||||
|
return Err(StoreError::PasteArtifactNotFound(artifact_id.to_string()));
|
||||||
|
}
|
||||||
|
let path = artifact_dir.join(format!("{artifact_id}.json"));
|
||||||
|
let bytes = match fs::read(path) {
|
||||||
|
Ok(bytes) => bytes,
|
||||||
|
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
|
||||||
|
return Err(StoreError::PasteArtifactNotFound(artifact_id.to_string()));
|
||||||
|
}
|
||||||
|
Err(error) => return Err(error.into()),
|
||||||
|
};
|
||||||
|
let stored: StoredPasteArtifact = serde_json::from_slice(&bytes)?;
|
||||||
|
verify(&stored, artifact_id)?;
|
||||||
|
Ok((stored.reference, stored.content))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn verify(stored: &StoredPasteArtifact, artifact_id: &str) -> Result<(), StoreError> {
|
||||||
|
let actual_digest = sha256_hex(&stored.content);
|
||||||
|
if stored.reference.artifact_id != artifact_id
|
||||||
|
|| stored.reference.byte_len != stored.content.len() as u64
|
||||||
|
|| stored.reference.char_count != stored.content.chars().count() as u64
|
||||||
|
|| stored.reference.line_count != line_count(&stored.content)
|
||||||
|
|| stored.reference.sha256 != actual_digest
|
||||||
|
{
|
||||||
|
return Err(StoreError::PasteArtifactIntegrity(artifact_id.to_string()));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sha256_hex(content: &str) -> String {
|
||||||
|
Sha256::digest(content.as_bytes())
|
||||||
|
.iter()
|
||||||
|
.map(|byte| format!("{byte:02x}"))
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn line_count(content: &str) -> u64 {
|
||||||
|
if content.is_empty() {
|
||||||
|
0
|
||||||
|
} else {
|
||||||
|
content.lines().count().max(1) as u64
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -13,7 +13,8 @@
|
|||||||
|
|
||||||
use crate::event_trace::TraceEntry;
|
use crate::event_trace::TraceEntry;
|
||||||
use crate::segment_log::LogEntry;
|
use crate::segment_log::LogEntry;
|
||||||
use crate::{SegmentId, SessionId};
|
use crate::{PasteArtifactLimits, SegmentId, SessionId};
|
||||||
|
use protocol::PasteArtifactRef;
|
||||||
|
|
||||||
/// Errors from the persistence store.
|
/// Errors from the persistence store.
|
||||||
#[derive(Debug, thiserror::Error)]
|
#[derive(Debug, thiserror::Error)]
|
||||||
@@ -29,6 +30,18 @@ pub enum StoreError {
|
|||||||
|
|
||||||
#[error("log corrupted at line {line}: {message}")]
|
#[error("log corrupted at line {line}: {message}")]
|
||||||
Corrupt { line: usize, message: String },
|
Corrupt { line: usize, message: String },
|
||||||
|
|
||||||
|
#[error("paste artifact storage is unavailable")]
|
||||||
|
PasteArtifactUnsupported,
|
||||||
|
|
||||||
|
#[error("paste artifact not found: {0}")]
|
||||||
|
PasteArtifactNotFound(String),
|
||||||
|
|
||||||
|
#[error("paste artifact integrity check failed: {0}")]
|
||||||
|
PasteArtifactIntegrity(String),
|
||||||
|
|
||||||
|
#[error("paste artifact size limit exceeded: {0}")]
|
||||||
|
PasteArtifactLimit(String),
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Sync persistence backend for segment logs.
|
/// Sync persistence backend for segment logs.
|
||||||
@@ -117,6 +130,26 @@ pub trait Store: Send + Sync {
|
|||||||
segment_id: SegmentId,
|
segment_id: SegmentId,
|
||||||
) -> Result<usize, StoreError>;
|
) -> Result<usize, StoreError>;
|
||||||
|
|
||||||
|
/// Store a large paste before its reference is committed to history.
|
||||||
|
fn write_paste_artifact(
|
||||||
|
&self,
|
||||||
|
_session_id: SessionId,
|
||||||
|
_source_entry_id: &str,
|
||||||
|
_content: &str,
|
||||||
|
_limits: PasteArtifactLimits,
|
||||||
|
) -> Result<PasteArtifactRef, StoreError> {
|
||||||
|
Err(StoreError::PasteArtifactUnsupported)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Read and verify one artifact owned by `session_id`.
|
||||||
|
fn read_paste_artifact(
|
||||||
|
&self,
|
||||||
|
_session_id: SessionId,
|
||||||
|
_artifact_id: &str,
|
||||||
|
) -> Result<(PasteArtifactRef, String), StoreError> {
|
||||||
|
Err(StoreError::PasteArtifactUnsupported)
|
||||||
|
}
|
||||||
|
|
||||||
/// 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,
|
||||||
|
|||||||
@@ -608,6 +608,24 @@ where
|
|||||||
) -> Result<usize, crate::StoreError> {
|
) -> Result<usize, crate::StoreError> {
|
||||||
self.session_store.read_entry_count(session_id, segment_id)
|
self.session_store.read_entry_count(session_id, segment_id)
|
||||||
}
|
}
|
||||||
|
fn write_paste_artifact(
|
||||||
|
&self,
|
||||||
|
session_id: SessionId,
|
||||||
|
source_entry_id: &str,
|
||||||
|
content: &str,
|
||||||
|
limits: crate::PasteArtifactLimits,
|
||||||
|
) -> Result<protocol::PasteArtifactRef, crate::StoreError> {
|
||||||
|
self.session_store
|
||||||
|
.write_paste_artifact(session_id, source_entry_id, content, limits)
|
||||||
|
}
|
||||||
|
fn read_paste_artifact(
|
||||||
|
&self,
|
||||||
|
session_id: SessionId,
|
||||||
|
artifact_id: &str,
|
||||||
|
) -> Result<(protocol::PasteArtifactRef, String), crate::StoreError> {
|
||||||
|
self.session_store
|
||||||
|
.read_paste_artifact(session_id, artifact_id)
|
||||||
|
}
|
||||||
fn append_trace(
|
fn append_trace(
|
||||||
&self,
|
&self,
|
||||||
session_id: SessionId,
|
session_id: SessionId,
|
||||||
|
|||||||
@@ -10,9 +10,11 @@
|
|||||||
//! every later operation must use that same ID.
|
//! every later operation must use that same ID.
|
||||||
|
|
||||||
use crate::event_trace::TraceEntry;
|
use crate::event_trace::TraceEntry;
|
||||||
|
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::{SegmentId, SessionId};
|
use crate::{PasteArtifactLimits, SegmentId, SessionId};
|
||||||
|
use protocol::PasteArtifactRef;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
use std::fs::{self, File, OpenOptions};
|
use std::fs::{self, File, OpenOptions};
|
||||||
use std::io::{Read, Seek, SeekFrom, Write};
|
use std::io::{Read, Seek, SeekFrom, Write};
|
||||||
@@ -25,6 +27,7 @@ const PREVIOUS_SESSION_SCHEMA_VERSION: u32 = 2;
|
|||||||
const LEGACY_SESSION_SCHEMA_VERSION: u32 = 1;
|
const LEGACY_SESSION_SCHEMA_VERSION: u32 = 1;
|
||||||
const SESSION_FILE: &str = "session.json";
|
const SESSION_FILE: &str = "session.json";
|
||||||
const SEGMENTS_DIR: &str = "segments";
|
const SEGMENTS_DIR: &str = "segments";
|
||||||
|
const PASTE_ARTIFACTS_DIR: &str = "artifacts/paste";
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[derive(Clone)]
|
||||||
pub struct WorkerSessionStore {
|
pub struct WorkerSessionStore {
|
||||||
@@ -317,6 +320,35 @@ impl Store for WorkerSessionStore {
|
|||||||
.count())
|
.count())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn write_paste_artifact(
|
||||||
|
&self,
|
||||||
|
session_id: SessionId,
|
||||||
|
source_entry_id: &str,
|
||||||
|
content: &str,
|
||||||
|
limits: PasteArtifactLimits,
|
||||||
|
) -> Result<PasteArtifactRef, StoreError> {
|
||||||
|
self.ensure_session(session_id, true)?;
|
||||||
|
let _guard = self
|
||||||
|
.append_lock
|
||||||
|
.lock()
|
||||||
|
.map_err(|_| std::io::Error::other("Worker Session append lock was poisoned"))?;
|
||||||
|
write_to_dir(
|
||||||
|
&self.root.join(PASTE_ARTIFACTS_DIR),
|
||||||
|
source_entry_id,
|
||||||
|
content,
|
||||||
|
limits,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_paste_artifact(
|
||||||
|
&self,
|
||||||
|
session_id: SessionId,
|
||||||
|
artifact_id: &str,
|
||||||
|
) -> Result<(PasteArtifactRef, String), StoreError> {
|
||||||
|
self.ensure_session(session_id, false)?;
|
||||||
|
read_from_dir(&self.root.join(PASTE_ARTIFACTS_DIR), artifact_id)
|
||||||
|
}
|
||||||
|
|
||||||
fn append_trace(
|
fn append_trace(
|
||||||
&self,
|
&self,
|
||||||
session_id: SessionId,
|
session_id: SessionId,
|
||||||
@@ -601,6 +633,45 @@ mod tests {
|
|||||||
assert_eq!(store.list_sessions().unwrap(), vec![session_id]);
|
assert_eq!(store.list_sessions().unwrap(), vec![session_id]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn worker_session_store_keeps_paste_artifacts_inside_retention_root() {
|
||||||
|
let root = tempfile::tempdir().unwrap();
|
||||||
|
let store = WorkerSessionStore::new(root.path().join("session")).unwrap();
|
||||||
|
let session_id = new_session_id();
|
||||||
|
store
|
||||||
|
.create_segment(session_id, new_segment_id(), &[])
|
||||||
|
.unwrap();
|
||||||
|
let content = "large paste body\n終端\n";
|
||||||
|
let reference = store
|
||||||
|
.write_paste_artifact(
|
||||||
|
session_id,
|
||||||
|
"entry-1",
|
||||||
|
content,
|
||||||
|
PasteArtifactLimits::default(),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
root.path()
|
||||||
|
.join(format!(
|
||||||
|
"session/{PASTE_ARTIFACTS_DIR}/{}.json",
|
||||||
|
reference.artifact_id
|
||||||
|
))
|
||||||
|
.is_file()
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
store
|
||||||
|
.read_paste_artifact(session_id, &reference.artifact_id)
|
||||||
|
.unwrap()
|
||||||
|
.1,
|
||||||
|
content
|
||||||
|
);
|
||||||
|
assert!(matches!(
|
||||||
|
store.read_paste_artifact(new_session_id(), &reference.artifact_id),
|
||||||
|
Err(StoreError::Corrupt { .. })
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn schema_v1_logs_are_rewritten_and_promoted_to_v3() {
|
fn schema_v1_logs_are_rewritten_and_promoted_to_v3() {
|
||||||
let root = tempfile::tempdir().unwrap();
|
let root = tempfile::tempdir().unwrap();
|
||||||
|
|||||||
+52
-2
@@ -61,6 +61,7 @@ impl FlowRefAtom {
|
|||||||
pub enum Atom {
|
pub enum Atom {
|
||||||
Char(char),
|
Char(char),
|
||||||
Paste(PasteRef),
|
Paste(PasteRef),
|
||||||
|
PasteArtifact(protocol::PasteArtifactRef),
|
||||||
FileRef(FileRefAtom),
|
FileRef(FileRefAtom),
|
||||||
FlowRef(FlowRefAtom),
|
FlowRef(FlowRefAtom),
|
||||||
}
|
}
|
||||||
@@ -72,6 +73,13 @@ impl Atom {
|
|||||||
match self {
|
match self {
|
||||||
Atom::Char(_) => None,
|
Atom::Char(_) => None,
|
||||||
Atom::Paste(p) => Some((Style::default().fg(Color::Magenta), p.label())),
|
Atom::Paste(p) => Some((Style::default().fg(Color::Magenta), p.label())),
|
||||||
|
Atom::PasteArtifact(artifact) => Some((
|
||||||
|
Style::default().fg(Color::Magenta),
|
||||||
|
format!(
|
||||||
|
"[Paste artifact {} | {} chars, {} lines]",
|
||||||
|
artifact.artifact_id, artifact.char_count, artifact.line_count
|
||||||
|
),
|
||||||
|
)),
|
||||||
Atom::FileRef(r) => Some((Style::default().fg(Color::Cyan), r.label())),
|
Atom::FileRef(r) => Some((Style::default().fg(Color::Cyan), r.label())),
|
||||||
Atom::FlowRef(r) => Some((Style::default().fg(Color::Yellow), r.label())),
|
Atom::FlowRef(r) => Some((Style::default().fg(Color::Yellow), r.label())),
|
||||||
}
|
}
|
||||||
@@ -102,7 +110,9 @@ enum WordKind {
|
|||||||
fn atom_class(atom: &Atom) -> AtomClass {
|
fn atom_class(atom: &Atom) -> AtomClass {
|
||||||
match atom {
|
match atom {
|
||||||
Atom::Char(c) => char_class(*c),
|
Atom::Char(c) => char_class(*c),
|
||||||
Atom::Paste(_) | Atom::FileRef(_) | Atom::FlowRef(_) => AtomClass::Chip,
|
Atom::Paste(_) | Atom::PasteArtifact(_) | Atom::FileRef(_) | Atom::FlowRef(_) => {
|
||||||
|
AtomClass::Chip
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -190,6 +200,9 @@ impl InputBuffer {
|
|||||||
content: content.clone(),
|
content: content.clone(),
|
||||||
}));
|
}));
|
||||||
}
|
}
|
||||||
|
protocol::Segment::PasteArtifact { artifact } => {
|
||||||
|
self.atoms.push(Atom::PasteArtifact(artifact.clone()));
|
||||||
|
}
|
||||||
protocol::Segment::FileRef { path } => {
|
protocol::Segment::FileRef { path } => {
|
||||||
self.atoms
|
self.atoms
|
||||||
.push(Atom::FileRef(FileRefAtom { path: path.clone() }));
|
.push(Atom::FileRef(FileRefAtom { path: path.clone() }));
|
||||||
@@ -225,6 +238,13 @@ impl InputBuffer {
|
|||||||
match atom {
|
match atom {
|
||||||
Atom::Char(c) => text.push(*c),
|
Atom::Char(c) => text.push(*c),
|
||||||
Atom::Paste(paste) => text.push_str(&paste.content),
|
Atom::Paste(paste) => text.push_str(&paste.content),
|
||||||
|
Atom::PasteArtifact(artifact) => {
|
||||||
|
text.push_str(&protocol::Segment::flatten_to_text(&[
|
||||||
|
protocol::Segment::PasteArtifact {
|
||||||
|
artifact: artifact.clone(),
|
||||||
|
},
|
||||||
|
]))
|
||||||
|
}
|
||||||
Atom::FileRef(file) => text.push_str(&file.path),
|
Atom::FileRef(file) => text.push_str(&file.path),
|
||||||
Atom::FlowRef(flow) => text.push_str(&flow.selector),
|
Atom::FlowRef(flow) => text.push_str(&flow.selector),
|
||||||
}
|
}
|
||||||
@@ -497,6 +517,12 @@ impl InputBuffer {
|
|||||||
content: p.content.clone(),
|
content: p.content.clone(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
Atom::PasteArtifact(artifact) => {
|
||||||
|
flush_text(&mut buf, &mut out);
|
||||||
|
out.push(protocol::Segment::PasteArtifact {
|
||||||
|
artifact: artifact.clone(),
|
||||||
|
});
|
||||||
|
}
|
||||||
Atom::FileRef(r) => {
|
Atom::FileRef(r) => {
|
||||||
flush_text(&mut buf, &mut out);
|
flush_text(&mut buf, &mut out);
|
||||||
out.push(protocol::Segment::FileRef {
|
out.push(protocol::Segment::FileRef {
|
||||||
@@ -902,6 +928,28 @@ mod submit_segments_tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn restored_paste_artifact_remains_a_typed_segment() {
|
||||||
|
let artifact = protocol::PasteArtifactRef {
|
||||||
|
artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b2".to_string(),
|
||||||
|
byte_len: 65_536,
|
||||||
|
char_count: 65_530,
|
||||||
|
line_count: 200,
|
||||||
|
sha256: "a".repeat(64),
|
||||||
|
source_entry_id: "entry-1".to_string(),
|
||||||
|
};
|
||||||
|
let original = Segment::PasteArtifact {
|
||||||
|
artifact: artifact.clone(),
|
||||||
|
};
|
||||||
|
let mut buf = InputBuffer::new();
|
||||||
|
buf.replace_with_segments(std::slice::from_ref(&original));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
buf.submit_segments(),
|
||||||
|
vec![Segment::PasteArtifact { artifact }]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn empty_buffer_yields_empty_segments() {
|
fn empty_buffer_yields_empty_segments() {
|
||||||
let buf = InputBuffer::new();
|
let buf = InputBuffer::new();
|
||||||
@@ -1219,7 +1267,9 @@ mod word_motion_tests {
|
|||||||
for a in &buf.atoms {
|
for a in &buf.atoms {
|
||||||
match a {
|
match a {
|
||||||
Atom::Char(c) => out.push(*c),
|
Atom::Char(c) => out.push(*c),
|
||||||
Atom::Paste(_) | Atom::FileRef(_) | Atom::FlowRef(_) => out.push_str("<P>"),
|
Atom::Paste(_) | Atom::PasteArtifact(_) | Atom::FileRef(_) | Atom::FlowRef(_) => {
|
||||||
|
out.push_str("<P>")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
out
|
out
|
||||||
|
|||||||
@@ -1296,6 +1296,13 @@ fn chip_span_for(seg: &Segment, fallback: Style) -> (Style, String) {
|
|||||||
Style::default().fg(Color::Magenta),
|
Style::default().fg(Color::Magenta),
|
||||||
format!("[Clipboard #{id} | {chars} chars, {line_count} lines]"),
|
format!("[Clipboard #{id} | {chars} chars, {line_count} lines]"),
|
||||||
),
|
),
|
||||||
|
Segment::PasteArtifact { artifact } => (
|
||||||
|
Style::default().fg(Color::Magenta),
|
||||||
|
format!(
|
||||||
|
"[Paste artifact {} | {} chars, {} lines]",
|
||||||
|
artifact.artifact_id, artifact.char_count, artifact.line_count
|
||||||
|
),
|
||||||
|
),
|
||||||
Segment::FileRef { path } => (Style::default().fg(Color::Cyan), format!("@{path}")),
|
Segment::FileRef { path } => (Style::default().fg(Color::Cyan), format!("@{path}")),
|
||||||
Segment::Flow { selector } => (
|
Segment::Flow { selector } => (
|
||||||
Style::default().fg(Color::Yellow),
|
Style::default().fg(Color::Yellow),
|
||||||
@@ -1314,6 +1321,10 @@ fn segment_display_text(seg: &Segment) -> String {
|
|||||||
Segment::Paste {
|
Segment::Paste {
|
||||||
id, chars, lines, ..
|
id, chars, lines, ..
|
||||||
} => format!("[Clipboard #{id} | {chars} chars, {lines} lines]"),
|
} => format!("[Clipboard #{id} | {chars} chars, {lines} lines]"),
|
||||||
|
Segment::PasteArtifact { artifact } => format!(
|
||||||
|
"[Paste artifact {} | {} chars, {} lines]",
|
||||||
|
artifact.artifact_id, artifact.char_count, artifact.line_count
|
||||||
|
),
|
||||||
Segment::FileRef { path } => format!("@{path}"),
|
Segment::FileRef { path } => format!("@{path}"),
|
||||||
Segment::Flow { selector } => format!("[Flow: {selector}]"),
|
Segment::Flow { selector } => format!("[Flow: {selector}]"),
|
||||||
Segment::Unknown => "[unknown segment]".to_owned(),
|
Segment::Unknown => "[unknown segment]".to_owned(),
|
||||||
|
|||||||
@@ -898,6 +898,20 @@ where
|
|||||||
crate::spawn::tool::ParentNotificationTarget::Buffer(worker.notify_buffer_handle())
|
crate::spawn::tool::ParentNotificationTarget::Buffer(worker.notify_buffer_handle())
|
||||||
});
|
});
|
||||||
let prompts = worker.prompts().clone();
|
let prompts = worker.prompts().clone();
|
||||||
|
let paste_store = worker.store().clone();
|
||||||
|
let paste_session_id = worker.session_id();
|
||||||
|
worker
|
||||||
|
.engine_mut()
|
||||||
|
.register_tool(crate::paste_artifact_tool::search_input_artifact_tool(
|
||||||
|
paste_store.clone(),
|
||||||
|
paste_session_id,
|
||||||
|
));
|
||||||
|
worker
|
||||||
|
.engine_mut()
|
||||||
|
.register_tool(crate::paste_artifact_tool::read_input_artifact_tool(
|
||||||
|
paste_store,
|
||||||
|
paste_session_id,
|
||||||
|
));
|
||||||
// Resolve the existing Worker–Workdir binding into the domain provider.
|
// Resolve the existing Worker–Workdir binding into the domain provider.
|
||||||
// Tools only consume the provider handle; they do not own its root, cwd,
|
// Tools only consume the provider handle; they do not own its root, cwd,
|
||||||
// scope, or lifecycle. No-workdir Workers expose no local tools.
|
// scope, or lifecycle. No-workdir Workers expose no local tools.
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ pub mod hook;
|
|||||||
pub(crate) mod in_flight;
|
pub(crate) mod in_flight;
|
||||||
pub mod ipc;
|
pub mod ipc;
|
||||||
pub mod model_client;
|
pub mod model_client;
|
||||||
|
mod paste_artifact_tool;
|
||||||
pub mod prompt;
|
pub mod prompt;
|
||||||
pub mod runtime;
|
pub mod runtime;
|
||||||
pub mod runtime_command;
|
pub mod runtime_command;
|
||||||
|
|||||||
@@ -0,0 +1,347 @@
|
|||||||
|
//! Bounded model-facing access to session-owned large paste artifacts.
|
||||||
|
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use schemars::JsonSchema;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use session_store::{SessionId, Store, StoreError};
|
||||||
|
|
||||||
|
const MAX_QUERY_BYTES: usize = 256;
|
||||||
|
const DEFAULT_SEARCH_RESULTS: usize = 20;
|
||||||
|
const MAX_SEARCH_RESULTS: usize = 100;
|
||||||
|
const MAX_SNIPPET_CHARS: usize = 300;
|
||||||
|
const DEFAULT_READ_BYTES: usize = 8 * 1024;
|
||||||
|
const MAX_READ_BYTES: usize = 16 * 1024;
|
||||||
|
|
||||||
|
const SEARCH_DESCRIPTION: &str = "Search one large pasted-input artifact owned by the current Worker. Returns bounded matching line snippets; never returns the whole artifact.";
|
||||||
|
const READ_DESCRIPTION: &str = "Read a bounded UTF-8 byte range from one large pasted-input artifact owned by the current Worker. Use next_offset for repeated calls instead of requesting the whole artifact.";
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct ArtifactAccess<St: Store + Clone> {
|
||||||
|
store: St,
|
||||||
|
session_id: SessionId,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize, JsonSchema)]
|
||||||
|
struct SearchInputArtifactInput {
|
||||||
|
/// Opaque artifact id from a large-paste history reference.
|
||||||
|
artifact_id: String,
|
||||||
|
/// Literal case-sensitive text to find.
|
||||||
|
query: String,
|
||||||
|
/// Maximum matching lines to return (1..=100).
|
||||||
|
max_results: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct SearchInputArtifactOutput {
|
||||||
|
artifact_id: String,
|
||||||
|
matches: Vec<SearchMatch>,
|
||||||
|
truncated: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct SearchMatch {
|
||||||
|
line: u64,
|
||||||
|
byte_offset: u64,
|
||||||
|
snippet: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct SearchInputArtifactTool<St: Store + Clone> {
|
||||||
|
access: ArtifactAccess<St>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl<St> Tool for SearchInputArtifactTool<St>
|
||||||
|
where
|
||||||
|
St: Store + Clone + Send + Sync + 'static,
|
||||||
|
{
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
input_json: &str,
|
||||||
|
_context: ToolExecutionContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
let input: SearchInputArtifactInput =
|
||||||
|
serde_json::from_str(input_json).map_err(|error| {
|
||||||
|
ToolError::InvalidArgument(format!("invalid SearchInputArtifact input: {error}"))
|
||||||
|
})?;
|
||||||
|
if input.query.is_empty() || input.query.len() > MAX_QUERY_BYTES {
|
||||||
|
return Err(ToolError::InvalidArgument(
|
||||||
|
"query must contain 1..=256 UTF-8 bytes".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let max_results = input
|
||||||
|
.max_results
|
||||||
|
.unwrap_or(DEFAULT_SEARCH_RESULTS)
|
||||||
|
.clamp(1, MAX_SEARCH_RESULTS);
|
||||||
|
let (_, content) = self
|
||||||
|
.access
|
||||||
|
.store
|
||||||
|
.read_paste_artifact(self.access.session_id, &input.artifact_id)
|
||||||
|
.map_err(tool_store_error)?;
|
||||||
|
let mut matches = Vec::new();
|
||||||
|
let mut truncated = false;
|
||||||
|
let mut byte_offset = 0_u64;
|
||||||
|
for (index, raw_line) in content.split_inclusive('\n').enumerate() {
|
||||||
|
let line = raw_line.strip_suffix('\n').unwrap_or(raw_line);
|
||||||
|
if line.contains(&input.query) {
|
||||||
|
if matches.len() == max_results {
|
||||||
|
truncated = true;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
matches.push(SearchMatch {
|
||||||
|
line: index as u64 + 1,
|
||||||
|
byte_offset,
|
||||||
|
snippet: truncate_chars(line, MAX_SNIPPET_CHARS),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
byte_offset += raw_line.len() as u64;
|
||||||
|
}
|
||||||
|
json_output(
|
||||||
|
format!("Found {} matching pasted-input line(s).", matches.len()),
|
||||||
|
&SearchInputArtifactOutput {
|
||||||
|
artifact_id: input.artifact_id,
|
||||||
|
matches,
|
||||||
|
truncated,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize, JsonSchema)]
|
||||||
|
struct ReadInputArtifactInput {
|
||||||
|
/// Opaque artifact id from a large-paste history reference.
|
||||||
|
artifact_id: String,
|
||||||
|
/// UTF-8 byte offset to start reading. Defaults to 0 and must be a character boundary.
|
||||||
|
offset: Option<u64>,
|
||||||
|
/// Maximum UTF-8 bytes to return (4..=16384). Defaults to 8192.
|
||||||
|
max_bytes: Option<usize>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct ReadInputArtifactOutput {
|
||||||
|
artifact_id: String,
|
||||||
|
offset: u64,
|
||||||
|
content: String,
|
||||||
|
next_offset: Option<u64>,
|
||||||
|
truncated: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct ReadInputArtifactTool<St: Store + Clone> {
|
||||||
|
access: ArtifactAccess<St>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl<St> Tool for ReadInputArtifactTool<St>
|
||||||
|
where
|
||||||
|
St: Store + Clone + Send + Sync + 'static,
|
||||||
|
{
|
||||||
|
async fn execute(
|
||||||
|
&self,
|
||||||
|
input_json: &str,
|
||||||
|
_context: ToolExecutionContext,
|
||||||
|
) -> Result<ToolOutput, ToolError> {
|
||||||
|
let input: ReadInputArtifactInput = serde_json::from_str(input_json).map_err(|error| {
|
||||||
|
ToolError::InvalidArgument(format!("invalid ReadInputArtifact input: {error}"))
|
||||||
|
})?;
|
||||||
|
let offset = input.offset.unwrap_or(0);
|
||||||
|
let max_bytes = input
|
||||||
|
.max_bytes
|
||||||
|
.unwrap_or(DEFAULT_READ_BYTES)
|
||||||
|
.clamp(4, MAX_READ_BYTES);
|
||||||
|
let (_, content) = self
|
||||||
|
.access
|
||||||
|
.store
|
||||||
|
.read_paste_artifact(self.access.session_id, &input.artifact_id)
|
||||||
|
.map_err(tool_store_error)?;
|
||||||
|
let offset = usize::try_from(offset).map_err(|_| {
|
||||||
|
ToolError::InvalidArgument("offset exceeds the artifact size".to_string())
|
||||||
|
})?;
|
||||||
|
if offset > content.len() || !content.is_char_boundary(offset) {
|
||||||
|
return Err(ToolError::InvalidArgument(
|
||||||
|
"offset must be a UTF-8 character boundary within the artifact".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let mut end = offset.saturating_add(max_bytes).min(content.len());
|
||||||
|
while end > offset && !content.is_char_boundary(end) {
|
||||||
|
end -= 1;
|
||||||
|
}
|
||||||
|
let output = content[offset..end].to_string();
|
||||||
|
let next_offset = (end < content.len()).then_some(end as u64);
|
||||||
|
let truncated = next_offset.is_some();
|
||||||
|
|
||||||
|
json_output(
|
||||||
|
format!("Read {} pasted-input byte(s).", output.len()),
|
||||||
|
&ReadInputArtifactOutput {
|
||||||
|
artifact_id: input.artifact_id,
|
||||||
|
offset: offset as u64,
|
||||||
|
content: output,
|
||||||
|
next_offset,
|
||||||
|
truncated,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn search_input_artifact_tool<St>(store: St, session_id: SessionId) -> ToolDefinition
|
||||||
|
where
|
||||||
|
St: Store + Clone + Send + Sync + 'static,
|
||||||
|
{
|
||||||
|
Arc::new(move || {
|
||||||
|
let schema = serde_json::to_value(schemars::schema_for!(SearchInputArtifactInput))
|
||||||
|
.unwrap_or_else(|_| serde_json::json!({}));
|
||||||
|
let meta = ToolMeta::new("SearchInputArtifact")
|
||||||
|
.description(SEARCH_DESCRIPTION)
|
||||||
|
.input_schema(schema);
|
||||||
|
let tool: Arc<dyn Tool> = Arc::new(SearchInputArtifactTool {
|
||||||
|
access: ArtifactAccess {
|
||||||
|
store: store.clone(),
|
||||||
|
session_id,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
(meta, tool)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn read_input_artifact_tool<St>(store: St, session_id: SessionId) -> ToolDefinition
|
||||||
|
where
|
||||||
|
St: Store + Clone + Send + Sync + 'static,
|
||||||
|
{
|
||||||
|
Arc::new(move || {
|
||||||
|
let schema = serde_json::to_value(schemars::schema_for!(ReadInputArtifactInput))
|
||||||
|
.unwrap_or_else(|_| serde_json::json!({}));
|
||||||
|
let meta = ToolMeta::new("ReadInputArtifact")
|
||||||
|
.description(READ_DESCRIPTION)
|
||||||
|
.input_schema(schema);
|
||||||
|
let tool: Arc<dyn Tool> = Arc::new(ReadInputArtifactTool {
|
||||||
|
access: ArtifactAccess {
|
||||||
|
store: store.clone(),
|
||||||
|
session_id,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
(meta, tool)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn json_output(summary: String, value: &impl Serialize) -> Result<ToolOutput, ToolError> {
|
||||||
|
let content = serde_json::to_string(value)
|
||||||
|
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
|
||||||
|
Ok(ToolOutput {
|
||||||
|
summary,
|
||||||
|
content: Some(content),
|
||||||
|
attachments: Vec::new(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn tool_store_error(error: StoreError) -> ToolError {
|
||||||
|
let message = match error {
|
||||||
|
StoreError::PasteArtifactNotFound(_) => "paste artifact not found",
|
||||||
|
StoreError::PasteArtifactIntegrity(_) | StoreError::Corrupt { .. } => {
|
||||||
|
"paste artifact failed its integrity check"
|
||||||
|
}
|
||||||
|
StoreError::PasteArtifactUnsupported => "paste artifact storage is unavailable",
|
||||||
|
_ => "paste artifact is unavailable",
|
||||||
|
};
|
||||||
|
ToolError::ExecutionFailed(message.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn truncate_chars(value: &str, limit: usize) -> String {
|
||||||
|
let mut chars = value.chars();
|
||||||
|
let truncated = chars.by_ref().take(limit).collect::<String>();
|
||||||
|
if chars.next().is_some() {
|
||||||
|
format!("{truncated}…")
|
||||||
|
} else {
|
||||||
|
truncated
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use agen::tool::ToolExecutionContext;
|
||||||
|
use session_store::{FsStore, PasteArtifactLimits, Store, new_session_id};
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn search_and_read_are_bounded_and_owner_scoped() {
|
||||||
|
let temp = tempfile::TempDir::new().unwrap();
|
||||||
|
let store = FsStore::new(temp.path()).unwrap();
|
||||||
|
let owner = new_session_id();
|
||||||
|
let other = new_session_id();
|
||||||
|
let content = (0..700)
|
||||||
|
.map(|index| format!("line {index}: needle {}", "x".repeat(80)))
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("\n");
|
||||||
|
let artifact = store
|
||||||
|
.write_paste_artifact(owner, "entry-1", &content, PasteArtifactLimits::default())
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let search = SearchInputArtifactTool {
|
||||||
|
access: ArtifactAccess {
|
||||||
|
store: store.clone(),
|
||||||
|
session_id: owner,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
let searched = search
|
||||||
|
.execute(
|
||||||
|
&serde_json::json!({
|
||||||
|
"artifact_id": artifact.artifact_id,
|
||||||
|
"query": "needle",
|
||||||
|
"max_results": 3
|
||||||
|
})
|
||||||
|
.to_string(),
|
||||||
|
ToolExecutionContext::default(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let searched: serde_json::Value =
|
||||||
|
serde_json::from_str(searched.content.as_deref().unwrap()).unwrap();
|
||||||
|
assert_eq!(searched["matches"].as_array().unwrap().len(), 3);
|
||||||
|
assert_eq!(searched["truncated"], true);
|
||||||
|
|
||||||
|
let read = ReadInputArtifactTool {
|
||||||
|
access: ArtifactAccess {
|
||||||
|
store: store.clone(),
|
||||||
|
session_id: owner,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
let read_output = read
|
||||||
|
.execute(
|
||||||
|
&serde_json::json!({
|
||||||
|
"artifact_id": artifact.artifact_id,
|
||||||
|
"offset": 2,
|
||||||
|
"max_bytes": 999999
|
||||||
|
})
|
||||||
|
.to_string(),
|
||||||
|
ToolExecutionContext::default(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
let read_output: serde_json::Value =
|
||||||
|
serde_json::from_str(read_output.content.as_deref().unwrap()).unwrap();
|
||||||
|
assert!(read_output["content"].as_str().unwrap().len() <= MAX_READ_BYTES);
|
||||||
|
assert_eq!(read_output["truncated"], true);
|
||||||
|
assert!(read_output["next_offset"].as_u64().is_some());
|
||||||
|
|
||||||
|
let foreign = ReadInputArtifactTool {
|
||||||
|
access: ArtifactAccess {
|
||||||
|
store,
|
||||||
|
session_id: other,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
let error = foreign
|
||||||
|
.execute(
|
||||||
|
&serde_json::json!({
|
||||||
|
"artifact_id": artifact.artifact_id,
|
||||||
|
"offset": 0,
|
||||||
|
"max_bytes": 1
|
||||||
|
})
|
||||||
|
.to_string(),
|
||||||
|
ToolExecutionContext::default(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap_err();
|
||||||
|
assert!(error.to_string().contains("paste artifact not found"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -62,6 +62,7 @@ pub(crate) fn metadata(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
pub(crate) fn history_entry(
|
pub(crate) fn history_entry(
|
||||||
item: Item,
|
item: Item,
|
||||||
origin: WorkerHistoryProvenance,
|
origin: WorkerHistoryProvenance,
|
||||||
@@ -69,6 +70,21 @@ pub(crate) fn history_entry(
|
|||||||
HistoryEntry::new(item, metadata(origin, None))
|
HistoryEntry::new(item, metadata(origin, None))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn history_entry_with_id(
|
||||||
|
item: Item,
|
||||||
|
entry_id: SessionHistoryEntryId,
|
||||||
|
origin: WorkerHistoryProvenance,
|
||||||
|
) -> HistoryEntry<SessionHistoryMetadata> {
|
||||||
|
HistoryEntry::new(
|
||||||
|
item,
|
||||||
|
SessionHistoryMetadata {
|
||||||
|
entry_id,
|
||||||
|
origin,
|
||||||
|
derivation: None,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn to_logged_history_entry(
|
pub(crate) fn to_logged_history_entry(
|
||||||
entry: &HistoryEntry<SessionHistoryMetadata>,
|
entry: &HistoryEntry<SessionHistoryMetadata>,
|
||||||
) -> LoggedHistoryEntry {
|
) -> LoggedHistoryEntry {
|
||||||
|
|||||||
+218
-12
@@ -15,8 +15,8 @@ use agen::{
|
|||||||
};
|
};
|
||||||
use arc_swap::ArcSwap;
|
use arc_swap::ArcSwap;
|
||||||
use session_store::{
|
use session_store::{
|
||||||
LogEntry, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, Store, StoreError,
|
LogEntry, PasteArtifactLimits, PromptRenderProvenance, SegmentId, SessionExtension, SessionId,
|
||||||
SystemItem, segment_log,
|
Store, StoreError, SystemItem, segment_log,
|
||||||
};
|
};
|
||||||
use session_store::{
|
use session_store::{
|
||||||
WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore, WorkerReclaimedChild,
|
WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore, WorkerReclaimedChild,
|
||||||
@@ -25,10 +25,12 @@ use session_store::{
|
|||||||
use tracing::{info, warn};
|
use tracing::{info, warn};
|
||||||
|
|
||||||
use crate::segment_log_sink::SegmentLogSink;
|
use crate::segment_log_sink::SegmentLogSink;
|
||||||
|
#[cfg(test)]
|
||||||
|
use crate::session_history::history_entry;
|
||||||
use crate::session_history::{
|
use crate::session_history::{
|
||||||
SessionHistoryDerivation, SessionHistoryMetadata, WorkerHistoryProvenance, history_entry,
|
SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata,
|
||||||
metadata as new_history_metadata, restore_history_entries, to_logged_history_entry,
|
WorkerHistoryProvenance, history_entry_with_id, metadata as new_history_metadata,
|
||||||
worker_subject,
|
restore_history_entries, to_logged_history_entry, worker_subject,
|
||||||
};
|
};
|
||||||
|
|
||||||
use manifest::{
|
use manifest::{
|
||||||
@@ -58,6 +60,7 @@ use crate::internal_worker::{
|
|||||||
};
|
};
|
||||||
|
|
||||||
const COMPACTION_EXTENSION_DOMAIN: &str = "yoi.compaction";
|
const COMPACTION_EXTENSION_DOMAIN: &str = "yoi.compaction";
|
||||||
|
const LARGE_PASTE_INLINE_MAX_BYTES: usize = 32 * 1024;
|
||||||
const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration";
|
const WORKER_ORCHESTRATION_INSTRUCTION_ID: &str = "worker.orchestration";
|
||||||
const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration";
|
const WORKER_ORCHESTRATION_PROMPT_REF: &str = "common.worker_orchestration";
|
||||||
|
|
||||||
@@ -2797,7 +2800,15 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
|||||||
St: Clone + 'static,
|
St: Clone + 'static,
|
||||||
F: FnOnce(),
|
F: FnOnce(),
|
||||||
{
|
{
|
||||||
let (input, pending_flow_state, flow_projection) = self.prepare_flow_input(input)?;
|
let (mut input, pending_flow_state, flow_projection) = self.prepare_flow_input(input)?;
|
||||||
|
let projected_entry_ids = if flow_projection.is_some() {
|
||||||
|
(0..input.len())
|
||||||
|
.map(|_| SessionHistoryEntryId::new())
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
} else {
|
||||||
|
vec![SessionHistoryEntryId::new()]
|
||||||
|
};
|
||||||
|
self.materialize_large_pastes(&mut input, &projected_entry_ids, flow_projection.is_some())?;
|
||||||
if let Some(state) = pending_flow_state.as_ref() {
|
if let Some(state) = pending_flow_state.as_ref() {
|
||||||
let payload = serde_json::to_value(state).map_err(|error| {
|
let payload = serde_json::to_value(state).map_err(|error| {
|
||||||
WorkerError::FlowInput(format!("serialize Flow runtime state: {error}"))
|
WorkerError::FlowInput(format!("serialize Flow runtime state: {error}"))
|
||||||
@@ -2830,7 +2841,8 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
|||||||
trigger: protocol::InvokeKind::UserSend,
|
trigger: protocol::InvokeKind::UserSend,
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let projected_input = self.projected_input_history(&input, flow_projection.as_ref());
|
let projected_input =
|
||||||
|
self.projected_input_history(&input, flow_projection.as_ref(), &projected_entry_ids);
|
||||||
|
|
||||||
// Persist original typed segments together with the exact ordered
|
// Persist original typed segments together with the exact ordered
|
||||||
// model-visible item+origin projection before any entry becomes live.
|
// model-visible item+origin projection before any entry becomes live.
|
||||||
@@ -3064,17 +3076,61 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn materialize_large_pastes(
|
||||||
|
&self,
|
||||||
|
input: &mut [Segment],
|
||||||
|
projected_entry_ids: &[SessionHistoryEntryId],
|
||||||
|
one_entry_per_segment: bool,
|
||||||
|
) -> Result<(), WorkerError> {
|
||||||
|
for (index, segment) in input.iter_mut().enumerate() {
|
||||||
|
if let Segment::PasteArtifact { artifact } = segment {
|
||||||
|
let (stored, _) = self
|
||||||
|
.store
|
||||||
|
.read_paste_artifact(self.session_id(), &artifact.artifact_id)?;
|
||||||
|
if &stored != artifact {
|
||||||
|
return Err(WorkerError::Store(StoreError::PasteArtifactIntegrity(
|
||||||
|
artifact.artifact_id.clone(),
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let Segment::Paste { content, .. } = segment else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
if content.len() <= LARGE_PASTE_INLINE_MAX_BYTES {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let entry_index = if one_entry_per_segment { index } else { 0 };
|
||||||
|
let source_entry_id = projected_entry_ids
|
||||||
|
.get(entry_index)
|
||||||
|
.expect("projected input id exists for every paste")
|
||||||
|
.0
|
||||||
|
.as_str();
|
||||||
|
let artifact = self.store.write_paste_artifact(
|
||||||
|
self.session_id(),
|
||||||
|
source_entry_id,
|
||||||
|
content,
|
||||||
|
PasteArtifactLimits::default(),
|
||||||
|
)?;
|
||||||
|
*segment = Segment::PasteArtifact { artifact };
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
fn projected_input_history(
|
fn projected_input_history(
|
||||||
&self,
|
&self,
|
||||||
input: &[Segment],
|
input: &[Segment],
|
||||||
flow_projection: Option<&PreparedFlowProjection>,
|
flow_projection: Option<&PreparedFlowProjection>,
|
||||||
|
entry_ids: &[SessionHistoryEntryId],
|
||||||
) -> Vec<HistoryEntry<SessionHistoryMetadata>> {
|
) -> Vec<HistoryEntry<SessionHistoryMetadata>> {
|
||||||
if let Some(flow) = flow_projection {
|
if let Some(flow) = flow_projection {
|
||||||
return input
|
return input
|
||||||
.iter()
|
.iter()
|
||||||
.map(|segment| match segment {
|
.zip(entry_ids)
|
||||||
Segment::Flow { .. } => history_entry(
|
.map(|(segment, entry_id)| match segment {
|
||||||
|
Segment::Flow { .. } => history_entry_with_id(
|
||||||
Item::user_message(flow.instructions.clone()),
|
Item::user_message(flow.instructions.clone()),
|
||||||
|
entry_id.clone(),
|
||||||
WorkerHistoryProvenance::FlowInstruction {
|
WorkerHistoryProvenance::FlowInstruction {
|
||||||
selector: flow.selector.clone(),
|
selector: flow.selector.clone(),
|
||||||
definition_id: flow.definition_id.clone(),
|
definition_id: flow.definition_id.clone(),
|
||||||
@@ -3083,8 +3139,9 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
|||||||
state_id: flow.state_id.clone(),
|
state_id: flow.state_id.clone(),
|
||||||
},
|
},
|
||||||
),
|
),
|
||||||
other => history_entry(
|
other => history_entry_with_id(
|
||||||
Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))),
|
Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))),
|
||||||
|
entry_id.clone(),
|
||||||
// Current public submit transport does not carry a
|
// Current public submit transport does not carry a
|
||||||
// trusted account/Worker subject envelope. Fail closed
|
// trusted account/Worker subject envelope. Fail closed
|
||||||
// instead of promoting role=user to HumanInput.
|
// instead of promoting role=user to HumanInput.
|
||||||
@@ -3094,8 +3151,12 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
|||||||
.collect();
|
.collect();
|
||||||
}
|
}
|
||||||
|
|
||||||
vec![history_entry(
|
vec![history_entry_with_id(
|
||||||
Item::user_message(Segment::flatten_to_text(input)),
|
Item::user_message(Segment::flatten_to_text(input)),
|
||||||
|
entry_ids
|
||||||
|
.first()
|
||||||
|
.expect("projected Worker input always has one entry id")
|
||||||
|
.clone(),
|
||||||
WorkerHistoryProvenance::LegacyUnknown,
|
WorkerHistoryProvenance::LegacyUnknown,
|
||||||
)]
|
)]
|
||||||
}
|
}
|
||||||
@@ -6405,6 +6466,11 @@ fn preview_segments(segments: &[Segment]) -> String {
|
|||||||
match segment {
|
match segment {
|
||||||
Segment::Text { content } => preview.push_str(content.trim()),
|
Segment::Text { content } => preview.push_str(content.trim()),
|
||||||
Segment::Paste { content, .. } => preview.push_str(content.trim()),
|
Segment::Paste { content, .. } => preview.push_str(content.trim()),
|
||||||
|
Segment::PasteArtifact { artifact } => {
|
||||||
|
preview.push_str("[Large paste artifact: ");
|
||||||
|
preview.push_str(&artifact.artifact_id);
|
||||||
|
preview.push(']');
|
||||||
|
}
|
||||||
Segment::FileRef { path } => {
|
Segment::FileRef { path } => {
|
||||||
preview.push('@');
|
preview.push('@');
|
||||||
preview.push_str(path);
|
preview.push_str(path);
|
||||||
@@ -7909,7 +7975,9 @@ mod build_summary_prompt_tests {
|
|||||||
FLOW_RUNTIME_EXTENSION_DOMAIN,
|
FLOW_RUNTIME_EXTENSION_DOMAIN,
|
||||||
serde_json::to_value(&state).unwrap(),
|
serde_json::to_value(&state).unwrap(),
|
||||||
);
|
);
|
||||||
let projected = worker.projected_input_history(&segments, projection.as_ref());
|
let projected_ids = vec![SessionHistoryEntryId::new(), SessionHistoryEntryId::new()];
|
||||||
|
let projected =
|
||||||
|
worker.projected_input_history(&segments, projection.as_ref(), &projected_ids);
|
||||||
worker
|
worker
|
||||||
.commit_entry(LogEntry::AnnotatedUserInput {
|
.commit_entry(LogEntry::AnnotatedUserInput {
|
||||||
ts: segment_log::now_millis(),
|
ts: segment_log::now_millis(),
|
||||||
@@ -7981,6 +8049,144 @@ mod build_summary_prompt_tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn large_paste_is_stored_before_compact_history_is_committed() {
|
||||||
|
let (_dir, worker) = rewind_test_worker().await;
|
||||||
|
let exact = "x".repeat(LARGE_PASTE_INLINE_MAX_BYTES);
|
||||||
|
let exact_ids = vec![SessionHistoryEntryId::new()];
|
||||||
|
let mut exact_input = vec![Segment::Paste {
|
||||||
|
id: 1,
|
||||||
|
chars: exact.len() as u32,
|
||||||
|
lines: 1,
|
||||||
|
content: exact.clone(),
|
||||||
|
}];
|
||||||
|
worker
|
||||||
|
.materialize_large_pastes(&mut exact_input, &exact_ids, false)
|
||||||
|
.unwrap();
|
||||||
|
assert!(matches!(&exact_input[0], Segment::Paste { content, .. } if content == &exact));
|
||||||
|
let mut empty_input = vec![Segment::Paste {
|
||||||
|
id: 0,
|
||||||
|
chars: 0,
|
||||||
|
lines: 0,
|
||||||
|
content: String::new(),
|
||||||
|
}];
|
||||||
|
worker
|
||||||
|
.materialize_large_pastes(&mut empty_input, &exact_ids, false)
|
||||||
|
.unwrap();
|
||||||
|
assert!(matches!(
|
||||||
|
&empty_input[0],
|
||||||
|
Segment::Paste { content, .. } if content.is_empty()
|
||||||
|
));
|
||||||
|
|
||||||
|
let body = format!("{}\n終端\n", "多".repeat(12_000));
|
||||||
|
let entry_id = SessionHistoryEntryId::new();
|
||||||
|
let mut input = vec![Segment::Paste {
|
||||||
|
id: 2,
|
||||||
|
chars: body.chars().count() as u32,
|
||||||
|
lines: 3,
|
||||||
|
content: body.clone(),
|
||||||
|
}];
|
||||||
|
worker
|
||||||
|
.materialize_large_pastes(&mut input, std::slice::from_ref(&entry_id), false)
|
||||||
|
.unwrap();
|
||||||
|
let artifact = match &input[0] {
|
||||||
|
Segment::PasteArtifact { artifact } => artifact.clone(),
|
||||||
|
other => panic!("expected stored paste reference, got {other:?}"),
|
||||||
|
};
|
||||||
|
assert_eq!(artifact.source_entry_id, entry_id.0);
|
||||||
|
assert_eq!(artifact.byte_len, body.len() as u64);
|
||||||
|
assert_eq!(artifact.char_count, body.chars().count() as u64);
|
||||||
|
assert_eq!(
|
||||||
|
worker
|
||||||
|
.store
|
||||||
|
.read_paste_artifact(worker.session_id(), &artifact.artifact_id)
|
||||||
|
.unwrap()
|
||||||
|
.1,
|
||||||
|
body
|
||||||
|
);
|
||||||
|
worker
|
||||||
|
.materialize_large_pastes(&mut input, &[SessionHistoryEntryId::new()], false)
|
||||||
|
.unwrap();
|
||||||
|
assert!(matches!(
|
||||||
|
&input[0],
|
||||||
|
Segment::PasteArtifact { artifact: retained }
|
||||||
|
if retained.source_entry_id == artifact.source_entry_id
|
||||||
|
));
|
||||||
|
|
||||||
|
let history = worker.projected_input_history(&input, None, &[entry_id]);
|
||||||
|
assert!(!history[0].item.as_text().unwrap().contains("終端"));
|
||||||
|
append_test_entry(
|
||||||
|
&worker,
|
||||||
|
LogEntry::Invoke {
|
||||||
|
ts: segment_log::now_millis(),
|
||||||
|
trigger: protocol::InvokeKind::UserSend,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
worker
|
||||||
|
.commit_entry(LogEntry::AnnotatedUserInput {
|
||||||
|
ts: segment_log::now_millis(),
|
||||||
|
segments: input.clone(),
|
||||||
|
extensions: Vec::new(),
|
||||||
|
history: history.iter().map(to_logged_history_entry).collect(),
|
||||||
|
})
|
||||||
|
.unwrap();
|
||||||
|
let location = worker.segment_state.location();
|
||||||
|
let entries = worker
|
||||||
|
.store
|
||||||
|
.read_all(location.session_id, location.segment_id)
|
||||||
|
.unwrap();
|
||||||
|
let persisted = serde_json::to_string(&entries).unwrap();
|
||||||
|
assert!(!persisted.contains("終端"));
|
||||||
|
let state = session_store::collect_state(&entries);
|
||||||
|
assert!(matches!(
|
||||||
|
&state.user_segments[0][0],
|
||||||
|
Segment::PasteArtifact { artifact: restored }
|
||||||
|
if restored.artifact_id == artifact.artifact_id
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn large_paste_storage_failure_commits_no_input() {
|
||||||
|
let temp = tempfile::TempDir::new().unwrap();
|
||||||
|
let store = session_store::FsStore::new(temp.path()).unwrap();
|
||||||
|
let mut worker = Worker::new(
|
||||||
|
minimal_manifest(),
|
||||||
|
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
|
||||||
|
store.clone(),
|
||||||
|
WorkerWorkspaceContext::unavailable(None, "test unavailable"),
|
||||||
|
WorkerFilesystemAuthority::None,
|
||||||
|
Scope::empty(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
worker.ensure_segment_head().unwrap();
|
||||||
|
std::fs::write(
|
||||||
|
temp.path()
|
||||||
|
.join(worker.session_id().to_string())
|
||||||
|
.join("artifacts"),
|
||||||
|
"block artifact directory creation",
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let result = worker
|
||||||
|
.run(vec![Segment::Paste {
|
||||||
|
id: 1,
|
||||||
|
chars: (LARGE_PASTE_INLINE_MAX_BYTES + 1) as u32,
|
||||||
|
lines: 1,
|
||||||
|
content: "x".repeat(LARGE_PASTE_INLINE_MAX_BYTES + 1),
|
||||||
|
}])
|
||||||
|
.await;
|
||||||
|
assert!(matches!(result, Err(WorkerError::Store(StoreError::Io(_)))));
|
||||||
|
let location = worker.segment_state.location();
|
||||||
|
let entries = store
|
||||||
|
.read_all(location.session_id, location.segment_id)
|
||||||
|
.unwrap();
|
||||||
|
assert!(
|
||||||
|
!entries
|
||||||
|
.iter()
|
||||||
|
.any(|entry| matches!(entry, LogEntry::AnnotatedUserInput { .. }))
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
async fn rewind_test_worker() -> (
|
async fn rewind_test_worker() -> (
|
||||||
tempfile::TempDir,
|
tempfile::TempDir,
|
||||||
Worker<NoopClient, session_store::FsStore>,
|
Worker<NoopClient, session_store::FsStore>,
|
||||||
|
|||||||
@@ -613,7 +613,19 @@ async fn feature_flags_default_to_core_tool_surface_only() {
|
|||||||
|
|
||||||
let request = wait_for_captured_request(&client_for_assert).await;
|
let request = wait_for_captured_request(&client_for_assert).await;
|
||||||
let names = request_tool_names(&request);
|
let names = request_tool_names(&request);
|
||||||
assert_eq!(names, vec!["Bash", "Edit", "Glob", "Grep", "Read", "Write"]);
|
assert_eq!(
|
||||||
|
names,
|
||||||
|
vec![
|
||||||
|
"Bash",
|
||||||
|
"Edit",
|
||||||
|
"Glob",
|
||||||
|
"Grep",
|
||||||
|
"Read",
|
||||||
|
"ReadInputArtifact",
|
||||||
|
"SearchInputArtifact",
|
||||||
|
"Write",
|
||||||
|
]
|
||||||
|
);
|
||||||
assert!(!names.iter().any(|name| name == "TaskCreate"));
|
assert!(!names.iter().any(|name| name == "TaskCreate"));
|
||||||
assert!(!names.iter().any(|name| name == "WebSearch"));
|
assert!(!names.iter().any(|name| name == "WebSearch"));
|
||||||
assert!(!names.iter().any(|name| name == "SubWorkerSpawn"));
|
assert!(!names.iter().any(|name| name == "SubWorkerSpawn"));
|
||||||
|
|||||||
@@ -133,7 +133,9 @@ message: string,
|
|||||||
*/
|
*/
|
||||||
timestamp_ms: number, };
|
timestamp_ms: number, };
|
||||||
|
|
||||||
export type Segment = { "kind": "text", content: string, } | { "kind": "paste", id: number, chars: number, lines: number, content: string, } | { "kind": "file_ref", path: string, } | { "kind": "flow", selector: string, } | { "kind": "unknown" };
|
export type PasteArtifactRef = { artifact_id: string, byte_len: number, char_count: number, line_count: number, sha256: string, source_entry_id: string, };
|
||||||
|
|
||||||
|
export type Segment = { "kind": "text", content: string, } | { "kind": "paste", id: number, chars: number, lines: number, content: string, } | { "kind": "paste_artifact", artifact: PasteArtifactRef, } | { "kind": "file_ref", path: string, } | { "kind": "flow", selector: string, } | { "kind": "unknown" };
|
||||||
|
|
||||||
export type WorkerEvent = { "kind": "turn_ended", worker_name: string, } | { "kind": "errored", worker_name: string, message: string, } | { "kind": "shut_down", worker_name: string, } | { "kind": "scope_sub_delegated",
|
export type WorkerEvent = { "kind": "turn_ended", worker_name: string, } | { "kind": "errored", worker_name: string, message: string, } | { "kind": "shut_down", worker_name: string, } | { "kind": "scope_sub_delegated",
|
||||||
/**
|
/**
|
||||||
|
|||||||
@@ -137,6 +137,24 @@ function snapshotEvent(cwd: string, entries: unknown[] = []): Event {
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Deno.test("large paste segments project compact artifact metadata", () => {
|
||||||
|
const body = "secret pasted body";
|
||||||
|
const text = segmentsToText([{
|
||||||
|
kind: "paste_artifact",
|
||||||
|
artifact: {
|
||||||
|
artifact_id: "019ca7c8-57b6-7f05-8edf-524147aba7b2",
|
||||||
|
byte_len: 65536,
|
||||||
|
char_count: 65530,
|
||||||
|
line_count: 200,
|
||||||
|
sha256: "a".repeat(64),
|
||||||
|
source_entry_id: "entry-1",
|
||||||
|
},
|
||||||
|
}]);
|
||||||
|
assert(text.includes("019ca7c8-57b6-7f05-8edf-524147aba7b2"), "artifact id is visible");
|
||||||
|
assert(text.includes("65536 bytes"), "bounded size metadata is visible");
|
||||||
|
assert(!text.includes(body), "artifact body is not projected");
|
||||||
|
});
|
||||||
|
|
||||||
Deno.test("console routing projects live errors but not completion replies", () => {
|
Deno.test("console routing projects live errors but not completion replies", () => {
|
||||||
const errorEvent = {
|
const errorEvent = {
|
||||||
event: "error",
|
event: "error",
|
||||||
|
|||||||
@@ -1069,6 +1069,8 @@ export function segmentsToText(segments: Segment[]): string {
|
|||||||
case "paste":
|
case "paste":
|
||||||
return segment.content ||
|
return segment.content ||
|
||||||
`[paste ${segment.id}: ${segment.chars} chars / ${segment.lines} lines]`;
|
`[paste ${segment.id}: ${segment.chars} chars / ${segment.lines} lines]`;
|
||||||
|
case "paste_artifact":
|
||||||
|
return `[Large paste artifact ${segment.artifact.artifact_id}: ${segment.artifact.byte_len} bytes, sha256 ${segment.artifact.sha256}]`;
|
||||||
case "file_ref":
|
case "file_ref":
|
||||||
return `@file ${segment.path}`;
|
return `@file ${segment.path}`;
|
||||||
case "unknown":
|
case "unknown":
|
||||||
|
|||||||
Reference in New Issue
Block a user