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

feat: integrate bounded subject memory surfaces

Browse Source
This commit is contained in:
Hare
2026-10-02 15:11:12 +09:00
parent aad200258c dfe837655a
commit 43e3b42055
20 changed files with 3314 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+45
View File
@@ -767,6 +767,17 @@ impl WorkerController {
// the actual registered set instead of a hand-maintained mirror.
worker.engine().tool_server_handle().flush_pending();
// Runtime-owned fresh Workers persist their initial model-visible head
// before controller exposure. Feature installation above must happen first
// so the resident contribution is loaded into that durable prompt. Restored
// Workers retain their established first-turn refresh lifecycle.
if runtime_managed && worker.needs_initial_session_head_materialization() {
worker
.materialize_durable_session_head()
.await
.map_err(|error| std::io::Error::other(error.to_string()))?;
}
// === 4. Initial runtime files + WorkerSharedState + WorkerHandle +
// SocketServer ===
let manifest_toml = toml::to_string_pretty(worker.manifest()).unwrap_or_default();
@@ -1313,6 +1324,16 @@ where
{
feature_registry.add_module(memory_lifecycle);
}
if let Some(surface_lifecycle) = crate::feature::builtin::memory_surface_lifecycle::SubjektivSurfaceLifecycleFeature::from_manifest(
worker.manifest_lifecycle_features_enabled(),
worker.workspace_client_handle(),
spawner_manifest.clone(),
worker.llm_client_handle(),
prompts.clone(),
spawner_workspace_context.clone(),
)? {
feature_registry.add_module(surface_lifecycle);
}
if let Some(subjektiv_lifecycle) =
crate::feature::builtin::memory_lifecycle::SubjektivLifecycleFeature::from_resolved_config(
worker.manifest_lifecycle_features_enabled(),
@@ -3455,6 +3476,30 @@ mod tests {
use tempfile::TempDir;
use tokio::net::UnixListener;
#[test]
fn runtime_session_head_is_materialized_after_feature_installation_before_exposure() {
let production = include_str!("controller.rs")
.split_once("#[cfg(test)]\nmod tests")
.map(|(production, _)| production)
.expect("controller test module marker");
let startup = production
.split_once("async fn spawn_initialized")
.map(|(_, startup)| startup)
.expect("controller startup implementation");
let feature_install = startup
.find("let fs_for_view = register_worker_tools(")
.expect("feature installation call");
let materialize = startup
.find(".materialize_durable_session_head()")
.expect("runtime session-head materialization");
let exposure = startup
.find("let handle = WorkerHandle")
.expect("controller handle exposure");
assert!(feature_install < materialize);
assert!(materialize < exposure);
}
#[test]
fn automatic_memory_extraction_destinations_are_mutually_exclusive() {
let mut memory = manifest::MemoryFeatureProfileConfig::default();
crates/worker/src/feature/builtin.rs
+2
View File
@@ -11,6 +11,8 @@ pub mod manage_worker;
pub mod memory;
pub(crate) mod memory_lifecycle;
pub mod memory_staging_output;
pub(crate) mod memory_surface_lifecycle;
pub(crate) mod memory_surface_output;
pub mod merge_request;
pub mod objective;
pub mod orchestration;
crates/worker/src/feature/builtin/memory.rs
+103 -28
View File
@@ -23,8 +23,8 @@ use crate::feature::{
ToolDeclaration,
};
use crate::worker::{
SystemPromptContributionSource, WorkspaceClient, WorkspaceClientError, WorkspaceRequest,
WorkspaceRequestMethod,
SystemPromptContribution, SystemPromptContributionSource, WorkspaceClient,
WorkspaceClientError, WorkspaceRequest, WorkspaceRequestMethod,
};
#[derive(Clone, Debug)]
@@ -93,6 +93,87 @@ impl dyn WorkspaceClient + '_ {
) -> Result<MemoryConsolidationOutput, WorkspaceMemoryBackendError> {
execute_memory_consolidation(self, operation, true).await
}
pub async fn prepare_subjektiv_memory_surface(
&self,
) -> Result<server_api::SubjektivSurfacePrepareResponse, WorkspaceMemoryBackendError> {
match execute_subjektiv_memory_operation(
self,
server_api::SubjektivMemoryBackendOperation::PrepareSurface(
server_api::SubjektivSurfacePrepareRequest {},
),
)
.await?
{
server_api::SubjektivMemoryBackendResponse::SurfacePrepared(output) => Ok(output),
other => Err(WorkspaceMemoryBackendError::Backend(format!(
"unexpected surface preparation response: {other:?}"
))),
}
}
pub async fn publish_subjektiv_memory_surface(
&self,
input: server_api::SubjektivSurfacePublishRequest,
) -> Result<server_api::SubjektivSurfacePublishResponse, WorkspaceMemoryBackendError> {
match execute_subjektiv_memory_operation(
self,
server_api::SubjektivMemoryBackendOperation::PublishSurface(input),
)
.await?
{
server_api::SubjektivMemoryBackendResponse::SurfacePublished(output) => Ok(output),
other => Err(WorkspaceMemoryBackendError::Backend(format!(
"unexpected surface publication response: {other:?}"
))),
}
}
pub async fn fail_subjektiv_memory_surface(
&self,
input: server_api::SubjektivSurfaceFailureRequest,
) -> Result<server_api::SubjektivSurfaceFailureResponse, WorkspaceMemoryBackendError> {
match execute_subjektiv_memory_operation(
self,
server_api::SubjektivMemoryBackendOperation::FailSurface(input),
)
.await?
{
server_api::SubjektivMemoryBackendResponse::SurfaceFailed(output) => Ok(output),
other => Err(WorkspaceMemoryBackendError::Backend(format!(
"unexpected surface failure response: {other:?}"
))),
}
}
}
async fn execute_subjektiv_memory_operation(
client: &dyn WorkspaceClient,
operation: server_api::SubjektivMemoryBackendOperation,
) -> Result<server_api::SubjektivMemoryBackendResponse, WorkspaceMemoryBackendError> {
let workspace_id =
client
.workspace_id()
.ok_or_else(|| WorkspaceMemoryBackendError::Unavailable {
reason: format!(
"workspace client kind `{}` has no workspace id",
client.kind()
),
})?;
let response = client.execute(WorkspaceRequest::json(
WorkspaceRequestMethod::Post,
format!("/api/w/{workspace_id}/subjektiv/memory"),
serde_json::to_string(&server_api::SubjektivMemoryBackendRequest { operation })?,
))?;
let status = reqwest::StatusCode::from_u16(response.status)
.unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR);
if !response.is_success() {
return Err(WorkspaceMemoryBackendError::Http {
status,
body: response.body,
});
}
Ok(serde_json::from_str(&response.body)?)
}
async fn execute_memory_backend(
@@ -571,7 +652,7 @@ struct WorkspaceResidentSummarySource {
#[async_trait]
impl SystemPromptContributionSource for WorkspaceResidentSummarySource {
async fn load(&self) -> Option<String> {
async fn load(&self) -> SystemPromptContribution {
match self
.client
.execute_memory_backend_operation(
@@ -581,14 +662,22 @@ impl SystemPromptContributionSource for WorkspaceResidentSummarySource {
)
.await
{
Ok(memory::backend::MemoryBackendOperationResult::ToolOutput(output)) => output.content,
Ok(memory::backend::MemoryBackendOperationResult::ResidentSummary(output))
if output.availability
== memory::backend::MemoryResidentSummaryAvailability::Ready =>
{
SystemPromptContribution::Ready(output.content.unwrap_or_default())
}
Ok(memory::backend::MemoryBackendOperationResult::ResidentSummary(_)) => {
SystemPromptContribution::Unavailable
}
Ok(other) => {
tracing::debug!(?other, "unexpected resident Memory Backend result");
None
SystemPromptContribution::Unavailable
}
Err(error) => {
tracing::debug!(%error, "resident Memory summary unavailable");
None
SystemPromptContribution::Unavailable
}
}
}
@@ -762,8 +851,8 @@ mod tests {
let body = serde_json::json!({
"status": "ok",
"result": {
"kind": "tool_output",
"summary": "resident Memory summary collected",
"kind": "resident_summary",
"availability": "ready",
"content": content,
}
})
@@ -895,12 +984,8 @@ mod tests {
.unwrap()
.unwrap();
assert_eq!(
plan.resident_summary_source
.unwrap()
.load()
.await
.as_deref(),
Some("# Durable Memory")
plan.resident_summary_source.unwrap().load().await,
SystemPromptContribution::Ready("# Durable Memory".into())
);
}
@@ -935,22 +1020,12 @@ mod tests {
.unwrap();
assert_eq!(
first
.resident_summary_source
.unwrap()
.load()
.await
.as_deref(),
Some("first resident summary")
first.resident_summary_source.unwrap().load().await,
SystemPromptContribution::Ready("first resident summary".into())
);
assert_eq!(
restored
.resident_summary_source
.unwrap()
.load()
.await
.as_deref(),
Some("updated resident summary")
restored.resident_summary_source.unwrap().load().await,
SystemPromptContribution::Ready("updated resident summary".into())
);
}
crates/worker/src/feature/builtin/memory_surface_lifecycle.rs
+973
View File
@@ -0,0 +1,973 @@
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use agen::llm_client::event::Event as LlmEvent;
use agen::llm_client::{ClientError, LlmClient, Request};
use arc_swap::ArcSwap;
use async_trait::async_trait;
use futures_util::Stream;
use crate::PromptCatalog;
use crate::Scope;
use crate::feature::background::{
BackgroundTaskCancellation, BackgroundTaskContext, BackgroundTaskSpec, BackgroundTaskTrigger,
FeatureBackgroundTask,
};
use crate::feature::builtin::memory::WorkspaceMemoryBackendError;
use crate::feature::builtin::memory_surface_output::{
MemorySurfaceOutputFeature, MemorySurfaceOutputState, submit_llm_tool_definition,
};
use crate::feature::{
BackgroundTaskDeclaration, FeatureDescriptor, FeatureInstallContext, FeatureInstallError,
FeatureModule, FeatureRegistryBuilder,
};
use crate::hook::{HookError, HookErrorCategory};
use crate::internal_worker::{
InternalWorkerAuthority, InternalWorkerIdentity, InternalWorkerSpec,
run_internal_worker_with_cancel_sender,
};
use crate::worker::{WorkerFilesystemAuthority, WorkerWorkspaceContext, WorkspaceClient};
use manifest::WorkerManifest;
const TASK_NAME: &str = "subjektiv-memory-surface";
const TASK_TIMEOUT: Duration = Duration::from_secs(300);
const MAX_GENERATION_ATTEMPTS: usize = 2;
const EDITOR_MAX_TURNS: u32 = 3;
const EDITOR_QUESTION: &str = "この主体が次の作業を始める際、毎回思い出しておくべきことは何か。継続的な制約、判断の前提、未解決事項、再発を避けたい教訓を、指定予算内でまとめる。";
#[derive(Clone)]
pub(crate) struct SubjektivSurfaceLifecycleFeature {
task: SubjektivSurfaceLifecycleTask,
}
#[derive(Clone)]
struct SubjektivSurfaceLifecycleTask {
workspace_client: Arc<dyn WorkspaceClient>,
manifest: WorkerManifest,
client: Box<dyn LlmClient>,
prompts: Arc<ArcSwap<PromptCatalog>>,
workspace_context: WorkerWorkspaceContext,
}
const INPUT_BUDGET_ERROR_MARKER: &str = "surface_editor_input_budget_exceeded";
struct InputBudgetLlmClient {
inner: Box<dyn LlmClient>,
input_token_budget: usize,
}
impl Clone for InputBudgetLlmClient {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone_boxed(),
input_token_budget: self.input_token_budget,
}
}
}
#[async_trait]
impl LlmClient for InputBudgetLlmClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
let estimated = estimated_editor_request_tokens(&request)
.map_err(|error| ClientError::Config(error.to_string()))?;
if estimated > self.input_token_budget {
return Err(ClientError::Config(format!(
"{INPUT_BUDGET_ERROR_MARKER}: request estimate {estimated} exceeds {}",
self.input_token_budget
)));
}
self.inner.stream(request).await
}
}
impl SubjektivSurfaceLifecycleFeature {
pub(crate) fn from_manifest(
lifecycle_enabled: bool,
workspace_client: Arc<dyn WorkspaceClient>,
manifest: WorkerManifest,
client: Box<dyn LlmClient>,
prompts: Arc<ArcSwap<PromptCatalog>>,
workspace_context: WorkerWorkspaceContext,
) -> std::io::Result<Option<Self>> {
let dedicated = manifest.profile.as_ref().is_some_and(|snapshot| {
matches!(
&snapshot.source,
manifest::ProfileSource::Registry {
source: manifest::ProfileRegistrySource::Builtin,
name,
..
} if name == "subjektiv-memory-consolidation"
)
});
if !lifecycle_enabled || !dedicated {
return Ok(None);
}
if !workspace_client.is_available() || workspace_client.workspace_id().is_none() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"subjektiv surface generation requires Backend Workspace API authority",
));
}
Ok(Some(Self {
task: SubjektivSurfaceLifecycleTask {
workspace_client,
manifest,
client,
prompts,
workspace_context,
},
}))
}
}
impl FeatureModule for SubjektivSurfaceLifecycleFeature {
fn descriptor(&self) -> FeatureDescriptor {
FeatureDescriptor::builtin("subjektiv-memory-surface-lifecycle", "subjektiv Memory Surface Lifecycle")
.with_description("Rebuilds one bounded subject Memory surface after each committed consolidation run.")
.with_background_task(BackgroundTaskDeclaration::worker_managed(
TASK_NAME,
"Generate and publish a current subject Memory surface from bounded confirmed revisions.",
))
}
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
let declaration = BackgroundTaskDeclaration::worker_managed(
TASK_NAME,
"Generate and publish a current subject Memory surface from bounded confirmed revisions.",
);
let mut spec = BackgroundTaskSpec::single_flight(declaration, TASK_TIMEOUT);
spec.trigger = BackgroundTaskTrigger::RunCommitted;
context.background_tasks().register(spec, self.task.clone())
}
}
struct SurfaceEditorFailure {
error: HookError,
reason_code: &'static str,
}
impl SurfaceEditorFailure {
fn input_budget(error: HookError) -> Self {
Self {
error,
reason_code: "input_budget_exhausted",
}
}
fn editor(error: HookError) -> Self {
Self {
error,
reason_code: "editor_failed",
}
}
}
#[async_trait]
impl FeatureBackgroundTask for SubjektivSurfaceLifecycleTask {
async fn run(
&self,
context: BackgroundTaskContext,
cancellation: BackgroundTaskCancellation,
) -> Result<(), HookError> {
for attempt in 0..MAX_GENERATION_ATTEMPTS {
context.generation_fence.ensure_current()?;
if cancellation.is_cancelled() {
return Ok(());
}
let generation = self
.workspace_client
.prepare_subjektiv_memory_surface()
.await
.map_err(surface_hook_error)?;
if generation.materials.is_empty() && generation.active_memory_count > 0 {
self.record_failure(&generation.generation_id, "input_budget_exhausted")
.await;
tracing::warn!(
active_memory_count = generation.active_memory_count,
"subject Memory surface materials could not fit the bounded input"
);
return Ok(());
}
let points = if generation.materials.is_empty() {
Vec::new()
} else {
match self.edit_surface(&generation, cancellation.clone()).await {
Ok(points) => points,
Err(failure) => {
self.record_failure(&generation.generation_id, failure.reason_code)
.await;
tracing::warn!(error = %failure.error, reason_code = failure.reason_code, "subject Memory surface editor failed");
return Ok(());
}
}
};
context.generation_fence.ensure_current()?;
if cancellation.is_cancelled() {
return Ok(());
}
let publish = self
.workspace_client
.publish_subjektiv_memory_surface(server_api::SubjektivSurfacePublishRequest {
generation_id: generation.generation_id.clone(),
points,
})
.await;
match publish {
Ok(output) => {
tracing::debug!(
snapshot_id = output.snapshot_id,
store_revision = output.built_from_store_revision,
empty = output.empty,
"published subject Memory surface"
);
return Ok(());
}
Err(WorkspaceMemoryBackendError::Http { status, .. })
if status == reqwest::StatusCode::CONFLICT
&& attempt + 1 < MAX_GENERATION_ATTEMPTS =>
{
// Confirmed Memory moved while editing. Discard the stale
// output and rebuild once from a fresh bounded generation.
continue;
}
Err(error) => {
self.record_failure(&generation.generation_id, "publish_failed")
.await;
tracing::warn!(%error, "subject Memory surface publication failed");
return Ok(());
}
}
}
Ok(())
}
}
impl SubjektivSurfaceLifecycleTask {
async fn edit_surface(
&self,
generation: &server_api::SubjektivSurfacePrepareResponse,
cancellation: BackgroundTaskCancellation,
) -> Result<Vec<server_api::SubjektivSurfacePoint>, SurfaceEditorFailure> {
let language = self
.manifest
.feature
.memory
.workspace_settings()
.map(|settings| settings.language)
.ok_or_else(|| {
SurfaceEditorFailure::editor(HookError::new(
HookErrorCategory::Internal,
"surface editor requires bound Workspace Memory settings",
))
})?;
let system_prompt = self
.prompts
.load_full()
.subjektiv_memory_surface_system(&language)
.map_err(|error| {
SurfaceEditorFailure::editor(HookError::new(
HookErrorCategory::Internal,
error.to_string(),
))
})?;
let (input, editor_materials) = bounded_editor_input(system_prompt.as_str(), generation)
.map_err(SurfaceEditorFailure::input_budget)?;
let output_state =
MemorySurfaceOutputState::new(editor_materials, generation.body_token_budget);
let features = FeatureRegistryBuilder::new()
.with_module(MemorySurfaceOutputFeature::new(output_state.clone()));
let cancellation_observer = cancellation.clone();
let cancel_observer = move |sender: tokio::sync::mpsc::Sender<()>| {
tokio::spawn(async move {
cancellation_observer.cancelled().await;
let _ = sender.send(()).await;
});
};
let result = run_internal_worker_with_cancel_sender(
InternalWorkerSpec {
identity: InternalWorkerIdentity {
kind: "subjektiv-memory-surface",
run_id: uuid::Uuid::now_v7(),
},
manifest: self.manifest.clone(),
client: Box::new(InputBudgetLlmClient {
inner: self.client.clone_boxed(),
input_token_budget: generation.input_token_budget,
}),
system_prompt,
input,
cache_key: None,
max_turns: Some(EDITOR_MAX_TURNS),
engine_configurator: None,
features,
required_tools: &["SubmitMemorySurface"],
authority: InternalWorkerAuthority {
workspace: self.workspace_context.clone(),
filesystem: WorkerFilesystemAuthority::None,
scope: Scope::empty(),
workdir_session: None,
},
},
cancel_observer,
)
.await;
if let Err(error) = result {
let message = format!("surface editor Internal Worker failed: {}", error.source);
let error = HookError::new(HookErrorCategory::Internal, message.clone());
return Err(if message.contains(INPUT_BUDGET_ERROR_MARKER) {
SurfaceEditorFailure::input_budget(error)
} else {
SurfaceEditorFailure::editor(error)
});
}
output_state
.submitted()
.ok_or_else(|| {
HookError::new(
HookErrorCategory::Internal,
"surface editor finished without valid structured output",
)
})
.map_err(SurfaceEditorFailure::editor)
}
async fn record_failure(&self, generation_id: &str, reason_code: &str) {
if let Err(error) = self
.workspace_client
.fail_subjektiv_memory_surface(server_api::SubjektivSurfaceFailureRequest {
generation_id: generation_id.to_string(),
reason_code: reason_code.to_string(),
})
.await
{
tracing::debug!(%error, "could not record subject Memory surface failure");
}
}
}
fn bounded_editor_input(
system_prompt: &str,
generation: &server_api::SubjektivSurfacePrepareResponse,
) -> Result<(String, Vec<server_api::SubjektivSurfaceMaterial>), HookError> {
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> {
// Serialize the actual prompt-bearing Agen request fields. This is the
// stable provider-independent boundary; provider-specific HTTP envelopes
// are intentionally outside the configured editor-input budget.
let request = serde_json::to_vec(&serde_json::json!({
"system_prompt": &request.system_prompt,
"items": &request.items,
"tools": &request.tools,
}))
.map_err(|error| HookError::new(HookErrorCategory::Internal, error.to_string()))?;
Ok(request.len().saturating_add(3) / 4)
}
fn surface_hook_error(error: WorkspaceMemoryBackendError) -> HookError {
HookError::new(HookErrorCategory::Internal, error.to_string())
}
#[cfg(test)]
mod tests {
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use agen::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
use agen::llm_client::{ClientError, Request};
use futures::Stream;
use memory::extract::CandidateKind;
use super::*;
use crate::feature::FeatureId;
use crate::feature::background::FeatureBackgroundTaskRegistryBuilder;
use crate::hook::HookInvocationContext;
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,
body_md,
why_useful: "budget fixture".into(),
staleness: None,
}
}
fn generation(
input_token_budget: usize,
materials: Vec<server_api::SubjektivSurfaceMaterial>,
) -> server_api::SubjektivSurfacePrepareResponse {
server_api::SubjektivSurfacePrepareResponse {
generation_id: "generation-1".into(),
store_revision: 1,
active_memory_count: materials.len(),
materials,
body_token_budget: 1_024,
input_token_budget,
per_kind_limit: 8,
total_material_limit: 24,
}
}
fn editor_request(system_prompt: &str, input: &str) -> Request {
Request::new()
.system(system_prompt)
.user(input)
.tool(submit_llm_tool_definition())
}
#[derive(Debug)]
struct SurfaceWorkspaceClient {
conflict_count: usize,
input_token_budget: usize,
prepare_calls: AtomicUsize,
publish_calls: AtomicUsize,
failure_reasons: Mutex<Vec<String>>,
}
impl SurfaceWorkspaceClient {
fn new(conflict_count: usize) -> Self {
Self {
conflict_count,
input_token_budget: 12_000,
prepare_calls: AtomicUsize::new(0),
publish_calls: AtomicUsize::new(0),
failure_reasons: Mutex::new(Vec::new()),
}
}
fn with_input_token_budget(input_token_budget: usize) -> Self {
Self {
input_token_budget,
..Self::new(0)
}
}
fn response(value: server_api::SubjektivMemoryBackendResponse) -> WorkspaceResponse {
WorkspaceResponse {
status: 200,
body: serde_json::to_string(&value).unwrap(),
}
}
}
impl WorkspaceClient for SurfaceWorkspaceClient {
fn workspace_id(&self) -> Option<&str> {
Some("workspace-1")
}
fn kind(&self) -> &str {
"surface-lifecycle-test"
}
fn is_available(&self) -> bool {
true
}
fn execute(
&self,
request: WorkspaceRequest,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
let request: server_api::SubjektivMemoryBackendRequest =
serde_json::from_str(request.body.as_deref().unwrap_or_default())
.map_err(|error| WorkspaceClientError::Request(error.to_string()))?;
match request.operation {
server_api::SubjektivMemoryBackendOperation::PrepareSurface(_) => {
let call = self.prepare_calls.fetch_add(1, Ordering::SeqCst) + 1;
let mut prepared =
generation(self.input_token_budget, vec![material("memory-1", 32)]);
prepared.generation_id = format!("generation-{call}");
prepared.store_revision = call as u64;
Ok(Self::response(
server_api::SubjektivMemoryBackendResponse::SurfacePrepared(prepared),
))
}
server_api::SubjektivMemoryBackendOperation::PublishSurface(input) => {
assert_eq!(input.points.len(), 1);
let call = self.publish_calls.fetch_add(1, Ordering::SeqCst) + 1;
if call <= self.conflict_count {
return Ok(WorkspaceResponse {
status: 409,
body: "store generation moved".into(),
});
}
Ok(Self::response(
server_api::SubjektivMemoryBackendResponse::SurfacePublished(
server_api::SubjektivSurfacePublishResponse {
snapshot_id: format!("surface-{call}"),
built_from_store_revision: call as u64,
empty: false,
},
),
))
}
server_api::SubjektivMemoryBackendOperation::FailSurface(input) => {
self.failure_reasons.lock().unwrap().push(input.reason_code);
Ok(Self::response(
server_api::SubjektivMemoryBackendResponse::SurfaceFailed(
server_api::SubjektivSurfaceFailureResponse {
store_revision: self.prepare_calls.load(Ordering::SeqCst) as u64,
status: "failed".into(),
},
),
))
}
_ => Err(WorkspaceClientError::Unavailable(
"surface lifecycle test only accepts surface operations".into(),
)),
}
}
}
#[derive(Clone)]
struct SurfaceEditorClient {
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl LlmClient for SurfaceEditorClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
let call = self.calls.fetch_add(1, Ordering::SeqCst);
let events = if call % 2 == 0 {
vec![
LlmEvent::tool_use_start(0, format!("submit-{call}"), "SubmitMemorySurface"),
LlmEvent::tool_input_delta(
0,
serde_json::json!({
"points": [{
"body_md": "- Keep the confirmed constraint.",
"memory_refs": [{"memory_id": "memory-1", "revision": 1}]
}]
})
.to_string(),
),
LlmEvent::tool_use_stop(0),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
} else {
vec![
LlmEvent::text_block_start(0),
LlmEvent::text_delta(0, "submitted"),
LlmEvent::text_block_stop(0, None),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
};
Ok(Box::pin(futures::stream::iter(events.into_iter().map(Ok))))
}
}
fn test_manifest() -> WorkerManifest {
let mut manifest = WorkerManifest::from_toml(
r#"
[worker]
name = "surface-lifecycle-test"
scope = "main"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
[[scope.allow]]
target = "/surface-lifecycle-test"
permission = "write"
"#,
)
.unwrap();
manifest.feature.memory.profile.enabled = true;
manifest
.feature
.memory
.bind_workspace_settings(manifest::WorkspaceMemorySettingsSnapshot {
workspace_id: "workspace-1".into(),
settings_revision: 1,
language: "English".into(),
})
.unwrap();
manifest
}
fn test_task(
workspace_client: Arc<dyn WorkspaceClient>,
calls: Arc<AtomicUsize>,
) -> SubjektivSurfaceLifecycleTask {
SubjektivSurfaceLifecycleTask {
workspace_client,
manifest: test_manifest(),
client: Box::new(SurfaceEditorClient { calls }),
prompts: Arc::new(ArcSwap::from(PromptCatalog::builtins_only().unwrap())),
workspace_context: WorkerWorkspaceContext::no_workspace(),
}
}
async fn run_task(task: SubjektivSurfaceLifecycleTask) {
let declaration = BackgroundTaskDeclaration::worker_managed(
TASK_NAME,
"surface lifecycle integration test",
);
let mut spec = BackgroundTaskSpec::single_flight(declaration, Duration::from_secs(5));
spec.trigger = BackgroundTaskTrigger::RunCommitted;
let mut builder = FeatureBackgroundTaskRegistryBuilder::default();
builder
.register(FeatureId::builtin("surface-lifecycle-test"), spec, task)
.unwrap();
let registry = builder.build();
registry
.start_run_committed(HookInvocationContext {
worker_id: "worker-1".into(),
session_id: "session-1".into(),
session_revision: 1,
run_id: Some("run-1".into()),
..Default::default()
})
.unwrap();
tokio::time::timeout(Duration::from_secs(5), async {
loop {
if !registry.diagnostics().is_empty() {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("surface lifecycle task should finish");
registry.shutdown().await.unwrap();
}
#[test]
fn editor_input_budget_counts_prompt_json_question_and_tool_schema() {
let generation = generation(
4_000,
vec![material("memory-1", 9_000), material("memory-2", 9_000)],
);
let (input, selected) = bounded_editor_input("system policy", &generation).unwrap();
assert_eq!(selected.len(), 1, "deterministically bounds one category");
assert_eq!(selected[0].memory_id, "memory-1");
assert!(input.contains(EDITOR_QUESTION));
assert!(
estimated_editor_request_tokens(&editor_request("system policy", &input)).unwrap()
<= generation.input_token_budget
);
let materials_only_tokens = (serde_json::to_vec(&selected)
.unwrap()
.len()
.saturating_add(3))
/ 4;
assert!(
estimated_editor_request_tokens(&editor_request("system policy", &input)).unwrap()
> materials_only_tokens,
"full accounting must include framing and the SubmitMemorySurface schema"
);
}
#[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));
let client = InputBudgetLlmClient {
inner: Box::new(SurfaceEditorClient {
calls: calls.clone(),
}),
input_token_budget: 32,
};
let request = editor_request("system policy", &"x".repeat(1_000));
let error = match client.stream(request).await {
Err(error) => error,
Ok(_) => panic!("over-budget request reached the provider"),
};
assert!(error.to_string().contains(INPUT_BUDGET_ERROR_MARKER));
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn correction_turns_remain_subject_to_editor_input_budget() {
let calls = Arc::new(AtomicUsize::new(0));
let client = InputBudgetLlmClient {
inner: Box::new(SurfaceEditorClient {
calls: calls.clone(),
}),
input_token_budget: 2_000,
};
let _stream = client
.stream(editor_request("system policy", "small initial request"))
.await
.unwrap();
let request =
editor_request("system policy", "small initial request").assistant("x".repeat(12_000));
let error = match client.stream(request).await {
Err(error) => error,
Ok(_) => panic!("over-budget request reached the provider"),
};
assert!(error.to_string().contains(INPUT_BUDGET_ERROR_MARKER));
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn lifecycle_rebuilds_once_after_a_store_generation_conflict() {
let workspace = Arc::new(SurfaceWorkspaceClient::new(1));
let editor_calls = Arc::new(AtomicUsize::new(0));
run_task(test_task(workspace.clone(), editor_calls.clone())).await;
assert_eq!(workspace.prepare_calls.load(Ordering::SeqCst), 2);
assert_eq!(workspace.publish_calls.load(Ordering::SeqCst), 2);
assert!(workspace.failure_reasons.lock().unwrap().is_empty());
assert_eq!(editor_calls.load(Ordering::SeqCst), 4);
}
#[tokio::test]
async fn lifecycle_bounds_conflict_retries_and_records_terminal_failure() {
let workspace = Arc::new(SurfaceWorkspaceClient::new(MAX_GENERATION_ATTEMPTS));
let editor_calls = Arc::new(AtomicUsize::new(0));
run_task(test_task(workspace.clone(), editor_calls.clone())).await;
assert_eq!(workspace.prepare_calls.load(Ordering::SeqCst), 2);
assert_eq!(workspace.publish_calls.load(Ordering::SeqCst), 2);
assert_eq!(editor_calls.load(Ordering::SeqCst), 4);
assert_eq!(
workspace.failure_reasons.lock().unwrap().as_slice(),
["publish_failed"]
);
}
#[tokio::test]
async fn lifecycle_records_input_budget_exhaustion_without_provider_or_publish() {
let workspace = Arc::new(SurfaceWorkspaceClient::with_input_token_budget(1));
let editor_calls = Arc::new(AtomicUsize::new(0));
run_task(test_task(workspace.clone(), editor_calls.clone())).await;
assert_eq!(workspace.prepare_calls.load(Ordering::SeqCst), 1);
assert_eq!(workspace.publish_calls.load(Ordering::SeqCst), 0);
assert_eq!(editor_calls.load(Ordering::SeqCst), 0);
assert_eq!(
workspace.failure_reasons.lock().unwrap().as_slice(),
["input_budget_exhausted"]
);
}
#[test]
fn editor_input_budget_never_relabels_trimmed_nonempty_materials_as_empty() {
let generation = generation(1, vec![material("memory-1", 32)]);
let error = bounded_editor_input("system policy", &generation).unwrap_err();
assert!(error.to_string().contains("input budget"));
}
}
crates/worker/src/feature/builtin/memory_surface_output.rs
+256
View File
@@ -0,0 +1,256 @@
use std::collections::HashSet;
use std::sync::{Arc, Mutex};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait;
use schemars::JsonSchema;
use serde::Deserialize;
use crate::feature::{
FeatureDescriptor, FeatureInstallContext, FeatureInstallError, FeatureModule, ToolContribution,
ToolDeclaration,
};
const SUBMIT_TOOL: &str = "SubmitMemorySurface";
const SUBMIT_DESCRIPTION: &str = "Submit a bounded Markdown Memory surface whose every point cites exact refs from the supplied confirmed-Memory materials. Reference validation proves existence and scope, not semantic correctness.";
#[derive(Clone)]
pub(crate) struct MemorySurfaceOutputState {
materials: Arc<Vec<server_api::SubjektivSurfaceMaterial>>,
body_token_budget: usize,
submitted: Arc<Mutex<Option<Vec<server_api::SubjektivSurfacePoint>>>>,
}
impl MemorySurfaceOutputState {
pub(crate) fn new(
materials: Vec<server_api::SubjektivSurfaceMaterial>,
body_token_budget: usize,
) -> Self {
Self {
materials: Arc::new(materials),
body_token_budget,
submitted: Arc::new(Mutex::new(None)),
}
}
pub(crate) fn submitted(&self) -> Option<Vec<server_api::SubjektivSurfacePoint>> {
self.submitted
.lock()
.expect("Memory surface output state poisoned")
.clone()
}
fn validate(&self, points: &[server_api::SubjektivSurfacePoint]) -> Result<(), ToolError> {
if self.materials.is_empty() {
return if points.is_empty() {
Ok(())
} else {
Err(ToolError::InvalidArgument(
"empty materials require an empty points array".into(),
))
};
}
if points.is_empty() {
return Err(ToolError::InvalidArgument(
"non-empty materials require at least one grounded point".into(),
));
}
if points.len() > 24 {
return Err(ToolError::InvalidArgument(format!(
"points count {} exceeds the configured limit 24",
points.len()
)));
}
let allowed = self
.materials
.iter()
.map(|material| (material.memory_id.as_str(), material.revision))
.collect::<HashSet<_>>();
let mut bodies = Vec::with_capacity(points.len());
for point in points {
let body = point.body_md.trim();
if body.is_empty() {
return Err(ToolError::InvalidArgument(
"surface point body_md must not be empty".into(),
));
}
if point.memory_refs.is_empty() {
return Err(ToolError::InvalidArgument(
"every surface point requires at least one memory_ref".into(),
));
}
let mut point_refs = HashSet::new();
for reference in &point.memory_refs {
if !point_refs.insert((reference.memory_id.as_str(), reference.revision)) {
return Err(ToolError::InvalidArgument(format!(
"duplicate surface point ref {}@{}",
reference.memory_id, reference.revision
)));
}
if !allowed.contains(&(reference.memory_id.as_str(), reference.revision)) {
return Err(ToolError::InvalidArgument(format!(
"surface ref {}@{} is not in the generation materials",
reference.memory_id, reference.revision
)));
}
}
bodies.push(body);
}
let bytes = bodies
.iter()
.map(|body| body.len())
.sum::<usize>()
.saturating_add(points.len().saturating_sub(1) * 2);
let tokens = bytes.saturating_add(3) / 4;
if tokens > self.body_token_budget {
return Err(ToolError::InvalidArgument(format!(
"surface body token estimate {tokens} exceeds budget {}",
self.body_token_budget
)));
}
Ok(())
}
}
#[derive(Clone)]
pub(crate) struct MemorySurfaceOutputFeature {
state: MemorySurfaceOutputState,
}
impl MemorySurfaceOutputFeature {
pub(crate) fn new(state: MemorySurfaceOutputState) -> Self {
Self { state }
}
}
impl FeatureModule for MemorySurfaceOutputFeature {
fn descriptor(&self) -> FeatureDescriptor {
FeatureDescriptor::builtin("memory-surface-output", "Memory Surface Output")
.with_description(
"Restricted structured output for one clean-context Memory surface editor.",
)
.with_tool(ToolDeclaration::new(SUBMIT_TOOL, SUBMIT_DESCRIPTION))
}
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
context.tools().register(ToolContribution::new(
SUBMIT_TOOL,
submit_definition(self.state.clone()),
))
}
}
fn submit_input_schema() -> serde_json::Value {
serde_json::to_value(schemars::schema_for!(SubmitMemorySurfaceParams))
.unwrap_or_else(|_| serde_json::json!({}))
}
pub(crate) fn submit_llm_tool_definition() -> agen::llm_client::ToolDefinition {
agen::llm_client::ToolDefinition::new(SUBMIT_TOOL)
.description(SUBMIT_DESCRIPTION)
.input_schema(submit_input_schema())
}
fn submit_definition(state: MemorySurfaceOutputState) -> ToolDefinition {
Arc::new(move || {
let definition = submit_llm_tool_definition();
let meta = ToolMeta::new(definition.name)
.description(
definition
.description
.expect("SubmitMemorySurface description is fixed"),
)
.input_schema(definition.input_schema);
let tool: Arc<dyn Tool> = Arc::new(SubmitMemorySurfaceTool {
state: state.clone(),
});
(meta, tool)
})
}
#[derive(Debug, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
struct SubmitMemorySurfaceParams {
points: Vec<server_api::SubjektivSurfacePoint>,
}
struct SubmitMemorySurfaceTool {
state: MemorySurfaceOutputState,
}
#[async_trait]
impl Tool for SubmitMemorySurfaceTool {
async fn execute(
&self,
input_json: &str,
_context: agen::tool::ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let params: SubmitMemorySurfaceParams =
serde_json::from_str(input_json).map_err(|error| {
ToolError::InvalidArgument(format!("invalid SubmitMemorySurface input: {error}"))
})?;
self.state.validate(&params.points)?;
let mut submitted = self
.state
.submitted
.lock()
.expect("Memory surface output state poisoned");
if submitted.is_some() {
return Err(ToolError::InvalidArgument(
"SubmitMemorySurface already accepted output for this generation".into(),
));
}
let count = params.points.len();
*submitted = Some(params.points);
Ok(ToolOutput {
summary: format!("Accepted {count} grounded Memory surface point(s)."),
content: None,
attachments: Vec::new(),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use memory::extract::CandidateKind;
fn material(id: &str) -> server_api::SubjektivSurfaceMaterial {
server_api::SubjektivSurfaceMaterial {
memory_id: id.into(),
revision: 1,
kind: CandidateKind::Constraint,
body_md: "Keep the boundary.".into(),
why_useful: "Avoid regression.".into(),
staleness: Some("Until the boundary changes.".into()),
}
}
fn point(id: &str) -> server_api::SubjektivSurfacePoint {
server_api::SubjektivSurfacePoint {
body_md: "- Keep the boundary unless it changes.".into(),
memory_refs: vec![server_api::SubjektivMemoryRevisionRef {
memory_id: id.into(),
revision: 1,
}],
}
}
#[test]
fn rejects_missing_or_material_external_refs_and_budget_overflow() {
let state = MemorySurfaceOutputState::new(vec![material("allowed")], 32);
assert!(state.validate(&[]).is_err());
assert!(state.validate(&[point("other")]).is_err());
let mut oversized = point("allowed");
oversized.body_md = "x".repeat(132);
assert!(state.validate(&[oversized]).is_err());
assert!(state.validate(&[point("allowed")]).is_ok());
}
#[test]
fn empty_materials_accept_only_normal_empty_output() {
let state = MemorySurfaceOutputState::new(Vec::new(), 32);
assert!(state.validate(&[]).is_ok());
assert!(state.validate(&[point("invented")]).is_err());
}
}
crates/worker/src/prompt/catalog.rs
+42
View File
@@ -95,6 +95,7 @@ pub enum WorkerPrompt {
MemoryExtractSystem,
MemoryConsolidationSystem,
SubjektivMemoryConsolidationSystem,
SubjektivMemorySurfaceSystem,
FlowVerifierSystem,
BackendJobSystem,
NotifyWrapper,
@@ -103,6 +104,7 @@ pub enum WorkerPrompt {
WorkingBoundariesSection,
AgentsMdSection,
ResidentMemorySummarySection,
ResidentMemoryRestoreSection,
WorkerOrchestrationGuidanceSection,
SubWorkerSpawnToolDescription,
}
@@ -116,6 +118,7 @@ impl WorkerPrompt {
Self::SubjektivMemoryConsolidationSystem => {
"internal.subjektiv_memory_consolidation_system"
}
Self::SubjektivMemorySurfaceSystem => "internal.subjektiv_memory_surface_system",
Self::FlowVerifierSystem => "internal.flow_verifier_system",
Self::BackendJobSystem => "internal.backend_job_system",
Self::NotifyWrapper => "internal.notify_wrapper",
@@ -124,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"
}
@@ -136,6 +140,7 @@ impl WorkerPrompt {
WorkerPrompt::MemoryExtractSystem,
WorkerPrompt::MemoryConsolidationSystem,
WorkerPrompt::SubjektivMemoryConsolidationSystem,
WorkerPrompt::SubjektivMemorySurfaceSystem,
WorkerPrompt::FlowVerifierSystem,
WorkerPrompt::BackendJobSystem,
WorkerPrompt::NotifyWrapper,
@@ -144,6 +149,7 @@ impl WorkerPrompt {
WorkerPrompt::WorkingBoundariesSection,
WorkerPrompt::AgentsMdSection,
WorkerPrompt::ResidentMemorySummarySection,
WorkerPrompt::ResidentMemoryRestoreSection,
WorkerPrompt::WorkerOrchestrationGuidanceSection,
WorkerPrompt::SubWorkerSpawnToolDescription,
];
@@ -463,6 +469,12 @@ impl PromptCatalog {
single("language", language),
)
}
pub fn subjektiv_memory_surface_system(&self, language: &str) -> Result<String, CatalogError> {
self.render(
WorkerPrompt::SubjektivMemorySurfaceSystem,
single("language", language),
)
}
pub fn flow_verifier_system(&self) -> Result<String, CatalogError> {
self.render(WorkerPrompt::FlowVerifierSystem, Value::UNDEFINED)
}
@@ -493,6 +505,19 @@ impl PromptCatalog {
single("summary", summary),
)
}
pub fn resident_memory_restore_section(
&self,
surface_ready: bool,
summary: Option<&str>,
) -> Result<String, CatalogError> {
self.render(
WorkerPrompt::ResidentMemoryRestoreSection,
minijinja::context! {
surface_ready => surface_ready,
summary => summary.unwrap_or_default(),
},
)
}
pub fn worker_orchestration_guidance_section(&self) -> Result<String, CatalogError> {
self.render(
WorkerPrompt::WorkerOrchestrationGuidanceSection,
@@ -927,6 +952,23 @@ mod tests {
.unwrap()
.contains("LANGUAGE_MARKER")
);
let surface = catalog
.subjektiv_memory_surface_system("SURFACE_LANGUAGE")
.unwrap();
assert!(surface.contains("SURFACE_LANGUAGE"));
for required in [
"conditions",
"negation",
"contradictions",
"staleness",
"SubmitMemorySurface",
"not semantic verification",
] {
assert!(
surface.contains(required),
"missing surface rule {required}"
);
}
assert!(
catalog
.notify_wrapper("NOTIFICATION_MARKER")
crates/worker/src/worker.rs
+285 -8
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;
@@ -2335,9 +2337,15 @@ impl WorkerSession {
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum SystemPromptContribution {
Ready(String),
Unavailable,
}
#[async_trait::async_trait]
pub(crate) trait SystemPromptContributionSource: Send + Sync {
async fn load(&self) -> Option<String>;
async fn load(&self) -> SystemPromptContribution;
}
#[derive(Debug, Clone)]
@@ -2501,6 +2509,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
@@ -2754,6 +2764,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(),
@@ -3821,6 +3832,66 @@ 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 contribution = if self.inject_resident_summary {
source.load().await
} else {
SystemPromptContribution::Unavailable
};
let (surface_ready, summary) = match &contribution {
SystemPromptContribution::Ready(summary) => (true, Some(summary.as_str())),
SystemPromptContribution::Unavailable => (false, None),
};
let prompt_catalog = self.prompts.load_full();
let body = prompt_catalog.resident_memory_restore_section(surface_ready, summary)?;
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
@@ -3828,9 +3899,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
@@ -3863,7 +3937,14 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
}
let resident_summary = if self.inject_resident_summary {
match &self.feature_resident_summary_source {
Some(source) => source.load().await,
Some(source) => match source.load().await {
SystemPromptContribution::Ready(summary) if !summary.is_empty() => {
Some(summary)
}
SystemPromptContribution::Ready(_) | SystemPromptContribution::Unavailable => {
None
}
},
None => None,
}
} else {
@@ -3919,6 +4000,10 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
.is_some_and(|state| state.pre_run_eligible(self.total_tokens().tokens))
}
pub(crate) fn needs_initial_session_head_materialization(&self) -> bool {
self.system_prompt_template.is_some()
}
/// Materialize the initial durable session head without running the model.
///
/// Runtime-managed Workers call this before creation is accepted so even a
@@ -6557,6 +6642,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(),
@@ -6646,6 +6732,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(),
@@ -6770,6 +6857,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(),
@@ -7151,6 +7239,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
@@ -8710,13 +8799,13 @@ mod build_summary_prompt_tests {
use super::*;
struct TestSystemPromptContributionSource {
value: Option<String>,
value: SystemPromptContribution,
load_count: Arc<AtomicUsize>,
}
#[async_trait::async_trait]
impl SystemPromptContributionSource for TestSystemPromptContributionSource {
async fn load(&self) -> Option<String> {
async fn load(&self) -> SystemPromptContribution {
self.load_count.fetch_add(1, Ordering::SeqCst);
self.value.clone()
}
@@ -8728,7 +8817,10 @@ mod build_summary_prompt_tests {
let load_count = Arc::new(AtomicUsize::new(0));
(
Arc::new(TestSystemPromptContributionSource {
value,
value: value.map_or(
SystemPromptContribution::Unavailable,
SystemPromptContribution::Ready,
),
load_count: Arc::clone(&load_count),
}),
load_count,
@@ -11083,6 +11175,191 @@ 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 empty_restore =
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 (empty_source, empty_loads) =
test_system_prompt_contribution_source(Some(String::new()));
empty_restore.install_system_prompt_contribution(Some(empty_source), None);
empty_restore
.materialize_durable_session_head()
.await
.unwrap();
assert_eq!(empty_loads.load(Ordering::SeqCst), 1);
let empty_boundary = empty_restore
.history()
.last()
.unwrap()
.as_text()
.unwrap()
.to_string();
assert!(
empty_boundary.contains("current ready resident memory surface is intentionally empty")
);
assert!(!empty_boundary.contains("No current ready resident memory surface"));
drop(empty_restore);
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();