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

fix: preserve resident surface startup states

Browse Source
This commit is contained in:
Hare
2026-10-02 15:01:09 +09:00
parent e662e2f465
commit dfe837655a
10 changed files with 230 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/memory/src/backend.rs
+18
View File
@@ -41,6 +41,7 @@ pub enum MemoryBackendHttpResponse {
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum MemoryBackendOperationResult {
ToolOutput(MemoryToolOutput),
ResidentSummary(MemoryResidentSummaryOutput),
Acknowledged(MemoryBackendAckOutput),
StagingWritten(MemoryStagingWriteOutput),
}
@@ -52,6 +53,23 @@ pub struct MemoryToolOutput {
pub content: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum MemoryResidentSummaryAvailability {
Ready,
Ungenerated,
Stale,
Failed,
Unavailable,
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct MemoryResidentSummaryOutput {
pub availability: MemoryResidentSummaryAvailability,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct MemoryQueryOperation {
#[serde(default, skip_serializing_if = "Option::is_none")]
crates/worker-runtime/src/worker_backend.rs
-4
View File
@@ -1266,10 +1266,6 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
}
}
worker
.materialize_durable_session_head()
.await
.map_err(|error| format!("materialize durable Worker session head: {error}"))?;
let workspace_client = worker.workspace_client_handle();
let started = prepared.start().await.map_err(|error| match error {
WorkerBootstrapError::Worker(source) => {
crates/worker/src/controller.rs
+35
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();
@@ -3465,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/memory.rs
+22 -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)]
@@ -652,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(
@@ -662,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
}
}
}
@@ -843,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,
}
})
@@ -976,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())
);
}
@@ -1016,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/prompt/catalog.rs
+5 -1
View File
@@ -507,11 +507,15 @@ impl PromptCatalog {
}
pub fn resident_memory_restore_section(
&self,
surface_ready: bool,
summary: Option<&str>,
) -> Result<String, CatalogError> {
self.render(
WorkerPrompt::ResidentMemoryRestoreSection,
single("summary", summary.unwrap_or_default()),
minijinja::context! {
surface_ready => surface_ready,
summary => summary.unwrap_or_default(),
},
)
}
pub fn worker_orchestration_guidance_section(&self) -> Result<String, CatalogError> {
crates/worker/src/worker.rs
+66 -8
View File
@@ -2336,9 +2336,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)]
@@ -3835,13 +3841,17 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.restore_resident_summary_refresh_pending = false;
return Ok(());
};
let summary = if self.inject_resident_summary {
let contribution = if self.inject_resident_summary {
source.load().await
} else {
None
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(summary.as_deref())?;
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 {
@@ -3925,7 +3935,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 {
@@ -3981,6 +3998,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
@@ -8770,13 +8791,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()
}
@@ -8788,7 +8809,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,
@@ -11246,6 +11270,40 @@ permission = "write"
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,
crates/workspace-server/src/memory_backend.rs
+11 -8
View File
@@ -1,8 +1,9 @@
use memory::backend::{
MemoryBackendAckOutput, MemoryBackendOperation, MemoryBackendOperationResult,
MemoryDocumentReadOperation, MemoryDocumentUpdateOperation, MemoryQueryOperation,
MemoryStagingCloseAction, MemoryStagingCloseOperation, MemoryStagingListOperation,
MemoryStagingReadOperation, MemoryStagingWriteOutput, MemoryToolOutput,
MemoryResidentSummaryAvailability, MemoryResidentSummaryOutput, MemoryStagingCloseAction,
MemoryStagingCloseOperation, MemoryStagingListOperation, MemoryStagingReadOperation,
MemoryStagingWriteOutput, MemoryToolOutput,
};
use memory::extract::{ExtractedCandidate, StagingEvidence, StagingRecord};
use memory::schema::{SourceEvidenceRef, SourceRef};
@@ -32,7 +33,7 @@ pub fn execute_memory_backend_operation_with_authority<A: MemoryAuthority>(
.map(MemoryBackendOperationResult::ToolOutput)
}
MemoryBackendOperation::ResidentSummary(_operation) => Ok(
MemoryBackendOperationResult::ToolOutput(execute_resident_summary(authority)?),
MemoryBackendOperationResult::ResidentSummary(execute_resident_summary(authority)?),
),
MemoryBackendOperation::AppendAudit(_operation) => Ok(
MemoryBackendOperationResult::Acknowledged(MemoryBackendAckOutput {
@@ -88,16 +89,18 @@ pub fn execute_memory_backend_operation_with_authority<A: MemoryAuthority>(
}
}
fn execute_resident_summary<A: MemoryAuthority>(authority: &A) -> Result<MemoryToolOutput> {
fn execute_resident_summary<A: MemoryAuthority>(
authority: &A,
) -> Result<MemoryResidentSummaryOutput> {
let document = authority.ensure_memory_document()?;
if document.body_md.trim().is_empty() {
Ok(MemoryToolOutput {
summary: "resident memory summary unavailable".to_string(),
Ok(MemoryResidentSummaryOutput {
availability: MemoryResidentSummaryAvailability::Unavailable,
content: None,
})
} else {
Ok(MemoryToolOutput {
summary: "resident memory document collected".to_string(),
Ok(MemoryResidentSummaryOutput {
availability: MemoryResidentSummaryAvailability::Ready,
content: Some(document.body_md),
})
}
crates/workspace-server/src/server.rs
+67 -18
View File
@@ -21,7 +21,7 @@ use flow::{FlowSourceKind, FlowSourceResolveRequest, ResolvedFlowSource};
use futures::{SinkExt, StreamExt};
use memory::backend::{
MemoryBackendHttpResponse, MemoryBackendOperation, MemoryBackendOperationResult,
MemoryConsolidateStagingOperation, MemoryConsolidationOutput, MemoryToolOutput,
MemoryConsolidateStagingOperation, MemoryConsolidationOutput,
};
use protocol::Segment;
use protocol::stream::{decode_method, encode_event};
@@ -16590,34 +16590,39 @@ fn execute_scoped_resident_summary(
.map_err(|error| Error::Store(error.to_string()))?
.resident_surface(subject_id)
.map_err(|error| Error::Store(error.to_string()))?;
let (summary, content) = match resident.availability {
Ok(MemoryBackendOperationResult::ResidentSummary(
resident_summary_output(resident),
))
}
fn resident_summary_output(
resident: crate::subjektiv::ResidentSurface,
) -> memory::backend::MemoryResidentSummaryOutput {
let (availability, content) = match resident.availability {
crate::subjektiv::SurfaceAvailability::Ready => {
let snapshot = resident.snapshot.expect("ready surface has a snapshot");
if snapshot.body_md.is_empty() {
("subject resident Memory surface is normally empty", None)
} else {
(
"subject resident Memory surface collected",
Some(snapshot.body_md),
)
}
}
crate::subjektiv::SurfaceAvailability::Ungenerated => {
("subject resident Memory surface is not generated", None)
(
memory::backend::MemoryResidentSummaryAvailability::Ready,
(!snapshot.body_md.is_empty()).then_some(snapshot.body_md),
)
}
crate::subjektiv::SurfaceAvailability::Ungenerated => (
memory::backend::MemoryResidentSummaryAvailability::Ungenerated,
None,
),
crate::subjektiv::SurfaceAvailability::Stale => (
"subject resident Memory surface is stale and was omitted",
memory::backend::MemoryResidentSummaryAvailability::Stale,
None,
),
crate::subjektiv::SurfaceAvailability::Failed => (
"subject resident Memory surface generation failed and was omitted",
memory::backend::MemoryResidentSummaryAvailability::Failed,
None,
),
};
Ok(MemoryBackendOperationResult::ToolOutput(MemoryToolOutput {
summary: summary.into(),
memory::backend::MemoryResidentSummaryOutput {
availability,
content,
}))
}
}
const SUBJEKTIV_SINGLETON_PREFIX: &str = "subjektiv:";
@@ -30146,6 +30151,50 @@ mod tests {
bytes
}
#[test]
fn resident_summary_transport_preserves_ready_empty_state() {
let output = resident_summary_output(crate::subjektiv::ResidentSurface {
availability: crate::subjektiv::SurfaceAvailability::Ready,
snapshot: Some(crate::subjektiv::SurfaceSnapshot {
schema_version: 1,
id: "surface-empty".into(),
subject_id: "subject-a".into(),
body_md: String::new(),
memory_refs: Vec::new(),
built_from_store_revision: 7,
created_at: "2026-10-02T00:00:00Z".into(),
}),
});
assert_eq!(
output.availability,
memory::backend::MemoryResidentSummaryAvailability::Ready
);
assert_eq!(output.content, None);
for (surface, expected) in [
(
crate::subjektiv::SurfaceAvailability::Ungenerated,
memory::backend::MemoryResidentSummaryAvailability::Ungenerated,
),
(
crate::subjektiv::SurfaceAvailability::Stale,
memory::backend::MemoryResidentSummaryAvailability::Stale,
),
(
crate::subjektiv::SurfaceAvailability::Failed,
memory::backend::MemoryResidentSummaryAvailability::Failed,
),
] {
let output = resident_summary_output(crate::subjektiv::ResidentSurface {
availability: surface,
snapshot: None,
});
assert_eq!(output.availability, expected);
assert_eq!(output.content, None);
}
}
#[test]
fn external_permission_levels_map_to_existing_bits_and_reject_legacy_partial_grants() {
let bits = |capabilities| {