Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: preserve original session coordinates through capture selection

Browse Source
This commit is contained in:
Hare
2026-10-22 23:17:25 +09:00
parent 994e75233b
commit 6776973e4c
9 changed files with 368 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/feature/builtin/flow_transition.rs
+1 -1
View File
@@ -474,7 +474,7 @@ where
self.segment_id,
&entries,
)?;
Ok(SessionCapture::from_history_entries(
Ok(SessionCapture::from_complete_history(
self.session_id.to_string(),
self.segment_id.to_string(),
history,
crates/worker/src/feature/builtin/memory_lifecycle.rs
+93 -6
View File
@@ -456,11 +456,13 @@ impl MemoryLifecycleTask {
.await;
return Ok(());
}
let view = SessionCapture::from_history_entries(
let view = SessionCapture::from_complete_history(
capture.session_id.clone(),
capture.segment_id.clone(),
capture.history[history_start..history_end].to_vec(),
);
capture.history.clone(),
)
.select_order_range(history_start as u64..history_end as u64)
.map_err(hook_internal)?;
let start_entry = pointer
.as_ref()
.map(|pointer| pointer.processed_through_entry + 1)
@@ -1557,7 +1559,7 @@ permission = "write"
#[tokio::test]
async fn run_committed_background_task_stages_non_empty_extraction_and_commits_pointer() {
let source = capture(2, 250);
let entry_ref = SessionCapture::from_history_entries(
let entry_ref = SessionCapture::from_complete_history(
source.session_id.clone(),
source.segment_id.clone(),
source.history.clone(),
@@ -1610,10 +1612,95 @@ permission = "write"
);
}
// Real background extraction -> common fixed selection -> StageMemoryCandidate
// -> typed backend operation and pointer write. Only model/backend transport
// are replaced; the second extraction uses the first run's saved pointer.
#[tokio::test]
async fn second_extraction_preserves_original_evidence_positions_after_nonzero_pointer() {
let first = capture(2, 250);
let mut second = capture(4, 750);
second.history[..2].clone_from_slice(&first.history);
second
.usage_history
.insert(0, first.usage_history[0].clone());
let selected_ref =
crate::SessionEntryRef::from_history_entry_id(&second.history[2].annotation.entry_id)
.to_string();
let extension_writes = Arc::new(Mutex::new(Vec::new()));
let workspace_client = Arc::new(RecordingWorkspaceClient::default());
let (event_tx, _) = broadcast::channel(64);
run_background_task(test_task(
first,
Box::new(ScriptClient::new(vec![
finish_empty_events("first-empty"),
completed_events(),
])),
extension_writes.clone(),
event_tx.clone(),
workspace_client.clone(),
))
.await;
{
let writes = extension_writes.lock().unwrap();
assert_eq!(writes.len(), 1);
let pointer: memory::ExtractPointerPayload =
serde_json::from_value(writes[0].2.clone()).unwrap();
assert_eq!(pointer.processed_through_history_len, 2);
second
.extensions
.push((writes[0].1.clone(), writes[0].2.clone()));
}
run_background_task(test_task(
second,
Box::new(ScriptClient::new(vec![
stage_candidate_events("second-stage", &selected_ref),
finish_events("second-finish", 1),
completed_events(),
])),
extension_writes.clone(),
event_tx,
workspace_client.clone(),
))
.await;
let operations = workspace_client
.requests
.lock()
.unwrap()
.iter()
.filter_map(|request| {
request.body.as_deref().and_then(|body| {
serde_json::from_str::<memory::backend::MemoryBackendOperation>(body).ok()
})
})
.filter_map(|operation| match operation {
memory::backend::MemoryBackendOperation::StageCandidate(operation) => {
Some(operation)
}
_ => None,
})
.collect::<Vec<_>>();
assert_eq!(operations.len(), 1);
let operation = &operations[0];
assert_eq!(operation.source.range, [2, 3]);
assert_eq!(operation.evidence[0].id, selected_ref);
assert_eq!(operation.evidence[0].entry_range, Some([2, 2]));
assert_eq!(
operation.source_refs[0].segment_id.as_deref(),
Some("segment-1")
);
assert_eq!(operation.source_refs[0].entry_range, Some([2, 2]));
let writes = extension_writes.lock().unwrap();
assert_eq!(writes.len(), 2);
let pointer: memory::ExtractPointerPayload =
serde_json::from_value(writes[1].2.clone()).unwrap();
assert_eq!(pointer.processed_through_history_len, 4);
assert_eq!(pointer.processed_through_entry, 3);
}
#[tokio::test]
async fn subjektiv_extraction_stages_once_and_commits_its_independent_pointer() {
let source = capture(2, 250);
let entry_ref = SessionCapture::from_history_entries(
let entry_ref = SessionCapture::from_complete_history(
source.session_id.clone(),
source.segment_id.clone(),
source.history.clone(),
@@ -1722,7 +1809,7 @@ permission = "write"
#[tokio::test]
async fn subjektiv_candidate_remains_staged_when_later_extraction_fails_without_pointer() {
let source = capture(2, 250);
let entry_ref = SessionCapture::from_history_entries(
let entry_ref = SessionCapture::from_complete_history(
source.session_id.clone(),
source.segment_id.clone(),
source.history.clone(),
crates/worker/src/feature/builtin/memory_staging_output.rs
+49
View File
@@ -639,4 +639,53 @@ mod tests {
.unwrap_err();
assert!(format!("{error:?}").contains("unknown SessionEntryRef"));
}
#[tokio::test]
async fn staging_roundtripped_mixed_sources_uses_original_evidence_ranges_not_view_order() {
let client = Arc::new(RecordingSubjektivClient::default());
let mut dto = SessionCapture::new(
"session",
"view",
vec![
Item::user_message("first"),
Item::assistant_message("second"),
],
)
.to_dto();
dto.entries[0].source.segment_id = "original-a".into();
dto.entries[0].source.entry_index = 80;
dto.entries[1].source.segment_id = "original-b".into();
dto.entries[1].source.entry_index = 3;
let limits = crate::session_capture::CaptureLimits::default();
let view = SessionCapture::decode(&dto.encode(limits).unwrap(), limits).unwrap();
let state = MemoryStagingOutputState::new_subjektiv(
view,
crate::subjektiv::test_connection(client.clone()),
SourceRef {
segment_id: "view".into(),
range: [0, 1],
},
"run".into(),
);
StageMemoryCandidateTool { state }.execute(
r#"{"kind":"decision","claim":"claim","why_useful":"useful","entry_refs":["E00000000","E00000001"]}"#,
agen::tool::ToolExecutionContext::direct(),
).await.unwrap();
let requests = client.requests.lock().unwrap();
let request: server_api::SubjektivStageCandidateRequest =
serde_json::from_str(requests[0].body.as_deref().unwrap()).unwrap();
for (index, segment, position) in [(0, "original-a", 80), (1, "original-b", 3)] {
assert_eq!(
request.operation.evidence[index].entry_range,
Some([position; 2])
);
assert_eq!(
request.operation.source_refs[index].segment_id.as_deref(),
Some(segment)
);
assert_eq!(
request.operation.source_refs[index].entry_range,
Some([position; 2])
);
}
}
}
crates/worker/src/feature/builtin/subjektiv_memory.rs
+3 -3
View File
@@ -571,7 +571,7 @@ impl SubjektivMemoryState {
.capture()
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
let session_id = capture.session_id.clone();
let view = SessionCapture::from_history_entries(
let view = SessionCapture::from_complete_history(
capture.session_id.clone(),
capture.segment_id,
capture.history,
@@ -680,7 +680,7 @@ fn replay_pending_committed_memory_with_policy(
.capture
.capture()
.map_err(|error| HookError::new(HookErrorCategory::Dependency, error.to_string()))?;
let view = SessionCapture::from_history_entries(
let view = SessionCapture::from_complete_history(
capture.session_id.clone(),
capture.segment_id.clone(),
capture.history.clone(),
@@ -815,7 +815,7 @@ fn resolve_entry_refs(
capture: &crate::feature::session::CommittedSessionCapture,
entry_refs: &[String],
) -> Result<Vec<SessionEntryEvidence>, ToolError> {
let view = SessionCapture::from_history_entries(
let view = SessionCapture::from_complete_history(
capture.session_id.clone(),
capture.segment_id.clone(),
capture.history.clone(),
crates/worker/src/session_capture.rs
+8 -7
View File
@@ -30,7 +30,7 @@ mod tests {
},
),
];
let capture = SessionCapture::from_history_entries("session", "segment", entries);
let capture = SessionCapture::from_complete_history("session", "segment", entries);
let overview = capture.overview();
assert_eq!(overview.len(), 1);
assert!(matches!(
@@ -54,8 +54,9 @@ mod tests {
},
);
let expected_ref = SessionEntryRef::from_history_entry_id(&retained.annotation.entry_id);
let before = SessionCapture::from_history_entries("session", "old", vec![retained.clone()]);
let after = SessionCapture::from_history_entries("session", "new", vec![retained]);
let before =
SessionCapture::from_complete_history("session", "old", vec![retained.clone()]);
let after = SessionCapture::from_complete_history("session", "new", vec![retained]);
assert_eq!(before.overview()[0].id, expected_ref);
assert_eq!(after.overview()[0].id, expected_ref);
assert_eq!(
@@ -96,7 +97,7 @@ mod tests {
.iter()
.map(|entry| SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id))
.collect::<Vec<_>>();
let capture = SessionCapture::from_history_entries("session", "segment", entries);
let capture = SessionCapture::from_complete_history("session", "segment", entries);
assert!(capture.overview().is_empty());
assert!(
capture
@@ -235,7 +236,7 @@ mod tests {
let snapshot: SessionSnapshot =
serde_json::from_value(serde_json::to_value(snapshot).unwrap()).unwrap();
for capture in [
SessionCapture::from_history_entries(
SessionCapture::from_complete_history(
session.to_string(),
segment.to_string(),
restored,
@@ -369,7 +370,7 @@ mod tests {
assert_eq!(entry.item, items[index]);
assert_eq!(entry.annotation, history[index].metadata);
}
let capture = SessionCapture::from_history_entries(
let capture = SessionCapture::from_complete_history(
session.to_string(),
fork.to_string(),
restored,
@@ -757,7 +758,7 @@ mod tests {
let restored =
crate::session_history::restore_history_entries(session, first, &log).unwrap();
let fixed =
SessionCapture::from_history_entries(session.to_string(), first.to_string(), restored);
SessionCapture::from_complete_history(session.to_string(), first.to_string(), restored);
let overview = fixed.overview()[0].clone();
let evidence = fixed.evidence_for(overview.id.as_str()).unwrap();
let extract = render_extract_input(&fixed);
crates/worker/src/worker.rs
+16 -6
View File
@@ -7338,6 +7338,20 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.cloned()
.collect::<Vec<_>>();
project_compactor_capture_attachments(&mut entries_to_summarise);
// Project before Host selection so omitted summary/reference inputs never
// rebase the original source positions of the retained prefix entries.
let mut exploration_entries = history_entries[..retain_from].to_vec();
project_compactor_capture_attachments(&mut exploration_entries);
let summarised_ids = entries_to_summarise
.iter()
.map(|entry| entry.annotation.entry_id.0.as_str())
.collect::<std::collections::HashSet<_>>();
let exploration_view = crate::session_capture::SessionCapture::from_complete_history(
self.session_id().to_string(),
self.segment_id().to_string(),
exploration_entries,
)
.retain_entries(|entry| summarised_ids.contains(entry.entry_id.as_str()));
let items_to_summarise = entries_to_summarise
.iter()
.map(|entry| entry.item.clone())
@@ -7500,11 +7514,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.with_module(
crate::feature::builtin::session_explore::SessionExploreFeature::new(
crate::feature::builtin::session_explore::SessionExploreState::new(
crate::session_capture::SessionCapture::from_history_entries(
self.session_id().to_string(),
self.segment_id().to_string(),
entries_to_summarise.clone(),
),
exploration_view,
),
),
)
@@ -17375,7 +17385,7 @@ permission = "write"
producer: "arbitrary-context-provider".into()
}
);
let view = crate::session_capture::SessionCapture::from_history_entries(
let view = crate::session_capture::SessionCapture::from_complete_history(
"session", "startup", replay,
);
let extract =