memory: consolidate staging through queue tools
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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"]})
|
||||
|
||||
+62
-35
@@ -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 },
|
||||
|
||||
|
||||
@@ -1,551 +0,0 @@
|
||||
//! consolidation (memory.consolidation) post-run trigger.
|
||||
//!
|
||||
//! Covers the gating, lock and cleanup behaviour without exercising the
|
||||
//! full sub-worker tool loop:
|
||||
//!
|
||||
//! - no `[memory]` section → no-op
|
||||
//! - `[memory]` present but no thresholds → no-op
|
||||
//! - staging empty → skip
|
||||
//! - staging below thresholds → skip + lock not acquired
|
||||
//! - staging above threshold → sub-worker runs, consumed entries removed
|
||||
//! - existing live lock → skip without error
|
||||
//!
|
||||
//! The sub-worker is fed a no-op LLM response (plain text) so it returns
|
||||
//! immediately. The post-run path then exercises lock acquisition,
|
||||
//! cleanup, and the empty-payload fast path.
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::Stream;
|
||||
use llm_engine::Engine;
|
||||
use llm_engine::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
|
||||
use llm_engine::llm_client::{ClientError, LlmClient, Request};
|
||||
use memory::WorkspaceLayout;
|
||||
use memory::extract::{CandidateKind, ExtractedCandidate, ExtractedPayload, write_staging};
|
||||
use memory::schema::SourceRef;
|
||||
use session_store::FsStore;
|
||||
use session_store::{CombinedStore, FsWorkerStore};
|
||||
|
||||
type TestStore = CombinedStore<FsStore, FsWorkerStore>;
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use worker::{Event, Worker};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockClient {
|
||||
responses: Arc<Vec<Vec<LlmEvent>>>,
|
||||
call_count: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl MockClient {
|
||||
fn new(responses: Vec<Vec<LlmEvent>>) -> Self {
|
||||
Self {
|
||||
responses: Arc::new(responses),
|
||||
call_count: Arc::new(AtomicUsize::new(0)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmClient for MockClient {
|
||||
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 count = self.call_count.fetch_add(1, Ordering::SeqCst);
|
||||
if count >= self.responses.len() {
|
||||
return Err(ClientError::Config("mock client exhausted".into()));
|
||||
}
|
||||
let events = self.responses[count].clone();
|
||||
let stream = futures::stream::iter(events.into_iter().map(Ok));
|
||||
Ok(Box::pin(stream))
|
||||
}
|
||||
}
|
||||
|
||||
fn done(text: &str) -> Vec<LlmEvent> {
|
||||
vec![
|
||||
LlmEvent::text_block_start(0),
|
||||
LlmEvent::text_delta(0, text),
|
||||
LlmEvent::text_block_stop(0, None),
|
||||
LlmEvent::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
]
|
||||
}
|
||||
|
||||
const NO_MEMORY_TOML: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
const MEMORY_NO_THRESHOLDS_TOML: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
const FILES_THRESHOLD_TOML: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
consolidation_threshold_files = 2
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
const ZERO_THRESHOLDS_TOML: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
consolidation_threshold_files = 0
|
||||
consolidation_threshold_bytes = 0
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
async fn make_worker_with(
|
||||
manifest_toml: &str,
|
||||
pwd: std::path::PathBuf,
|
||||
client: MockClient,
|
||||
) -> Worker<MockClient, TestStore> {
|
||||
let manifest = worker::WorkerManifest::from_toml(manifest_toml).unwrap();
|
||||
|
||||
let store_tmp = tempfile::tempdir().unwrap();
|
||||
let store = CombinedStore::new(
|
||||
FsStore::new(store_tmp.path()).unwrap(),
|
||||
FsWorkerStore::new(store_tmp.path().join("pods")).unwrap(),
|
||||
);
|
||||
std::mem::forget(store_tmp);
|
||||
|
||||
let scope = worker::Scope::writable(&pwd).unwrap();
|
||||
let worker = Engine::new(client);
|
||||
Worker::new(
|
||||
manifest,
|
||||
worker,
|
||||
store,
|
||||
worker::WorkerWorkspaceContext::local_filesystem(None),
|
||||
worker::WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()),
|
||||
scope,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn staging_payload(claim: String) -> ExtractedPayload {
|
||||
ExtractedPayload {
|
||||
candidates: vec![ExtractedCandidate {
|
||||
kind: CandidateKind::Lesson,
|
||||
claim,
|
||||
why_useful: "useful for consolidation trigger tests".into(),
|
||||
staleness: None,
|
||||
evidence_ids: Vec::new(),
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
fn write_n_staging(layout: &WorkspaceLayout, n: usize) -> Vec<uuid::Uuid> {
|
||||
let mut ids = Vec::new();
|
||||
for i in 0..n {
|
||||
let id = write_staging(
|
||||
layout,
|
||||
SourceRef {
|
||||
segment_id: format!("s-{i}"),
|
||||
range: [i as u64, i as u64],
|
||||
},
|
||||
staging_payload(format!("candidate-{i}")),
|
||||
)
|
||||
.unwrap()
|
||||
.remove(0)
|
||||
.id;
|
||||
ids.push(id);
|
||||
}
|
||||
ids
|
||||
}
|
||||
|
||||
fn attach_event_receiver(worker: &mut Worker<MockClient, TestStore>) -> broadcast::Receiver<Event> {
|
||||
let (tx, rx) = broadcast::channel(16);
|
||||
worker.attach_event_tx(tx);
|
||||
rx
|
||||
}
|
||||
|
||||
fn collect_memory_worker_reasons(rx: &mut broadcast::Receiver<Event>) -> Vec<String> {
|
||||
let mut reasons = Vec::new();
|
||||
loop {
|
||||
match rx.try_recv() {
|
||||
Ok(Event::MemoryWorker(event)) => reasons.push(event.reason),
|
||||
Ok(_) => {}
|
||||
Err(broadcast::error::TryRecvError::Empty) => break,
|
||||
Err(err) => panic!("unexpected broadcast receive error: {err}"),
|
||||
}
|
||||
}
|
||||
reasons
|
||||
}
|
||||
|
||||
fn read_audit_jsonl(layout: &WorkspaceLayout) -> Vec<serde_json::Value> {
|
||||
let text = std::fs::read_to_string(layout.audit_current_log_path()).unwrap();
|
||||
text.lines()
|
||||
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_memory_section_is_a_noop() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(NO_MEMORY_TOML, pwd.path().to_path_buf(), client).await;
|
||||
worker
|
||||
.try_post_run_consolidate()
|
||||
.await
|
||||
.expect("missing memory section must skip cleanly");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_thresholds_is_a_noop() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 5);
|
||||
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker =
|
||||
make_worker_with(MEMORY_NO_THRESHOLDS_TOML, pwd.path().to_path_buf(), client).await;
|
||||
worker
|
||||
.try_post_run_consolidate()
|
||||
.await
|
||||
.expect("consolidation disabled when both thresholds are None");
|
||||
|
||||
// No staging entries removed.
|
||||
assert_eq!(memory::consolidate::list_staging_entries(&layout).len(), 5);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn zero_thresholds_treated_as_disabled() {
|
||||
// Without the `Some(0) → None` collapse, `total_files >= 0` and
|
||||
// `total_bytes >= 0` would always evaluate true and consolidation would
|
||||
// fire on every post-run with any staging activity.
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 5);
|
||||
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(ZERO_THRESHOLDS_TOML, pwd.path().to_path_buf(), client).await;
|
||||
worker
|
||||
.try_post_run_consolidate()
|
||||
.await
|
||||
.expect("zero thresholds must collapse to disabled, not fire on every staging entry");
|
||||
|
||||
assert_eq!(
|
||||
memory::consolidate::list_staging_entries(&layout).len(),
|
||||
5,
|
||||
"staging must be untouched when both thresholds are zero"
|
||||
);
|
||||
let lock_path = layout.staging_dir().join(".consolidation.lock");
|
||||
assert!(!lock_path.exists(), "no lock should be acquired");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_staging_skips() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
// No mock calls expected.
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_staging_skip_is_audit_only() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
|
||||
let mut rx = attach_event_receiver(&mut worker);
|
||||
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
|
||||
assert!(collect_memory_worker_reasons(&mut rx).is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_only_staging_is_distinct_from_no_staging() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
std::fs::create_dir_all(layout.staging_dir()).unwrap();
|
||||
let invalid_id = uuid::Uuid::now_v7();
|
||||
let invalid_path = layout.staging_dir().join(format!("{invalid_id}.json"));
|
||||
std::fs::write(
|
||||
&invalid_path,
|
||||
serde_json::json!({
|
||||
"source": {
|
||||
"session_id": "legacy-session",
|
||||
"range": [0, 1]
|
||||
},
|
||||
"requests": []
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
|
||||
let mut rx = attach_event_receiver(&mut worker);
|
||||
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
|
||||
assert!(invalid_path.exists(), "invalid staging is not auto-deleted");
|
||||
let reasons = collect_memory_worker_reasons(&mut rx);
|
||||
assert_eq!(reasons, vec!["no_valid_staging_entries invalid=1"]);
|
||||
|
||||
let audit = read_audit_jsonl(&layout);
|
||||
let last = audit.last().unwrap();
|
||||
assert_eq!(last["reason"], "no_valid_staging_entries invalid=1");
|
||||
assert_eq!(last["consolidation"]["staging_count"], 0);
|
||||
assert_eq!(last["consolidation"]["invalid_staging_count"], 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn below_threshold_skip_is_audit_only() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 1); // threshold is 2
|
||||
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
|
||||
let mut rx = attach_event_receiver(&mut worker);
|
||||
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
|
||||
assert!(collect_memory_worker_reasons(&mut rx).is_empty());
|
||||
let audit = read_audit_jsonl(&layout);
|
||||
let reason = audit.last().unwrap()["reason"]
|
||||
.as_str()
|
||||
.expect("audit reason must be a string");
|
||||
assert!(reason.starts_with("threshold_not_reached "));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn completed_event_survives_terminal_empty_drain_skip() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 2); // threshold is 2 — fires.
|
||||
|
||||
let client = MockClient::new(vec![done("ok")]);
|
||||
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
|
||||
let mut rx = attach_event_receiver(&mut worker);
|
||||
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
|
||||
let reasons = collect_memory_worker_reasons(&mut rx);
|
||||
assert_eq!(reasons.len(), 2);
|
||||
assert!(reasons[0].starts_with("staging_threshold_reached files=2 bytes="));
|
||||
assert_eq!(reasons[1], "completed_no_record_changes");
|
||||
let audit = read_audit_jsonl(&layout);
|
||||
assert_eq!(audit.last().unwrap()["reason"], "no_staging_entries");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn below_threshold_skips_and_does_not_take_lock() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 1); // threshold is 2
|
||||
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
|
||||
// Staging untouched.
|
||||
assert_eq!(memory::consolidate::list_staging_entries(&layout).len(), 1);
|
||||
// Lock file must not exist.
|
||||
let lock_path = layout.staging_dir().join(".consolidation.lock");
|
||||
assert!(!lock_path.exists(), "lock file should not be created");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fires_on_threshold_and_cleans_up_consumed_entries() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 2); // threshold is 2 — fires.
|
||||
|
||||
// Sub-worker is given a single text-only response. The consolidation prompt
|
||||
// tells it to call memory tools; the mock skips those, but `Engine::run`
|
||||
// returns Ok regardless once the LLM closes with a final text.
|
||||
let client = MockClient::new(vec![done("ok")]);
|
||||
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
|
||||
// Consumed entries removed.
|
||||
assert!(
|
||||
memory::consolidate::list_staging_entries(&layout).is_empty(),
|
||||
"consumed staging entries must be cleaned up"
|
||||
);
|
||||
// Lock removed too.
|
||||
let lock_path = layout.staging_dir().join(".consolidation.lock");
|
||||
assert!(!lock_path.exists(), "lock file must be removed on success");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn in_flight_guard_skips_reentry_without_clearing() {
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 2);
|
||||
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(
|
||||
FILES_THRESHOLD_TOML,
|
||||
pwd.path().to_path_buf(),
|
||||
client.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Pre-set the in-flight flag as if another concurrent caller had
|
||||
// entered run_consolidate_once. The CAS at the top of
|
||||
// try_post_run_consolidate must take the early return without
|
||||
// touching staging or the LLM, and must leave the flag intact for
|
||||
// the holder to clear.
|
||||
let in_flight = worker.consolidation_in_flight_handle();
|
||||
in_flight.store(true, Ordering::Release);
|
||||
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
|
||||
assert!(
|
||||
in_flight.load(Ordering::Acquire),
|
||||
"reentry skip must not clear the in-flight flag — that's the holder's job"
|
||||
);
|
||||
assert_eq!(
|
||||
memory::consolidate::list_staging_entries(&layout).len(),
|
||||
2,
|
||||
"staging must remain untouched on reentry skip"
|
||||
);
|
||||
assert_eq!(
|
||||
client.call_count.load(Ordering::SeqCst),
|
||||
0,
|
||||
"no LLM calls should fire on reentry skip"
|
||||
);
|
||||
|
||||
// Sanity: when the flag is cleared, the same worker fires normally and
|
||||
// resets the flag itself (i.e. it isn't accidentally sticky).
|
||||
in_flight.store(false, Ordering::Release);
|
||||
let client2 = MockClient::new(vec![done("ok")]);
|
||||
let mut worker2 =
|
||||
make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client2).await;
|
||||
worker2.try_post_run_consolidate().await.unwrap();
|
||||
assert!(
|
||||
!worker2
|
||||
.consolidation_in_flight_handle()
|
||||
.load(Ordering::Acquire),
|
||||
"in-flight flag must be cleared after a normal run"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn coalesce_loop_terminates_with_one_iteration_when_snapshot_drains_staging() {
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
// Coalesce semantics from `docs/plan/memory.md` §並走防止: a single
|
||||
// run consumes the snapshot taken at acquire time; the loop
|
||||
// re-evaluates against any post-snapshot extract additions. With no
|
||||
// concurrent additions, the second iteration sees an empty staging
|
||||
// and bails out — exercised here by counting LLM calls.
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 4);
|
||||
|
||||
// Provide just one mock response. If the loop wrongly re-enters
|
||||
// run_consolidate_once after Completed, the second sub-worker run
|
||||
// would exhaust the mock and surface as an error.
|
||||
let client = MockClient::new(vec![done("ok")]);
|
||||
let mut worker = make_worker_with(
|
||||
FILES_THRESHOLD_TOML,
|
||||
pwd.path().to_path_buf(),
|
||||
client.clone(),
|
||||
)
|
||||
.await;
|
||||
worker.try_post_run_consolidate().await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
client.call_count.load(Ordering::SeqCst),
|
||||
1,
|
||||
"Coalesce must terminate once the staging snapshot is drained — got an extra LLM call"
|
||||
);
|
||||
assert!(
|
||||
memory::consolidate::list_staging_entries(&layout).is_empty(),
|
||||
"staging must be empty after the single iteration"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn live_lock_held_by_other_worker_skips() {
|
||||
let pwd = tempfile::tempdir().unwrap();
|
||||
let layout = WorkspaceLayout::new(pwd.path().to_path_buf());
|
||||
write_n_staging(&layout, 3);
|
||||
|
||||
// Pre-acquire lock with this test's PID — definitely alive — and
|
||||
// *don't* release it. The consolidation path must skip without error.
|
||||
let _live_lock = memory::consolidate::StagingLock::acquire(
|
||||
&layout,
|
||||
std::process::id(),
|
||||
"other-worker",
|
||||
Vec::new(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let client = MockClient::new(vec![]);
|
||||
let mut worker = make_worker_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client).await;
|
||||
worker
|
||||
.try_post_run_consolidate()
|
||||
.await
|
||||
.expect("InUse lock must surface as graceful skip");
|
||||
|
||||
// Staging untouched: lock holder owns the snapshot, not us.
|
||||
assert_eq!(memory::consolidate::list_staging_entries(&layout).len(), 3);
|
||||
}
|
||||
Reference in New Issue
Block a user