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/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,