feat: store large paste inputs as artifacts
This commit is contained in:
@@ -898,6 +898,20 @@ where
|
||||
crate::spawn::tool::ParentNotificationTarget::Buffer(worker.notify_buffer_handle())
|
||||
});
|
||||
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.
|
||||
// Tools only consume the provider handle; they do not own its root, cwd,
|
||||
// scope, or lifecycle. No-workdir Workers expose no local tools.
|
||||
|
||||
@@ -9,6 +9,7 @@ pub mod hook;
|
||||
pub(crate) mod in_flight;
|
||||
pub mod ipc;
|
||||
pub mod model_client;
|
||||
mod paste_artifact_tool;
|
||||
pub mod prompt;
|
||||
pub mod runtime;
|
||||
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(
|
||||
item: Item,
|
||||
origin: WorkerHistoryProvenance,
|
||||
@@ -69,6 +70,21 @@ pub(crate) fn history_entry(
|
||||
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(
|
||||
entry: &HistoryEntry<SessionHistoryMetadata>,
|
||||
) -> LoggedHistoryEntry {
|
||||
|
||||
+218
-12
@@ -15,8 +15,8 @@ use agen::{
|
||||
};
|
||||
use arc_swap::ArcSwap;
|
||||
use session_store::{
|
||||
LogEntry, PromptRenderProvenance, SegmentId, SessionExtension, SessionId, Store, StoreError,
|
||||
SystemItem, segment_log,
|
||||
LogEntry, PasteArtifactLimits, PromptRenderProvenance, SegmentId, SessionExtension, SessionId,
|
||||
Store, StoreError, SystemItem, segment_log,
|
||||
};
|
||||
use session_store::{
|
||||
WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore, WorkerReclaimedChild,
|
||||
@@ -25,10 +25,12 @@ use session_store::{
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::segment_log_sink::SegmentLogSink;
|
||||
#[cfg(test)]
|
||||
use crate::session_history::history_entry;
|
||||
use crate::session_history::{
|
||||
SessionHistoryDerivation, SessionHistoryMetadata, WorkerHistoryProvenance, history_entry,
|
||||
metadata as new_history_metadata, restore_history_entries, to_logged_history_entry,
|
||||
worker_subject,
|
||||
SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata,
|
||||
WorkerHistoryProvenance, history_entry_with_id, metadata as new_history_metadata,
|
||||
restore_history_entries, to_logged_history_entry, worker_subject,
|
||||
};
|
||||
|
||||
use manifest::{
|
||||
@@ -58,6 +60,7 @@ use crate::internal_worker::{
|
||||
};
|
||||
|
||||
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_PROMPT_REF: &str = "common.worker_orchestration";
|
||||
|
||||
@@ -2797,7 +2800,15 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
||||
St: Clone + 'static,
|
||||
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() {
|
||||
let payload = serde_json::to_value(state).map_err(|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,
|
||||
})?;
|
||||
|
||||
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
|
||||
// model-visible item+origin projection before any entry becomes live.
|
||||
@@ -3064,17 +3076,61 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
||||
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(
|
||||
&self,
|
||||
input: &[Segment],
|
||||
flow_projection: Option<&PreparedFlowProjection>,
|
||||
entry_ids: &[SessionHistoryEntryId],
|
||||
) -> Vec<HistoryEntry<SessionHistoryMetadata>> {
|
||||
if let Some(flow) = flow_projection {
|
||||
return input
|
||||
.iter()
|
||||
.map(|segment| match segment {
|
||||
Segment::Flow { .. } => history_entry(
|
||||
.zip(entry_ids)
|
||||
.map(|(segment, entry_id)| match segment {
|
||||
Segment::Flow { .. } => history_entry_with_id(
|
||||
Item::user_message(flow.instructions.clone()),
|
||||
entry_id.clone(),
|
||||
WorkerHistoryProvenance::FlowInstruction {
|
||||
selector: flow.selector.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(),
|
||||
},
|
||||
),
|
||||
other => history_entry(
|
||||
other => history_entry_with_id(
|
||||
Item::user_message(Segment::flatten_to_text(std::slice::from_ref(other))),
|
||||
entry_id.clone(),
|
||||
// Current public submit transport does not carry a
|
||||
// trusted account/Worker subject envelope. Fail closed
|
||||
// instead of promoting role=user to HumanInput.
|
||||
@@ -3094,8 +3151,12 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
|
||||
.collect();
|
||||
}
|
||||
|
||||
vec![history_entry(
|
||||
vec![history_entry_with_id(
|
||||
Item::user_message(Segment::flatten_to_text(input)),
|
||||
entry_ids
|
||||
.first()
|
||||
.expect("projected Worker input always has one entry id")
|
||||
.clone(),
|
||||
WorkerHistoryProvenance::LegacyUnknown,
|
||||
)]
|
||||
}
|
||||
@@ -6405,6 +6466,11 @@ fn preview_segments(segments: &[Segment]) -> String {
|
||||
match segment {
|
||||
Segment::Text { 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 } => {
|
||||
preview.push('@');
|
||||
preview.push_str(path);
|
||||
@@ -7909,7 +7975,9 @@ mod build_summary_prompt_tests {
|
||||
FLOW_RUNTIME_EXTENSION_DOMAIN,
|
||||
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
|
||||
.commit_entry(LogEntry::AnnotatedUserInput {
|
||||
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() -> (
|
||||
tempfile::TempDir,
|
||||
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 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 == "WebSearch"));
|
||||
assert!(!names.iter().any(|name| name == "SubWorkerSpawn"));
|
||||
|
||||
Reference in New Issue
Block a user