feat(session): share immutable captures and bounded body transport

This commit is contained in:
2026-10-22 22:25:47 +09:00
parent 8b6e3070dd
commit 994e75233b
13 changed files with 2923 additions and 955 deletions
@@ -26,7 +26,7 @@ use flow::{
use manifest::{Scope, WorkerManifest};
use schemars::JsonSchema;
use serde::Deserialize;
use session_store::{SegmentId, SessionId, Store, collect_state};
use session_store::{SegmentId, SessionId, Store};
use uuid::Uuid;
use crate::feature::builtin::session_explore::{SessionExploreFeature, SessionExploreState};
@@ -469,10 +469,15 @@ where
.store
.read_all(self.session_id, self.segment_id)
.map_err(|error| format!("read committed parent session: {error}"))?;
let restored = collect_state(&entries);
Ok(SessionCapture::new(
let history = crate::session_history::restore_history_entries(
self.session_id,
self.segment_id,
&entries,
)?;
Ok(SessionCapture::from_history_entries(
self.session_id.to_string(),
self.segment_id.to_string(),
restored.history,
history,
))
}
}
@@ -457,6 +457,7 @@ impl MemoryLifecycleTask {
return Ok(());
}
let view = SessionCapture::from_history_entries(
capture.session_id.clone(),
capture.segment_id.clone(),
capture.history[history_start..history_end].to_vec(),
);
@@ -1556,11 +1557,14 @@ 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(source.segment_id.clone(), source.history.clone())
.overview()[0]
.id
.to_string();
let entry_ref = SessionCapture::from_history_entries(
source.session_id.clone(),
source.segment_id.clone(),
source.history.clone(),
)
.overview()[0]
.id
.to_string();
let client = ScriptClient::new(vec![
stage_candidate_events("stage-1", &entry_ref),
finish_events("finish-1", 1),
@@ -1609,11 +1613,14 @@ permission = "write"
#[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(source.segment_id.clone(), source.history.clone())
.overview()[0]
.id
.to_string();
let entry_ref = SessionCapture::from_history_entries(
source.session_id.clone(),
source.segment_id.clone(),
source.history.clone(),
)
.overview()[0]
.id
.to_string();
let client = ScriptClient::new(vec![
stage_candidate_events("stage-1", &entry_ref),
finish_events("finish-1", 1),
@@ -1715,11 +1722,14 @@ 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(source.segment_id.clone(), source.history.clone())
.overview()[0]
.id
.to_string();
let entry_ref = SessionCapture::from_history_entries(
source.session_id.clone(),
source.segment_id.clone(),
source.history.clone(),
)
.overview()[0]
.id
.to_string();
let client = ScriptClient::new(vec![stage_candidate_events("stage-1", &entry_ref)]);
let extension_writes = Arc::new(Mutex::new(Vec::new()));
let (event_tx, _) = broadcast::channel(32);
@@ -510,7 +510,11 @@ mod tests {
fn state() -> MemoryStagingOutputState {
MemoryStagingOutputState::new(
SessionCapture::new("segment-1", vec![Item::user_message("durable decision")]),
SessionCapture::new(
"session",
"segment-1",
vec![Item::user_message("durable decision")],
),
crate::worker::marker_workspace_client(None, "test-backend"),
SourceRef {
segment_id: "segment-1".to_string(),
@@ -537,6 +541,7 @@ mod tests {
#[test]
fn render_input_uses_session_entry_refs_and_new_tool_names() {
let view = SessionCapture::new(
"session",
"segment-1",
vec![
Item::user_message("preference"),
@@ -586,7 +591,11 @@ mod tests {
async fn subjektiv_destination_sends_committed_session_without_subject_scope() {
let client = Arc::new(RecordingSubjektivClient::default());
let state = MemoryStagingOutputState::new_subjektiv(
SessionCapture::new("segment-1", vec![Item::user_message("durable decision")]),
SessionCapture::new(
"session",
"segment-1",
vec![Item::user_message("durable decision")],
),
{
let mut host = crate::subjektiv::test_connection(client.clone());
host.context.session_id = "session-committed".into();
@@ -17,14 +17,14 @@ use crate::feature::{
ToolDeclaration,
};
use crate::session_capture::{
ReadDetail, ReadOptions, ReadSelector, ReferenceKind, SearchOptions, SessionCapture,
SessionEntryRef, ToolPart,
CapturePurpose, RawField, ReadDetail, ReadOptions, ReadSelector, ReferenceKind, SearchOptions,
SessionCapture, SessionEntryRef, ToolPart,
};
const SHOW_OVERVIEW_DESCRIPTION: &str =
"Show a sparse, bounded index of real user and assistant session entries.";
const SEARCH_ENTRIES_DESCRIPTION: &str = "Search or compactly list a bounded range of committed session entries. Reasoning entries are never exposed.";
const READ_ENTRY_DESCRIPTION: &str = "Read one committed session entry by stable SessionEntryRef. Compact mode is the default; full mode includes bounded tool input or output.";
const READ_ENTRY_DESCRIPTION: &str = "Read one committed session entry by stable SessionEntryRef. Compact mode is the default; full mode is a bounded rendering, not raw source. Raw mode requires raw_field and supports UTF-8 byte_offset/next_offset without omission; use capture_id to bind continuation to this exact capture.";
const DEFAULT_PAGE_LIMIT: usize = 20;
const MAX_PAGE_LIMIT: usize = 100;
const DEFAULT_READ_ITEMS: usize = 10;
@@ -173,6 +173,38 @@ struct ReadEntryParams {
mode: String,
#[serde(default)]
max_items: Option<usize>,
#[serde(default)]
raw_field: Option<RawEntryField>,
#[serde(default)]
part_index: Option<usize>,
#[serde(default)]
byte_offset: Option<usize>,
#[serde(default)]
max_bytes: Option<usize>,
#[serde(default)]
capture_id: Option<String>,
}
#[derive(Debug, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
enum RawEntryField {
MessagePart,
PublicSystem,
ToolArguments,
ToolSummary,
ToolContent,
}
impl RawEntryField {
fn into_common(self, part_index: usize) -> RawField {
match self {
Self::MessagePart => RawField::MessagePart(part_index),
Self::PublicSystem => RawField::PublicSystem,
Self::ToolArguments => RawField::ToolArguments,
Self::ToolSummary => RawField::ToolSummary,
Self::ToolContent => RawField::ToolContent,
}
}
}
fn default_read_mode() -> String {
@@ -243,25 +275,20 @@ impl Tool for SearchEntriesTool {
let from = params.from.as_deref().map(parse_entry_ref).transpose()?;
let through = params.through.as_deref().map(parse_entry_ref).transpose()?;
let view = self.state.view();
if let (Some(from), Some(through)) = (&from, &through) {
if view.source_index_for_ref(from) > view.source_index_for_ref(through) {
return Err(ToolError::InvalidArgument(
"SearchEntries from must not be after through".to_string(),
));
}
}
let limit = bounded_limit(params.limit, DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT);
let hits = view.search(&SearchOptions {
query: params.query,
kind,
tool_part,
tool_name: params.tool_name,
limit: Some(limit),
min_entry_index: None,
from,
through,
offset: params.offset,
});
let hits = view
.search_checked(&SearchOptions {
query: params.query,
kind,
tool_part,
tool_name: params.tool_name,
limit: Some(limit),
min_entry_index: None,
from,
through,
offset: params.offset,
})
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
let entries = hits
.iter()
.map(|hit| {
@@ -272,6 +299,7 @@ impl Tool for SearchEntriesTool {
"tool_name": hit.tool_name,
"label": hit.label,
"text": hit.summary,
"text_kind": "preview",
})
})
.collect::<Vec<_>>();
@@ -299,12 +327,62 @@ impl Tool for ReadEntryTool {
) -> Result<ToolOutput, ToolError> {
let params: ReadEntryParams = parse_input("ReadEntry", input_json)?;
let entry_ref = parse_entry_ref(&params.entry_ref)?;
let view = self.state.view();
if params
.capture_id
.as_deref()
.is_some_and(|id| id != view.capture_id())
{
return Err(ToolError::InvalidArgument(
"capture identity does not match".into(),
));
}
if params.mode == "raw" {
let field = params
.raw_field
.ok_or_else(|| ToolError::InvalidArgument("raw mode requires raw_field".into()))?;
if params.part_index.is_some() && !matches!(field, RawEntryField::MessagePart) {
return Err(ToolError::InvalidArgument(
"part_index is only valid for message_part".into(),
));
}
let reference = view
.scoped_ref(&entry_ref)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
let raw = view
.read_raw(
&reference,
CapturePurpose::Observation,
field.into_common(params.part_index.unwrap_or(0)),
params.byte_offset.unwrap_or(0),
params
.max_bytes
.unwrap_or(MAX_READ_BYTES)
.clamp(1, MAX_READ_BYTES),
)
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?;
return json_output(
"Read raw session source bytes.".into(),
serde_json::json!({
"text_kind": "raw", "raw": raw,
}),
);
}
if params.raw_field.is_some()
|| params.part_index.is_some()
|| params.byte_offset.is_some()
|| params.max_bytes.is_some()
{
return Err(ToolError::InvalidArgument(
"raw fields require raw mode".into(),
));
}
let detail = match params.mode.as_str() {
"compact" => ReadDetail::Compact,
"full" => ReadDetail::Full,
other => {
return Err(ToolError::InvalidArgument(format!(
"invalid ReadEntry mode {other:?}; expected compact or full"
"invalid ReadEntry mode {other:?}; expected compact, full, or raw"
)));
}
};
@@ -347,6 +425,7 @@ impl Tool for ReadEntryTool {
serde_json::json!({
"entries": entries,
"truncated": read.truncated,
"text_kind": "rendered",
}),
)
}
@@ -408,6 +487,7 @@ mod tests {
fn state() -> SessionExploreState {
SessionExploreState::new(SessionCapture::new(
"session",
"segment-1",
vec![
Item::user_message("first"),
@@ -479,4 +559,64 @@ mod tests {
.unwrap();
assert!(output.content.unwrap().contains("second"));
}
#[tokio::test]
async fn search_rejects_unknown_or_reversed_capture_endpoints() {
let search = search_entries_definition(state())().1;
for input in [
r#"{"from":"E99999999"}"#,
r#"{"through":"E99999999"}"#,
r#"{"from":"E00000003","through":"E00000000"}"#,
] {
assert!(
search
.execute(input, agen::tool::ToolExecutionContext::direct())
.await
.is_err()
);
}
}
#[tokio::test]
async fn raw_tool_arguments_reassemble_without_preview_or_utf8_loss() {
let arguments = "日本語\n\"\\".repeat(100);
let state = SessionExploreState::new(SessionCapture::new(
"session",
"segment",
vec![Item::tool_call("call", "Read", arguments.clone())],
));
let capture_id = state.view().capture_id().to_owned();
let tool = read_entry_definition(state)().1;
let mut offset = 0;
let mut assembled = String::new();
loop {
let input = serde_json::json!({
"entry_ref": "E00000000", "mode": "raw", "raw_field": "tool_arguments",
"byte_offset": offset, "max_bytes": 17, "capture_id": capture_id,
});
let output = tool
.execute(
&input.to_string(),
agen::tool::ToolExecutionContext::direct(),
)
.await
.unwrap();
let value: serde_json::Value = serde_json::from_str(&output.content.unwrap()).unwrap();
assert_eq!(value["text_kind"], "raw");
let raw = &value["raw"];
assert_eq!(raw["offset"], offset);
let text = raw["text"].as_str().unwrap();
assert!(text.len() <= 17);
assembled.push_str(text);
match raw["next_offset"].as_u64() {
Some(next) => {
assert!(next > offset as u64);
offset = next as usize;
}
None => break,
}
}
assert_eq!(assembled, arguments);
assert!(tool.execute(r#"{"entry_ref":"E00000000","mode":"raw","raw_field":"tool_arguments","capture_id":"other"}"#,
agen::tool::ToolExecutionContext::direct()).await.is_err());
}
}
@@ -39,13 +39,66 @@ use crate::hook::{
HookErrorCategory, HookExecutionPolicy, HookFailurePolicy, HookPreRequestAction, PreLlmRequest,
PreRequestContext, RunCommitted, RunCommittedContext,
};
use crate::session_capture::{SessionCapture, SessionEntryEvidence};
use crate::session_capture::{CaptureBody, SessionCapture, SessionEntryEvidence};
use crate::subjektiv::SubjektivHostConnection;
#[cfg(test)]
use crate::worker::WorkspaceClient;
#[cfg(test)]
use server_api::SubjektivMemoryBackendRequest;
// Pending-commit receipt interpretation belongs to Memory, not the shared capture.
struct CommittedMemoryCall {
call_id: String,
name: String,
arguments: String,
evidence: SessionEntryEvidence,
}
fn committed_pending_tool_calls(view: &SessionCapture, names: &[&str]) -> Vec<CommittedMemoryCall> {
let pending: HashSet<_> = view
.entries()
.filter_map(|entry| {
let CaptureBody::ToolResult {
call_id,
content: Some(content),
is_error: false,
..
} = &entry.body
else {
return None;
};
(serde_json::from_str::<serde_json::Value>(content)
.ok()
.and_then(|value| value.get("status").cloned())
.as_ref()
.and_then(serde_json::Value::as_str)
== Some("pending_commit"))
.then_some(call_id.as_str())
})
.collect();
view.entries()
.filter_map(|entry| {
let CaptureBody::ToolCall {
call_id,
name,
arguments,
} = &entry.body
else {
return None;
};
if !names.contains(&name.as_str()) || !pending.contains(call_id.as_str()) {
return None;
}
Some(CommittedMemoryCall {
call_id: call_id.clone(),
name: name.clone(),
arguments: arguments.clone(),
evidence: view.evidence_for_tool_call(call_id)?,
})
})
.collect()
}
const QUERY_TOOL: &str = "SubjektivMemoryQuery";
const READ_TOOL: &str = "SubjektivMemoryRead";
const LIST_CHANGES_TOOL: &str = "SubjektivMemoryListChanges";
@@ -518,14 +571,18 @@ impl SubjektivMemoryState {
.capture()
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
let session_id = capture.session_id.clone();
let view = SessionCapture::from_history_entries(capture.segment_id, capture.history);
let committed = view
.committed_pending_tool_calls(&[REMEMBER_TOOL, PROPOSE_CHANGE_TOOL])
.into_iter()
.any(|call| {
receipt_id(&session_id, &call.call_id)
.is_ok_and(|committed_receipt| committed_receipt == expected_receipt_id)
});
let view = SessionCapture::from_history_entries(
capture.session_id.clone(),
capture.segment_id,
capture.history,
);
let committed =
committed_pending_tool_calls(&view, &[REMEMBER_TOOL, PROPOSE_CHANGE_TOOL])
.into_iter()
.any(|call| {
receipt_id(&session_id, &call.call_id)
.is_ok_and(|committed_receipt| committed_receipt == expected_receipt_id)
});
if committed {
status.status = SubjektivMemoryReceiptStatus::PendingCommit;
}
@@ -623,10 +680,13 @@ 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(capture.segment_id.clone(), capture.history.clone());
let view = SessionCapture::from_history_entries(
capture.session_id.clone(),
capture.segment_id.clone(),
capture.history.clone(),
);
let mut failures = Vec::new();
for call in view.committed_pending_tool_calls(&[REMEMBER_TOOL, PROPOSE_CHANGE_TOOL]) {
for call in committed_pending_tool_calls(&view, &[REMEMBER_TOOL, PROPOSE_CHANGE_TOOL]) {
let receipt_id = match receipt_id(&capture.session_id, &call.call_id) {
Ok(value) => value,
Err(error) => {
@@ -755,8 +815,11 @@ fn resolve_entry_refs(
capture: &crate::feature::session::CommittedSessionCapture,
entry_refs: &[String],
) -> Result<Vec<SessionEntryEvidence>, ToolError> {
let view =
SessionCapture::from_history_entries(capture.segment_id.clone(), capture.history.clone());
let view = SessionCapture::from_history_entries(
capture.session_id.clone(),
capture.segment_id.clone(),
capture.history.clone(),
);
let mut seen = HashSet::new();
entry_refs
.iter()
@@ -570,15 +570,8 @@ impl Tool for SearchSessionEntriesTool {
let view = latest_view(&*self.provider, &params.subject).await?;
let from = params.from.as_deref().map(parse_entry_ref).transpose()?;
let through = params.through.as_deref().map(parse_entry_ref).transpose()?;
if let (Some(from), Some(through)) = (&from, &through) {
if view.source_index_for_ref(from) > view.source_index_for_ref(through) {
return Err(ToolError::InvalidArgument(
"SearchSessionEntries from must not be after through".to_string(),
));
}
}
let entries = view
.search(&SearchOptions {
.search_checked(&SearchOptions {
query: params.query,
kind: params.kind.as_deref().map(parse_kind).transpose()?,
tool_part: params
@@ -593,6 +586,7 @@ impl Tool for SearchSessionEntriesTool {
through,
offset: params.offset,
})
.map_err(|error| ToolError::InvalidArgument(error.to_string()))?
.into_iter()
.map(|entry| {
serde_json::json!({
@@ -604,6 +598,7 @@ impl Tool for SearchSessionEntriesTool {
"tool_name": entry.tool_name,
"label": entry.label,
"text": entry.summary,
"text_kind": "preview",
})
})
.collect::<Vec<_>>();
@@ -687,7 +682,10 @@ async fn latest_view(
.capture_worker_session(subject)
.await
.map_err(tool_error)?;
// Observation providers expose a stable subject scope rather than a raw
// Session handle. This scope is never used to select recording material.
Ok(SessionCapture::from_session_snapshot(
format!("observation:{}", capture.segment_id),
capture.segment_id,
capture.session,
))
File diff suppressed because it is too large Load Diff
+4 -2
View File
@@ -7501,6 +7501,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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(),
),
@@ -17374,8 +17375,9 @@ permission = "write"
producer: "arbitrary-context-provider".into()
}
);
let view =
crate::session_capture::SessionCapture::from_history_entries("startup", replay);
let view = crate::session_capture::SessionCapture::from_history_entries(
"session", "startup", replay,
);
let extract =
crate::feature::builtin::memory_staging_output::render_extract_input(&view);
assert!(!extract.contains("PRIOR_JOURNAL_REFERENCE"));