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

memory: consolidate staging through queue tools

Browse Source
This commit is contained in:
Hare
2026-07-25 09:17:08 +09:00
parent 4de801f8d0
commit a44673fa6b
13 changed files with 701 additions and 1660 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+12 -4
View File
@@ -646,10 +646,18 @@ where
base_url,
} = workspace_client
{
for definition in crate::feature::builtin::memory::workspace_http_memory_tools(
workspace_id,
base_url,
) {
let definitions = if spawner_name == "memory-consolidation" {
crate::feature::builtin::memory::workspace_http_memory_consolidation_tools(
workspace_id,
base_url,
)
} else {
crate::feature::builtin::memory::workspace_http_memory_tools(
workspace_id,
base_url,
)
};
for definition in definitions {
worker.register_tool(definition);
}
} else {
crates/worker/src/feature/builtin/memory.rs
+109 -2
View File
@@ -12,9 +12,11 @@ use llm_engine::tool::{
};
use memory::backend::{
MemoryBackendHttpResponse, MemoryBackendOperation, MemoryBackendOperationResult,
MemoryDeleteOperation, MemoryEditOperation, MemoryQueryOperation, MemoryReadOperation,
MemoryToolOutput, MemoryWriteOperation,
MemoryConsolidateStagingOperation, MemoryConsolidationOutput, MemoryDeleteOperation,
MemoryEditOperation, MemoryQueryOperation, MemoryReadOperation, MemoryStagingCloseOperation,
MemoryStagingListOperation, MemoryStagingReadOperation, MemoryToolOutput, MemoryWriteOperation,
};
use schemars::JsonSchema;
use serde::de::DeserializeOwned;
use serde_json::json;
@@ -91,6 +93,28 @@ impl WorkspaceClient {
}
}
}
pub async fn request_memory_staging_consolidation(
&self,
operation: MemoryConsolidateStagingOperation,
) -> Result<MemoryConsolidationOutput, WorkspaceMemoryBackendError> {
match self {
WorkspaceClient::Http {
workspace_id,
base_url,
} => execute_http_memory_consolidation(workspace_id, base_url, operation).await,
WorkspaceClient::Available { kind } => Err(WorkspaceMemoryBackendError::Unavailable {
reason: format!(
"workspace client kind `{kind}` does not expose the Backend Workspace API"
),
}),
WorkspaceClient::Unavailable { reason } => {
Err(WorkspaceMemoryBackendError::Unavailable {
reason: reason.clone(),
})
}
}
}
}
async fn execute_http_memory_backend(
@@ -121,6 +145,29 @@ async fn execute_http_memory_backend(
}
}
async fn execute_http_memory_consolidation(
workspace_id: &str,
base_url: &str,
operation: MemoryConsolidateStagingOperation,
) -> Result<MemoryConsolidationOutput, WorkspaceMemoryBackendError> {
let url = format!(
"{}/api/w/{}/memory/consolidation",
base_url.trim_end_matches('/'),
workspace_id
);
let response = reqwest::Client::new()
.post(url)
.json(&operation)
.send()
.await?;
let status = response.status();
let body = response.text().await?;
if !status.is_success() {
return Err(WorkspaceMemoryBackendError::Http { status, body });
}
serde_json::from_str::<MemoryConsolidationOutput>(&body).map_err(Into::into)
}
pub fn workspace_http_memory_tools(
workspace_id: impl Into<String>,
base_url: impl Into<String>,
@@ -185,6 +232,58 @@ pub fn workspace_http_memory_tools(
]
}
pub fn workspace_http_memory_consolidation_tools(
workspace_id: impl Into<String>,
base_url: impl Into<String>,
) -> Vec<ToolDefinition> {
let workspace_id = workspace_id.into();
let base_url = base_url.into();
let mut tools = workspace_http_memory_tools(workspace_id.clone(), base_url.clone());
let backend = WorkspaceHttpMemoryBackend::new(workspace_id, base_url);
tools.extend([
memory_tool(
"MemoryStagingList",
STAGING_LIST_DESCRIPTION,
schema_for::<MemoryStagingListOperation>(),
backend.clone(),
|input| {
Ok(MemoryBackendOperation::StagingList(parse_input::<
MemoryStagingListOperation,
>(
input
)?))
},
),
memory_tool(
"MemoryStagingRead",
STAGING_READ_DESCRIPTION,
schema_for::<MemoryStagingReadOperation>(),
backend.clone(),
|input| {
Ok(MemoryBackendOperation::StagingRead(parse_input::<
MemoryStagingReadOperation,
>(
input
)?))
},
),
memory_tool(
"MemoryStagingClose",
STAGING_CLOSE_DESCRIPTION,
schema_for::<MemoryStagingCloseOperation>(),
backend,
|input| {
Ok(MemoryBackendOperation::StagingClose(parse_input::<
MemoryStagingCloseOperation,
>(
input
)?))
},
),
]);
tools
}
type OperationBuilder = fn(&str) -> Result<MemoryBackendOperation, ToolError>;
fn memory_tool(
@@ -229,6 +328,10 @@ fn parse_input<T: DeserializeOwned>(input: &str) -> Result<T, ToolError> {
serde_json::from_str(input).map_err(|error| ToolError::InvalidArgument(error.to_string()))
}
fn schema_for<T: JsonSchema>() -> serde_json::Value {
serde_json::to_value(schemars::schema_for!(T)).expect("memory tool schema should serialize")
}
fn tool_output(output: MemoryToolOutput) -> ToolOutput {
ToolOutput {
summary: output.summary,
@@ -243,6 +346,10 @@ const EDIT_DESCRIPTION: &str =
"Replace text in a durable memory record through Workspace authority.";
const DELETE_DESCRIPTION: &str = "Delete a durable memory record through Workspace authority.";
const QUERY_DESCRIPTION: &str = "Query durable memory records through Workspace authority.";
const STAGING_LIST_DESCRIPTION: &str =
"List pending Memory staging candidates without loading full record payloads.";
const STAGING_READ_DESCRIPTION: &str = "Read one pending Memory staging candidate by candidate_id.";
const STAGING_CLOSE_DESCRIPTION: &str = "Close one staging candidate with a required reason; records disposition and deletes the staging record.";
fn kind_schema() -> serde_json::Value {
json!({"type":"string","enum":["summary","decision","request"]})
crates/worker/src/worker.rs
+62 -35
View File
@@ -520,10 +520,7 @@ pub struct Worker<C: LlmClient, St: Store> {
/// the flag survives across `try_post_run_extract` calls without a
/// `&mut self` race.
extract_in_flight: Arc<AtomicBool>,
/// consolidation (memory.consolidation) in-process reentry guard. The
/// staging-side `StagingLock` already provides cross-process
/// exclusion, but this AtomicBool keeps a careless concurrent caller
/// inside the same Worker from racing on the staging snapshot.
/// consolidation (memory.consolidation) in-process reentry guard.
consolidation_in_flight: Arc<AtomicBool>,
/// Last completed extract boundary. `None` means no extract has
/// run yet on this session — next extract starts from entry 0.
@@ -3274,11 +3271,10 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
Ok(ExtractDecision::Completed)
}
/// consolidation (memory.consolidation) trigger.
/// Request Backend-managed Memory staging consolidation after a Worker turn.
///
/// Worker no longer has direct Workspace filesystem authority. Until consolidation is
/// exposed as a Backend Workspace Authority operation, the Worker must not inspect
/// staging, acquire staging locks, or register local memory tools directly.
/// Worker has no local Workspace memory authority. It only asks the Backend
/// Workspace to notify or spawn the dedicated consolidater Worker.
pub async fn try_post_run_consolidate(&mut self) -> Result<(), WorkerError> {
let Some(memory_cfg) = self.manifest.memory.clone() else {
return Ok(());
@@ -3289,30 +3285,64 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
.unwrap_or(&self.manifest.model);
let files_threshold = memory_cfg.consolidation_threshold_files.filter(|n| *n > 0);
let bytes_threshold = memory_cfg.consolidation_threshold_bytes.filter(|n| *n > 0);
let reason = if files_threshold.is_none() && bytes_threshold.is_none() {
"consolidation_threshold_disabled"
} else {
"consolidation_backend_operation_unavailable"
};
WorkerAuditBase::new(
memory::audit::AuditWorker::MemoryConsolidation,
memory::audit::AuditTrigger::StagingBacklog,
Some(model_audit_from_manifest(model)),
)
.emit(
self.workspace_client(),
self.event_tx.as_ref(),
memory::audit::WorkerLifecycleStatus::Skipped,
reason,
None,
None,
None,
)
.await;
if reason == "consolidation_backend_operation_unavailable" {
tracing::debug!(
"workspace memory consolidation skipped: backend operation is unavailable"
);
if files_threshold.is_none() && bytes_threshold.is_none() {
WorkerAuditBase::new(
memory::audit::AuditWorker::MemoryConsolidation,
memory::audit::AuditTrigger::StagingBacklog,
Some(model_audit_from_manifest(model)),
)
.emit(
self.workspace_client(),
self.event_tx.as_ref(),
memory::audit::WorkerLifecycleStatus::Skipped,
"consolidation_threshold_disabled",
None,
None,
None,
)
.await;
return Ok(());
}
match self
.workspace_client()
.request_memory_staging_consolidation(
memory::backend::MemoryConsolidateStagingOperation {
force: false,
threshold_files: files_threshold,
threshold_bytes: bytes_threshold,
},
)
.await
{
Ok(output) => {
tracing::debug!(
status = output.status.as_str(),
summary = output.summary.as_str(),
"requested backend memory staging consolidation"
);
}
Err(error) => {
tracing::warn!(
error = %error,
"failed to request backend memory staging consolidation"
);
WorkerAuditBase::new(
memory::audit::AuditWorker::MemoryConsolidation,
memory::audit::AuditTrigger::StagingBacklog,
Some(model_audit_from_manifest(model)),
)
.emit(
self.workspace_client(),
self.event_tx.as_ref(),
memory::audit::WorkerLifecycleStatus::Skipped,
"consolidation_backend_operation_failed",
None,
None,
None,
)
.await;
}
}
Ok(())
}
@@ -4587,9 +4617,6 @@ pub enum WorkerError {
#[error("feature install failed: {0}")]
FeatureInstall(String),
#[error("memory consolidation lock acquisition failed: {0}")]
ConsolidationLock(#[source] memory::consolidate::LockError),
#[error("session {segment_id} has no entries to restore")]
SegmentEmpty { segment_id: SegmentId },