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

fix: refresh memory surfaces on restore

Browse Source
This commit is contained in:
Hare
2026-10-02 14:32:37 +09:00
parent 504ea00ae3
commit e662e2f465
8 changed files with 458 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/session-store/src/system_item.rs
+12
View File
@@ -157,6 +157,16 @@ pub enum SystemItem {
prompt_provenance: Option<PromptRenderProvenance>,
},
/// Restore-boundary refresh of the current resident Memory projection.
/// The exact rendered body is retained so replay is byte-identical; a
/// refresh with no ready summary explicitly supersedes older resident
/// context without deleting or rewriting history.
ResidentSummaryRefresh {
body: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
prompt_provenance: Option<PromptRenderProvenance>,
},
/// Synthetic note inserted after an interrupted turn before the next
/// user input. `body` is the exact LLM-context text explaining that the
/// previous turn was cut short.
@@ -181,6 +191,7 @@ impl SystemItem {
format!("Ignored legacy procedure item: /{slug}")
}
SystemItem::TaskReminder { body, .. } => body.clone(),
SystemItem::ResidentSummaryRefresh { body, .. } => body.clone(),
SystemItem::Interrupt { body, .. } => body.clone(),
}
}
@@ -202,6 +213,7 @@ impl SystemItem {
SystemItem::LegacyKnowledgeIgnored { .. } => "legacy_knowledge_ignored",
SystemItem::LegacyIgnored { .. } => "legacy_ignored",
SystemItem::TaskReminder { .. } => "task_reminder",
SystemItem::ResidentSummaryRefresh { .. } => "resident_summary_refresh",
SystemItem::Interrupt { .. } => "interrupt",
}
}
crates/tui/src/app.rs
+1
View File
@@ -2584,6 +2584,7 @@ impl App {
}
session_store::SystemItem::FileAttachment { body, .. }
| session_store::SystemItem::SkillActivation { body, .. }
| session_store::SystemItem::ResidentSummaryRefresh { body, .. }
| session_store::SystemItem::Interrupt { body, .. } => {
self.task_store.apply_system_message_text(&body);
self.blocks.push(Block::SystemMessage { text: body });
crates/worker/src/feature/builtin/memory_surface_lifecycle.rs
+194 -37
View File
@@ -358,42 +358,120 @@ fn bounded_editor_input(
system_prompt: &str,
generation: &server_api::SubjektivSurfacePrepareResponse,
) -> Result<(String, Vec<server_api::SubjektivSurfaceMaterial>), HookError> {
let mut materials = generation.materials.clone();
loop {
let input = serde_json::to_string(&serde_json::json!({
"question": EDITOR_QUESTION,
"body_token_budget": generation.body_token_budget,
"input_token_budget": generation.input_token_budget,
"materials": &materials,
}))
.map_err(|error| HookError::new(HookErrorCategory::Internal, error.to_string()))?;
let request = Request::new()
.system(system_prompt)
.user(input.clone())
.tool(submit_llm_tool_definition());
let estimated_tokens = estimated_editor_request_tokens(&request)?;
if estimated_tokens <= generation.input_token_budget {
if materials.is_empty() && !generation.materials.is_empty() {
return Err(HookError::new(
HookErrorCategory::Internal,
format!(
"surface editor fixed prompt and material overhead left no grounded material within the {} token input budget",
generation.input_token_budget
),
));
}
return Ok((input, materials));
}
if materials.pop().is_none() {
return Err(HookError::new(
HookErrorCategory::Internal,
format!(
"surface editor fixed prompt, JSON framing, and tool schema exceed the {} token input budget (estimated {estimated_tokens})",
generation.input_token_budget
),
));
let empty_tokens = editor_request_tokens_for_materials(system_prompt, generation, &[])?;
if empty_tokens > generation.input_token_budget {
return Err(HookError::new(
HookErrorCategory::Internal,
format!(
"surface editor fixed prompt, JSON framing, and tool schema exceed the {} token input budget (estimated {empty_tokens})",
generation.input_token_budget
),
));
}
let mut populated_kinds = Vec::new();
for material in &generation.materials {
if !populated_kinds.contains(&material.kind) {
populated_kinds.push(material.kind.clone());
}
}
let material_token_budget = generation.input_token_budget - empty_tokens;
let per_kind_token_budget = material_token_budget / populated_kinds.len().max(1);
let mut accepted = vec![false; generation.materials.len()];
let mut deferred = Vec::new();
for (index, material) in generation.materials.iter().enumerate() {
let kind_materials = generation
.materials
.iter()
.zip(&accepted)
.filter(|(selected, accepted)| **accepted && selected.kind == material.kind)
.map(|(selected, _)| selected)
.chain(std::iter::once(material))
.collect::<Vec<_>>();
let all_materials = generation
.materials
.iter()
.zip(&accepted)
.filter(|(_, accepted)| **accepted)
.map(|(selected, _)| selected)
.chain(std::iter::once(material))
.collect::<Vec<_>>();
let kind_cost =
editor_request_tokens_for_materials(system_prompt, generation, &kind_materials)?
.saturating_sub(empty_tokens);
let all_tokens =
editor_request_tokens_for_materials(system_prompt, generation, &all_materials)?;
if kind_cost <= per_kind_token_budget && all_tokens <= generation.input_token_budget {
accepted[index] = true;
} else {
deferred.push(index);
}
}
// Equal per-kind shares reserve room in the final normalized request. Any
// unused share is reclaimed in the original canonical round-robin order.
for index in deferred {
let material = &generation.materials[index];
let all_materials = generation
.materials
.iter()
.zip(&accepted)
.filter(|(_, accepted)| **accepted)
.map(|(selected, _)| selected)
.chain(std::iter::once(material))
.collect::<Vec<_>>();
if editor_request_tokens_for_materials(system_prompt, generation, &all_materials)?
<= generation.input_token_budget
{
accepted[index] = true;
}
}
let materials = generation
.materials
.iter()
.zip(accepted)
.filter_map(|(material, accepted)| accepted.then_some(material.clone()))
.collect::<Vec<_>>();
if materials.is_empty() && !generation.materials.is_empty() {
return Err(HookError::new(
HookErrorCategory::Internal,
format!(
"surface editor fixed prompt and material overhead left no grounded material within the {} token input budget",
generation.input_token_budget
),
));
}
let material_refs = materials.iter().collect::<Vec<_>>();
let input = serialize_editor_input(generation, &material_refs)?;
Ok((input, materials))
}
fn editor_request_tokens_for_materials(
system_prompt: &str,
generation: &server_api::SubjektivSurfacePrepareResponse,
materials: &[&server_api::SubjektivSurfaceMaterial],
) -> Result<usize, HookError> {
let input = serialize_editor_input(generation, materials)?;
let request = Request::new()
.system(system_prompt)
.user(input)
.tool(submit_llm_tool_definition());
estimated_editor_request_tokens(&request)
}
fn serialize_editor_input(
generation: &server_api::SubjektivSurfacePrepareResponse,
materials: &[&server_api::SubjektivSurfaceMaterial],
) -> Result<String, HookError> {
serde_json::to_string(&serde_json::json!({
"question": EDITOR_QUESTION,
"body_token_budget": generation.body_token_budget,
"input_token_budget": generation.input_token_budget,
"materials": materials,
}))
.map_err(|error| HookError::new(HookErrorCategory::Internal, error.to_string()))
}
fn estimated_editor_request_tokens(request: &Request) -> Result<usize, HookError> {
@@ -431,11 +509,19 @@ mod tests {
use crate::worker::{WorkspaceClientError, WorkspaceRequest, WorkspaceResponse};
fn material(id: &str, body_bytes: usize) -> server_api::SubjektivSurfaceMaterial {
material_with_body(id, CandidateKind::Constraint, "x".repeat(body_bytes))
}
fn material_with_body(
id: &str,
kind: CandidateKind,
body_md: String,
) -> server_api::SubjektivSurfaceMaterial {
server_api::SubjektivSurfaceMaterial {
memory_id: id.into(),
revision: 1,
kind: CandidateKind::Constraint,
body_md: "x".repeat(body_bytes),
kind,
body_md,
why_useful: "budget fixture".into(),
staleness: None,
}
@@ -703,7 +789,7 @@ permission = "write"
);
let (input, selected) = bounded_editor_input("system policy", &generation).unwrap();
assert_eq!(selected.len(), 1, "deterministically trims the tail");
assert_eq!(selected.len(), 1, "deterministically bounds one category");
assert_eq!(selected[0].memory_id, "memory-1");
assert!(input.contains(EDITOR_QUESTION));
assert!(
@@ -722,6 +808,77 @@ permission = "write"
);
}
#[test]
fn final_editor_budget_preserves_feasible_nonempty_categories() {
let materials = vec![
material_with_body(
"preference-large",
CandidateKind::Preference,
"\\\"".repeat(8_000),
),
material_with_body(
"constraint-small",
CandidateKind::Constraint,
"constraint".into(),
),
material_with_body("decision-small", CandidateKind::Decision, "decision".into()),
material_with_body(
"assumption-small",
CandidateKind::WorkingAssumption,
"assumption".into(),
),
material_with_body(
"question-small",
CandidateKind::OpenQuestion,
"question".into(),
),
material_with_body("lesson-small", CandidateKind::Lesson, "lesson".into()),
];
let mut generation = generation(9_999, materials);
for _ in 0..3 {
let large = editor_request_tokens_for_materials(
"system policy",
&generation,
&[&generation.materials[0]],
)
.unwrap();
let small = editor_request_tokens_for_materials(
"system policy",
&generation,
&generation.materials[1..].iter().collect::<Vec<_>>(),
)
.unwrap();
generation.input_token_budget = large.max(small);
}
let all = editor_request_tokens_for_materials(
"system policy",
&generation,
&generation.materials.iter().collect::<Vec<_>>(),
)
.unwrap();
assert!(all > generation.input_token_budget);
let (input, selected) = bounded_editor_input("system policy", &generation).unwrap();
assert!(
!selected
.iter()
.any(|item| item.kind == CandidateKind::Preference)
);
for kind in [
CandidateKind::Constraint,
CandidateKind::Decision,
CandidateKind::WorkingAssumption,
CandidateKind::OpenQuestion,
CandidateKind::Lesson,
] {
assert!(selected.iter().any(|item| item.kind == kind));
}
assert!(
estimated_editor_request_tokens(&editor_request("system policy", &input)).unwrap()
<= generation.input_token_budget
);
}
#[tokio::test]
async fn actual_editor_request_is_rejected_before_provider_when_over_budget() {
let calls = Arc::new(AtomicUsize::new(0));
crates/worker/src/prompt/catalog.rs
+12
View File
@@ -104,6 +104,7 @@ pub enum WorkerPrompt {
WorkingBoundariesSection,
AgentsMdSection,
ResidentMemorySummarySection,
ResidentMemoryRestoreSection,
WorkerOrchestrationGuidanceSection,
SubWorkerSpawnToolDescription,
}
@@ -126,6 +127,7 @@ impl WorkerPrompt {
Self::WorkingBoundariesSection => "internal.working_boundaries_section",
Self::AgentsMdSection => "internal.agents_md_section",
Self::ResidentMemorySummarySection => "internal.resident_memory_summary_section",
Self::ResidentMemoryRestoreSection => "internal.resident_memory_restore_section",
Self::WorkerOrchestrationGuidanceSection => {
"internal.worker_orchestration_guidance_section"
}
@@ -147,6 +149,7 @@ impl WorkerPrompt {
WorkerPrompt::WorkingBoundariesSection,
WorkerPrompt::AgentsMdSection,
WorkerPrompt::ResidentMemorySummarySection,
WorkerPrompt::ResidentMemoryRestoreSection,
WorkerPrompt::WorkerOrchestrationGuidanceSection,
WorkerPrompt::SubWorkerSpawnToolDescription,
];
@@ -502,6 +505,15 @@ impl PromptCatalog {
single("summary", summary),
)
}
pub fn resident_memory_restore_section(
&self,
summary: Option<&str>,
) -> Result<String, CatalogError> {
self.render(
WorkerPrompt::ResidentMemoryRestoreSection,
single("summary", summary.unwrap_or_default()),
)
}
pub fn worker_orchestration_guidance_section(&self) -> Result<String, CatalogError> {
self.render(
WorkerPrompt::WorkerOrchestrationGuidanceSection,
crates/worker/src/worker.rs
+222 -3
View File
@@ -411,7 +411,9 @@ use crate::ipc::alerter::Alerter;
use crate::ipc::interceptor::WorkerInterceptor;
use crate::ipc::notify_buffer::NotifyBuffer;
use crate::prompt::agents_md::read_agents_md;
use crate::prompt::catalog::{CatalogError, PromptCatalog, WorkspacePromptProjection};
use crate::prompt::catalog::{
CatalogError, PromptCatalog, WorkerPrompt, WorkspacePromptProjection,
};
use crate::prompt::source::PromptCatalogSource;
use crate::prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate};
use crate::runtime::dir;
@@ -2500,6 +2502,8 @@ pub struct Worker<C: LlmClient, St: Store> {
inject_resident_summary: bool,
/// Deferred resident prompt source installed by an enabled Feature.
feature_resident_summary_source: Option<Arc<dyn SystemPromptContributionSource>>,
/// One-shot append-only resident refresh required at a restore boundary.
restore_resident_summary_refresh_pending: bool,
/// Complete system prompt replacement installed by an enabled Feature.
feature_system_prompt_override: Option<String>,
/// Typed user submissions in submit order. K-th entry corresponds to
@@ -2753,6 +2757,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
prompts,
inject_resident_summary: true,
feature_resident_summary_source: None,
restore_resident_summary_refresh_pending: false,
feature_system_prompt_override: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
@@ -3819,6 +3824,62 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
}
async fn ensure_restore_resident_summary_refreshed(&mut self) -> Result<(), WorkerError>
where
St: Clone + 'static,
{
if !self.restore_resident_summary_refresh_pending {
return Ok(());
}
let Some(source) = self.feature_resident_summary_source.clone() else {
self.restore_resident_summary_refresh_pending = false;
return Ok(());
};
let summary = if self.inject_resident_summary {
source.load().await
} else {
None
};
let prompt_catalog = self.prompts.load_full();
let body = prompt_catalog.resident_memory_restore_section(summary.as_deref())?;
let prompt_provenance =
Some(self.prompt_render_provenance(WorkerPrompt::ResidentMemoryRestoreSection.key()));
let item = SystemItem::ResidentSummaryRefresh {
body: body.clone(),
prompt_provenance,
};
let metadata = new_history_metadata(
WorkerHistoryProvenance::BackendInstruction { operation_id: None },
None,
);
self.commit_entry(LogEntry::AnnotatedSystemItem {
ts: segment_log::now_millis(),
entry: session_store::LoggedSystemHistoryEntry {
item,
metadata: metadata.clone(),
},
extensions: Vec::new(),
})?;
let history_entry = HistoryEntry::new(Item::system_message(body), metadata);
let mut annotate = history_annotator(
self.log_writer_handle(),
vec![history_entry.clone()],
self.pending_committed_history.clone(),
);
let (engine, session) = (
self.engine.as_mut().expect("worker present"),
&mut self.session,
);
engine.append_history_with(
session.history_mut(),
std::iter::once(history_entry.item),
&mut annotate,
)?;
session.note_mutation();
self.restore_resident_summary_refresh_pending = false;
Ok(())
}
/// Render the manifest-supplied instruction template exactly once,
/// just before the first LLM turn, append the fixed trailing
/// section (scope summary + optional AGENTS.md), and hand the
@@ -3826,9 +3887,12 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Subsequent invocations are no-ops: the template field is
/// consumed with `Option::take()`, so the materialised value
/// persists across all later turns and compaction.
async fn ensure_system_prompt_materialized(&mut self) -> Result<(), WorkerError> {
async fn ensure_system_prompt_materialized(&mut self) -> Result<(), WorkerError>
where
St: Clone + 'static,
{
let Some(template) = self.system_prompt_template.take() else {
return Ok(());
return self.ensure_restore_resident_summary_refreshed().await;
};
if let Some(rendered) = self.feature_system_prompt_override.take() {
self.engine
@@ -6549,6 +6613,7 @@ where
prompts: common.prompts,
inject_resident_summary: true,
feature_resident_summary_source: None,
restore_resident_summary_refresh_pending: false,
feature_system_prompt_override: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
@@ -6638,6 +6703,7 @@ where
prompts: common.prompts,
inject_resident_summary: true,
feature_resident_summary_source: None,
restore_resident_summary_refresh_pending: false,
feature_system_prompt_override: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
@@ -6762,6 +6828,7 @@ where
prompts: common.prompts,
inject_resident_summary: true,
feature_resident_summary_source: None,
restore_resident_summary_refresh_pending: false,
feature_system_prompt_override: None,
user_segments: Vec::new(),
sink: SegmentLogSink::new(),
@@ -7143,6 +7210,7 @@ where
prompts: common.prompts,
inject_resident_summary: true,
feature_resident_summary_source: None,
restore_resident_summary_refresh_pending: true,
feature_system_prompt_override: None,
user_segments: state.user_segments,
// Seed the mirror with the entries we just replayed so a
@@ -11075,6 +11143,157 @@ mod build_summary_prompt_tests {
assert!(worker.history().is_empty());
}
#[tokio::test]
async fn restored_worker_appends_latest_resident_surface_without_rewriting_history() {
let dir = tempfile::tempdir().unwrap();
let cwd = dir.path().join("workspace");
std::fs::create_dir_all(&cwd).unwrap();
let runtime = dir.path().join("runtime");
std::fs::create_dir_all(&runtime).unwrap();
let _runtime =
crate::runtime::worker_allocation::test_util::RuntimeDirSandbox::new(&runtime);
let store = session_store::CombinedStore::new(
session_store::FsStore::new(dir.path().join("sessions")).unwrap(),
session_store::FsWorkerStore::new(dir.path().join("workers")).unwrap(),
);
let manifest = WorkerManifest::from_toml(&format!(
r#"
[worker]
name = "restore-resident-{}"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
instruction = "default"
[[scope.allow]]
target = "{}"
permission = "write"
"#,
uuid::Uuid::now_v7(),
cwd.display()
))
.unwrap();
let authority = WorkerFilesystemAuthority::local(cwd.clone(), cwd.clone());
let scope = Scope::writable(&cwd).unwrap();
let mut original = Worker::new(
manifest.clone(),
Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
store.clone(),
WorkerWorkspaceContext::no_workspace(),
authority.clone(),
scope,
)
.await
.unwrap();
original.set_system_prompt_template(
SystemPromptTemplate::parse(
"default",
crate::prompt::source::PromptCatalogSource::builtins_only(),
)
.unwrap(),
);
let (old_source, _) =
test_system_prompt_contribution_source(Some("old subject surface".into()));
original.install_system_prompt_contribution(Some(old_source), None);
original.materialize_durable_session_head().await.unwrap();
let session_id = original.session_id();
let segment_id = original.segment_id();
assert!(
original
.engine()
.get_system_prompt()
.unwrap()
.contains("old subject surface")
);
drop(original);
let mut restored =
Worker::<Box<dyn LlmClient>, _>::restore_from_manifest_with_context_and_model_client(
session_id,
segment_id,
manifest.clone(),
store.clone(),
crate::prompt::source::PromptCatalogSource::builtins_only(),
WorkerWorkspaceContext::no_workspace(),
authority.clone(),
Some(Box::new(NoopClient)),
)
.await
.unwrap();
let history_before = restored.history().to_vec();
let (current_source, current_loads) =
test_system_prompt_contribution_source(Some("latest subject surface".into()));
restored.install_system_prompt_contribution(Some(current_source), None);
restored.materialize_durable_session_head().await.unwrap();
assert_eq!(current_loads.load(Ordering::SeqCst), 1);
assert_eq!(&restored.history()[..history_before.len()], history_before);
let refresh = restored
.history()
.last()
.unwrap()
.as_text()
.unwrap()
.to_string();
assert!(refresh.contains("latest subject surface"));
assert!(refresh.contains("supersedes every resident memory summary"));
assert!(!refresh.contains("old subject surface"));
let refreshed_len = restored.history().len();
restored.materialize_durable_session_head().await.unwrap();
assert_eq!(restored.history().len(), refreshed_len);
drop(restored);
let mut stale_restore =
Worker::<Box<dyn LlmClient>, _>::restore_from_manifest_with_context_and_model_client(
session_id,
segment_id,
manifest,
store.clone(),
crate::prompt::source::PromptCatalogSource::builtins_only(),
WorkerWorkspaceContext::no_workspace(),
authority,
Some(Box::new(NoopClient)),
)
.await
.unwrap();
let stale_history_before = stale_restore.history().to_vec();
let (stale_source, stale_loads) = test_system_prompt_contribution_source(None);
stale_restore.install_system_prompt_contribution(Some(stale_source), None);
stale_restore
.materialize_durable_session_head()
.await
.unwrap();
assert_eq!(stale_loads.load(Ordering::SeqCst), 1);
assert_eq!(
&stale_restore.history()[..stale_history_before.len()],
stale_history_before
);
let tombstone = stale_restore
.history()
.last()
.unwrap()
.as_text()
.unwrap()
.to_string();
assert!(tombstone.contains("No current ready resident memory surface"));
assert!(!tombstone.contains("latest subject surface"));
let entries = store.read_all(session_id, segment_id).unwrap();
assert!(matches!(
entries.last(),
Some(LogEntry::AnnotatedSystemItem {
entry: session_store::LoggedSystemHistoryEntry {
item: SystemItem::ResidentSummaryRefresh { .. },
..
},
..
})
));
}
#[tokio::test]
async fn memory_consolidation_prompt_uses_bound_workspace_language() {
let dir = tempfile::tempdir().unwrap();