refactor: move memory extraction into lifecycle feature
This commit is contained in:
@@ -197,8 +197,8 @@ async fn finish_controller_run<C, St>(
|
||||
{
|
||||
// history / user_segments are no longer mirrored on WorkerSharedState —
|
||||
// clients reconstruct them from `Event::Snapshot` + live
|
||||
// `Event::Entry` deliveries driven by the session-log sink. We
|
||||
// flip the status and kick post-run memory jobs here.
|
||||
// `Event::Entry` deliveries driven by the session-log sink. The
|
||||
// lifecycle hook/task registry observes the terminal commit separately.
|
||||
//
|
||||
// In-flight blocks are run-local streaming state, not durable transcript.
|
||||
// Any block not cleared by a committed AssistantItem must be discarded at
|
||||
@@ -206,7 +206,6 @@ async fn finish_controller_run<C, St>(
|
||||
// partial text/tool arguments after newer entries.
|
||||
worker.clear_in_flight_events();
|
||||
set_controller_status(shared_state, runtime_dir, working_event_tx, new_status).await;
|
||||
worker.spawn_post_run_memory_jobs();
|
||||
}
|
||||
|
||||
/// Pending turn launch staged by an event handler for the next outer-loop
|
||||
@@ -995,6 +994,34 @@ where
|
||||
let worker_enabled = feature_config.worker.enabled;
|
||||
let sub_worker_enabled = feature_config.sub_worker.enabled;
|
||||
let mut feature_registry = FeatureRegistryBuilder::new();
|
||||
if feature_config.memory.enabled {
|
||||
let config = memory_config.clone().ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"[feature.memory].enabled = true requires a [memory] configuration section",
|
||||
)
|
||||
})?;
|
||||
let workspace_client = worker.workspace_client_handle();
|
||||
if !workspace_client.is_available() || workspace_client.workspace_id().is_none() {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
"Memory extraction requires Backend Workspace API authority",
|
||||
));
|
||||
}
|
||||
feature_registry.add_module(
|
||||
crate::feature::builtin::memory_lifecycle::MemoryExtractionLifecycleFeature::new(
|
||||
config,
|
||||
worker.committed_session_capture_handle(),
|
||||
worker.session_extension_handle(),
|
||||
workspace_client,
|
||||
spawner_manifest.clone(),
|
||||
worker.llm_client_handle(),
|
||||
prompts.clone(),
|
||||
spawner_workspace_context.clone(),
|
||||
worker.working_event_sender(),
|
||||
),
|
||||
);
|
||||
}
|
||||
if sub_worker_enabled && !worker_enabled {
|
||||
feature_registry.add_module(
|
||||
crate::feature::builtin::manage_worker::sub_worker_control_feature(
|
||||
@@ -1717,11 +1744,6 @@ async fn controller_loop<C, St>(
|
||||
// Memory/Workdir teardown so they cannot observe a partially closed Worker.
|
||||
worker.stop_feature_runtime("controller shutdown").await;
|
||||
|
||||
// Background memory jobs own extract/consolidate workers after a
|
||||
// turn completes. Join them before closing the Workdir session so no
|
||||
// Worker-owned task can outlive its operation attachment.
|
||||
worker.wait_for_memory_jobs().await;
|
||||
|
||||
if let Some(session) = worker.workdir_session()
|
||||
&& let Err(error) = session.close().await
|
||||
{
|
||||
|
||||
@@ -2212,6 +2212,7 @@ pub mod background;
|
||||
pub mod builtin;
|
||||
pub mod mcp;
|
||||
pub mod plugin;
|
||||
pub(crate) mod session;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
@@ -8,7 +8,8 @@ pub mod flow_transition;
|
||||
pub mod manage_workdir;
|
||||
pub mod manage_worker;
|
||||
pub mod memory;
|
||||
pub mod memory_extract;
|
||||
pub(crate) mod memory_lifecycle;
|
||||
pub mod memory_staging_output;
|
||||
pub mod merge_request;
|
||||
pub mod objective;
|
||||
pub mod orchestration;
|
||||
@@ -19,8 +20,6 @@ pub mod ticket;
|
||||
pub mod worker_observation;
|
||||
pub mod workspace_worker_discovery;
|
||||
|
||||
pub(crate) use memory_extract::{MemoryExtractFeature, MemoryExtractState, render_extract_input};
|
||||
pub(crate) use session_explore::{SessionExploreFeature, SessionExploreState};
|
||||
pub use task::{TaskFeature, task_tools_feature};
|
||||
pub use ticket::{
|
||||
TicketFeature, TicketFeatureAccess, ticket_tools_feature, ticket_tools_feature_with_access,
|
||||
|
||||
@@ -0,0 +1,775 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use agen::llm_client::LlmClient;
|
||||
use arc_swap::ArcSwap;
|
||||
use async_trait::async_trait;
|
||||
use memory::extract;
|
||||
use memory::schema::SourceRef;
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use crate::PromptCatalog;
|
||||
use crate::Scope;
|
||||
use crate::WorkerRunResult;
|
||||
use crate::feature::background::{
|
||||
BackgroundTaskCancellation, BackgroundTaskContext, BackgroundTaskSpec, BackgroundTaskTrigger,
|
||||
FeatureBackgroundTask,
|
||||
};
|
||||
use crate::feature::builtin::memory_staging_output::{
|
||||
MemoryStagingOutputFeature, MemoryStagingOutputState, render_extract_input,
|
||||
};
|
||||
use crate::feature::builtin::session_explore::{SessionExploreFeature, SessionExploreState};
|
||||
use crate::feature::session::{
|
||||
CommittedSessionCapture, CommittedSessionCaptureHandle, SessionExtensionHandle,
|
||||
};
|
||||
use crate::feature::{
|
||||
BackgroundTaskDeclaration, FeatureDescriptor, FeatureInstallContext, FeatureInstallError,
|
||||
FeatureModule, FeatureRegistryBuilder,
|
||||
};
|
||||
use crate::hook::{HookError, HookErrorCategory};
|
||||
use crate::internal_worker::{
|
||||
InternalWorkerAuthority, InternalWorkerError, InternalWorkerIdentity, InternalWorkerResult,
|
||||
InternalWorkerSpec, run_internal_worker_with_cancel_sender,
|
||||
};
|
||||
use crate::session_capture::SessionCapture;
|
||||
use crate::worker::{WorkerFilesystemAuthority, WorkerWorkspaceContext, WorkspaceClient};
|
||||
use agen::token_counter::total_tokens_at;
|
||||
use manifest::WorkerManifest;
|
||||
use protocol::Event;
|
||||
|
||||
const TASK_NAME: &str = "memory-extraction";
|
||||
const TASK_TIMEOUT: Duration = Duration::from_secs(300);
|
||||
|
||||
/// Parent-Worker lifecycle Feature that observes committed runs and schedules
|
||||
/// bounded extraction work. It owns the Memory pointer, audit, restricted
|
||||
/// Internal Worker, and staging disposition; Worker core owns only generic
|
||||
/// hook/task/session plumbing.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MemoryExtractionLifecycleFeature {
|
||||
task: MemoryExtractionTask,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MemoryExtractionTask {
|
||||
config: manifest::MemoryConfig,
|
||||
capture: CommittedSessionCaptureHandle,
|
||||
extensions: SessionExtensionHandle,
|
||||
workspace_client: Arc<dyn WorkspaceClient>,
|
||||
manifest: WorkerManifest,
|
||||
client: Box<dyn LlmClient>,
|
||||
prompts: Arc<ArcSwap<PromptCatalog>>,
|
||||
workspace_context: WorkerWorkspaceContext,
|
||||
event_tx: Option<broadcast::Sender<Event>>,
|
||||
}
|
||||
|
||||
impl MemoryExtractionLifecycleFeature {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn new(
|
||||
config: manifest::MemoryConfig,
|
||||
capture: CommittedSessionCaptureHandle,
|
||||
extensions: SessionExtensionHandle,
|
||||
workspace_client: Arc<dyn WorkspaceClient>,
|
||||
manifest: WorkerManifest,
|
||||
client: Box<dyn LlmClient>,
|
||||
prompts: Arc<ArcSwap<PromptCatalog>>,
|
||||
workspace_context: WorkerWorkspaceContext,
|
||||
event_tx: Option<broadcast::Sender<Event>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
task: MemoryExtractionTask {
|
||||
config,
|
||||
capture,
|
||||
extensions,
|
||||
workspace_client,
|
||||
manifest,
|
||||
client,
|
||||
prompts,
|
||||
workspace_context,
|
||||
event_tx,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FeatureModule for MemoryExtractionLifecycleFeature {
|
||||
fn descriptor(&self) -> FeatureDescriptor {
|
||||
FeatureDescriptor::builtin("memory-extraction-lifecycle", "Memory Extraction Lifecycle")
|
||||
.with_description(
|
||||
"Observes terminal committed runs and schedules bounded Memory extraction.",
|
||||
)
|
||||
.with_background_task(BackgroundTaskDeclaration::worker_managed(
|
||||
TASK_NAME,
|
||||
"Extract provenance-preserving Memory candidates after committed runs.",
|
||||
))
|
||||
}
|
||||
|
||||
fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> {
|
||||
context
|
||||
.background_tasks()
|
||||
.register(memory_extraction_task_spec(), self.task.clone())
|
||||
}
|
||||
}
|
||||
|
||||
fn memory_extraction_task_spec() -> BackgroundTaskSpec {
|
||||
let declaration = BackgroundTaskDeclaration::worker_managed(
|
||||
TASK_NAME,
|
||||
"Extract provenance-preserving Memory candidates after committed runs.",
|
||||
);
|
||||
let mut spec = BackgroundTaskSpec::single_flight(declaration, TASK_TIMEOUT);
|
||||
spec.trigger = BackgroundTaskTrigger::RunCommitted;
|
||||
spec
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl FeatureBackgroundTask for MemoryExtractionTask {
|
||||
async fn run(
|
||||
&self,
|
||||
context: BackgroundTaskContext,
|
||||
cancellation: BackgroundTaskCancellation,
|
||||
) -> Result<(), HookError> {
|
||||
context.generation_fence.ensure_current()?;
|
||||
let capture = self.capture.capture().map_err(hook_internal)?;
|
||||
let pointer = extract_pointer(&capture)?;
|
||||
if !extraction_threshold_reached(&capture, pointer.as_ref(), &self.config) {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let history_start = pointer
|
||||
.as_ref()
|
||||
.map(|pointer| pointer.processed_through_history_len)
|
||||
.unwrap_or(0)
|
||||
.min(capture.history.len());
|
||||
let history_end = capture.history.len();
|
||||
if history_start >= history_end || capture.entry_count == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
let view = SessionCapture::from_history_entries(
|
||||
capture.segment_id.clone(),
|
||||
capture.history[history_start..history_end].to_vec(),
|
||||
);
|
||||
let start_entry = pointer
|
||||
.as_ref()
|
||||
.map(|pointer| pointer.processed_through_entry + 1)
|
||||
.unwrap_or(0);
|
||||
let source = SourceRef {
|
||||
segment_id: capture.segment_id.clone(),
|
||||
range: [start_entry as u64, (capture.entry_count - 1) as u64],
|
||||
};
|
||||
let audit = WorkerAuditBase::new(
|
||||
memory::audit::AuditWorker::MemoryExtract,
|
||||
memory::audit::AuditTrigger::TokenThreshold,
|
||||
self.config
|
||||
.extract_model
|
||||
.as_ref()
|
||||
.or(Some(&self.manifest.model))
|
||||
.map(model_audit_from_manifest),
|
||||
)
|
||||
.with_memory_settings(&self.config);
|
||||
let extract_audit_base = memory::audit::ExtractAudit {
|
||||
session_id: Some(capture.session_id.clone()),
|
||||
segment_id: Some(capture.segment_id.clone()),
|
||||
entry_range: Some([start_entry as u64, (capture.entry_count - 1) as u64]),
|
||||
history_range: Some([history_start as u64, history_end as u64]),
|
||||
..Default::default()
|
||||
};
|
||||
audit
|
||||
.emit(
|
||||
self.workspace_client.as_ref(),
|
||||
self.event_tx.as_ref(),
|
||||
memory::audit::WorkerLifecycleStatus::Started,
|
||||
"token_threshold_reached",
|
||||
None,
|
||||
Some(extract_audit_base.clone()),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let output_state = MemoryStagingOutputState::new(
|
||||
view.clone(),
|
||||
Arc::clone(&self.workspace_client),
|
||||
source,
|
||||
audit.run_id.to_string(),
|
||||
);
|
||||
let client = if let Some(model) = self.config.extract_model.as_ref() {
|
||||
match crate::model_client::build_client(model) {
|
||||
Ok(client) => client,
|
||||
Err(error) => {
|
||||
self.record_preparation_failure(&audit, &extract_audit_base, error.to_string())
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
} else {
|
||||
self.client.clone_boxed()
|
||||
};
|
||||
let Some(memory_language) = self
|
||||
.config
|
||||
.workspace_settings()
|
||||
.map(|snapshot| snapshot.language)
|
||||
else {
|
||||
self.record_preparation_failure(
|
||||
&audit,
|
||||
&extract_audit_base,
|
||||
"Memory extraction requires a bound Workspace Memory settings snapshot",
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
};
|
||||
let system_prompt = match self
|
||||
.prompts
|
||||
.load_full()
|
||||
.memory_extract_system(&memory_language)
|
||||
{
|
||||
Ok(prompt) => prompt,
|
||||
Err(error) => {
|
||||
self.record_preparation_failure(&audit, &extract_audit_base, error.to_string())
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
let mut manifest = self.manifest.clone();
|
||||
if let Some(model) = self.config.extract_model.clone() {
|
||||
manifest.model = model;
|
||||
}
|
||||
|
||||
let cancel_observer = move |sender: tokio::sync::mpsc::Sender<()>| {
|
||||
tokio::spawn(async move {
|
||||
cancellation.cancelled().await;
|
||||
let _ = sender.send(()).await;
|
||||
});
|
||||
};
|
||||
let features = FeatureRegistryBuilder::new()
|
||||
.with_module(SessionExploreFeature::new(SessionExploreState::new(
|
||||
view.clone(),
|
||||
)))
|
||||
.with_module(MemoryStagingOutputFeature::new(output_state.clone()));
|
||||
let result = run_internal_worker_with_cancel_sender(
|
||||
InternalWorkerSpec {
|
||||
identity: InternalWorkerIdentity {
|
||||
kind: "memory-extract",
|
||||
run_id: audit.run_id,
|
||||
},
|
||||
manifest,
|
||||
client,
|
||||
system_prompt,
|
||||
input: render_extract_input(&view),
|
||||
cache_key: Some(capture.segment_id.clone()),
|
||||
max_turns: self
|
||||
.config
|
||||
.extract_worker_max_turns
|
||||
.or(manifest::defaults::MEMORY_EXTRACT_WORKER_MAX_TURNS),
|
||||
engine_configurator: None,
|
||||
features,
|
||||
required_tools: &[
|
||||
"ShowOverview",
|
||||
"SearchEntries",
|
||||
"ReadEntry",
|
||||
"StageMemoryCandidate",
|
||||
"FinishMemoryExtraction",
|
||||
],
|
||||
authority: InternalWorkerAuthority {
|
||||
workspace: self.workspace_context.clone(),
|
||||
filesystem: WorkerFilesystemAuthority::None,
|
||||
scope: Scope::empty(),
|
||||
workdir_session: None,
|
||||
},
|
||||
},
|
||||
cancel_observer,
|
||||
)
|
||||
.await;
|
||||
|
||||
let usage_event = match &result {
|
||||
Ok(run) => {
|
||||
tracing::debug!(
|
||||
worker_kind = run.identity.kind,
|
||||
run_id = %run.identity.run_id,
|
||||
history_entries = run.history_entries,
|
||||
"memory extraction Internal Worker completed"
|
||||
);
|
||||
run.usage.as_ref()
|
||||
}
|
||||
Err(error) => error.usage.as_ref(),
|
||||
};
|
||||
let usage_audit = usage_event.map(|event| memory::audit::UsageAudit {
|
||||
input_tokens: event.input_tokens,
|
||||
output_tokens: event.output_tokens,
|
||||
total_tokens: event.total_tokens,
|
||||
cache_read_input_tokens: event.cache_read_input_tokens,
|
||||
cache_creation_input_tokens: event.cache_creation_input_tokens,
|
||||
});
|
||||
let staging_ids = output_state.staged();
|
||||
let pointer_staging_id = staging_ids.first().cloned().unwrap_or_default();
|
||||
let extract_audit = Some(memory::audit::ExtractAudit {
|
||||
staging_count: staging_ids.len(),
|
||||
staging_paths: staging_ids,
|
||||
..extract_audit_base
|
||||
});
|
||||
|
||||
match extraction_disposition(&result, output_state.is_finished()) {
|
||||
ExtractionDisposition::Cancelled(reason) => {
|
||||
audit
|
||||
.emit(
|
||||
self.workspace_client.as_ref(),
|
||||
self.event_tx.as_ref(),
|
||||
memory::audit::WorkerLifecycleStatus::Cancelled,
|
||||
reason,
|
||||
usage_audit,
|
||||
extract_audit,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
ExtractionDisposition::Failed(reason) => {
|
||||
audit
|
||||
.emit(
|
||||
self.workspace_client.as_ref(),
|
||||
self.event_tx.as_ref(),
|
||||
memory::audit::WorkerLifecycleStatus::Failed,
|
||||
reason,
|
||||
usage_audit,
|
||||
extract_audit,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
ExtractionDisposition::Completed => {}
|
||||
}
|
||||
|
||||
context.generation_fence.ensure_current()?;
|
||||
let next_pointer = memory::ExtractPointerPayload {
|
||||
processed_through_entry: capture.entry_count - 1,
|
||||
processed_through_history_len: capture.history.len(),
|
||||
staging_id: pointer_staging_id,
|
||||
};
|
||||
let payload = serde_json::to_value(&next_pointer).map_err(hook_internal)?;
|
||||
if !self
|
||||
.extensions
|
||||
.append_if_current(&capture.location(), extract::EXTRACT_DOMAIN, payload)
|
||||
.map_err(hook_internal)?
|
||||
{
|
||||
audit
|
||||
.emit(
|
||||
self.workspace_client.as_ref(),
|
||||
self.event_tx.as_ref(),
|
||||
memory::audit::WorkerLifecycleStatus::Cancelled,
|
||||
"session changed before memory-extract pointer commit",
|
||||
usage_audit,
|
||||
extract_audit,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
audit
|
||||
.emit(
|
||||
self.workspace_client.as_ref(),
|
||||
self.event_tx.as_ref(),
|
||||
memory::audit::WorkerLifecycleStatus::Completed,
|
||||
"memory-extract completed",
|
||||
usage_audit,
|
||||
extract_audit,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl MemoryExtractionTask {
|
||||
async fn record_preparation_failure(
|
||||
&self,
|
||||
audit: &WorkerAuditBase,
|
||||
extract: &memory::audit::ExtractAudit,
|
||||
reason: impl Into<String>,
|
||||
) {
|
||||
audit
|
||||
.emit(
|
||||
self.workspace_client.as_ref(),
|
||||
self.event_tx.as_ref(),
|
||||
memory::audit::WorkerLifecycleStatus::Failed,
|
||||
reason,
|
||||
None,
|
||||
Some(extract.clone()),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
enum ExtractionDisposition {
|
||||
Completed,
|
||||
Failed(String),
|
||||
Cancelled(String),
|
||||
}
|
||||
|
||||
fn extraction_disposition(
|
||||
result: &Result<InternalWorkerResult, InternalWorkerError>,
|
||||
finish_called: bool,
|
||||
) -> ExtractionDisposition {
|
||||
match result {
|
||||
Err(error) => {
|
||||
// Preserve the Internal Worker result's immutable identity/history
|
||||
// evidence for diagnostics even though the public audit reason is
|
||||
// intentionally bounded to the typed source error.
|
||||
tracing::debug!(
|
||||
worker_kind = error.identity.kind,
|
||||
run_id = %error.identity.run_id,
|
||||
history_entries = error.history_entries,
|
||||
"memory extraction Internal Worker failed"
|
||||
);
|
||||
ExtractionDisposition::Failed(error.source.to_string())
|
||||
}
|
||||
Ok(run) => match &run.lifecycle {
|
||||
WorkerRunResult::RolledBack => {
|
||||
ExtractionDisposition::Cancelled("memory-extract cancelled".to_string())
|
||||
}
|
||||
WorkerRunResult::Interrupted { message, .. } => {
|
||||
ExtractionDisposition::Failed(message.clone())
|
||||
}
|
||||
WorkerRunResult::Finished | WorkerRunResult::Paused | WorkerRunResult::LimitReached
|
||||
if finish_called =>
|
||||
{
|
||||
ExtractionDisposition::Completed
|
||||
}
|
||||
WorkerRunResult::Finished | WorkerRunResult::Paused | WorkerRunResult::LimitReached => {
|
||||
ExtractionDisposition::Failed(
|
||||
"memory-extract did not call FinishMemoryExtraction".to_string(),
|
||||
)
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn now_millis() -> u64 {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis()
|
||||
.try_into()
|
||||
.unwrap_or(u64::MAX)
|
||||
}
|
||||
|
||||
fn hook_internal(error: impl std::fmt::Display) -> HookError {
|
||||
HookError::new(HookErrorCategory::Internal, error.to_string())
|
||||
}
|
||||
|
||||
fn extract_pointer(
|
||||
capture: &CommittedSessionCapture,
|
||||
) -> Result<Option<memory::ExtractPointerPayload>, HookError> {
|
||||
let pointer = memory::extract::fold_pointer(&capture.extensions);
|
||||
if pointer.is_none()
|
||||
&& capture
|
||||
.extensions
|
||||
.iter()
|
||||
.any(|(domain, _)| domain == extract::EXTRACT_DOMAIN)
|
||||
{
|
||||
return Err(hook_internal(
|
||||
"latest committed Memory extraction pointer is malformed",
|
||||
));
|
||||
}
|
||||
Ok(pointer)
|
||||
}
|
||||
|
||||
fn extraction_threshold_reached(
|
||||
capture: &CommittedSessionCapture,
|
||||
pointer: Option<&memory::ExtractPointerPayload>,
|
||||
config: &manifest::MemoryConfig,
|
||||
) -> bool {
|
||||
if capture.history.is_empty() {
|
||||
return false;
|
||||
}
|
||||
let history_pointer = pointer
|
||||
.map(|pointer| pointer.processed_through_history_len)
|
||||
.unwrap_or(0)
|
||||
.min(capture.history.len());
|
||||
let items = capture
|
||||
.history
|
||||
.iter()
|
||||
.map(|entry| entry.item.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let current = total_tokens_at(&items, &capture.usage_history, capture.history.len()).tokens;
|
||||
let baseline = total_tokens_at(&items, &capture.usage_history, history_pointer).tokens;
|
||||
let Some(threshold) = config.extract_threshold.filter(|threshold| *threshold > 0) else {
|
||||
return false;
|
||||
};
|
||||
current.saturating_sub(baseline) >= threshold
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct WorkerAuditBase {
|
||||
run_id: uuid::Uuid,
|
||||
worker: memory::audit::AuditWorker,
|
||||
trigger: memory::audit::AuditTrigger,
|
||||
memory_settings: Option<memory::audit::MemorySettingsAudit>,
|
||||
model: Option<memory::audit::ModelAudit>,
|
||||
}
|
||||
|
||||
impl WorkerAuditBase {
|
||||
fn new(
|
||||
worker: memory::audit::AuditWorker,
|
||||
trigger: memory::audit::AuditTrigger,
|
||||
model: Option<memory::audit::ModelAudit>,
|
||||
) -> Self {
|
||||
Self {
|
||||
run_id: uuid::Uuid::now_v7(),
|
||||
worker,
|
||||
trigger,
|
||||
memory_settings: None,
|
||||
model,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_memory_settings(mut self, config: &manifest::MemoryConfig) -> Self {
|
||||
self.memory_settings =
|
||||
config
|
||||
.workspace_settings()
|
||||
.map(|snapshot| memory::audit::MemorySettingsAudit {
|
||||
workspace_id: snapshot.workspace_id,
|
||||
settings_revision: snapshot.settings_revision,
|
||||
language: snapshot.language,
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
async fn emit(
|
||||
&self,
|
||||
workspace_client: &dyn WorkspaceClient,
|
||||
event_tx: Option<&broadcast::Sender<Event>>,
|
||||
status: memory::audit::WorkerLifecycleStatus,
|
||||
reason: impl Into<String>,
|
||||
usage: Option<memory::audit::UsageAudit>,
|
||||
extract: Option<memory::audit::ExtractAudit>,
|
||||
consolidation: Option<memory::audit::ConsolidationAudit>,
|
||||
) {
|
||||
let reason = reason.into();
|
||||
let payload = memory::audit::WorkerLifecycleAudit {
|
||||
run_id: self.run_id,
|
||||
worker: self.worker.clone(),
|
||||
status,
|
||||
trigger: self.trigger,
|
||||
reason: reason.clone(),
|
||||
memory_settings: self.memory_settings.clone(),
|
||||
model: self.model.clone(),
|
||||
usage,
|
||||
extract,
|
||||
consolidation,
|
||||
};
|
||||
let _ = workspace_client
|
||||
.execute_memory_backend_operation(memory::backend::MemoryBackendOperation::AppendAudit(
|
||||
memory::backend::MemoryAppendAuditOperation {
|
||||
event: memory::audit::AuditEvent::new(
|
||||
memory::audit::AuditPayload::WorkerLifecycle(payload),
|
||||
),
|
||||
},
|
||||
))
|
||||
.await;
|
||||
if let Some(tx) = event_tx {
|
||||
let _ = tx.send(Event::MemoryWorker(protocol::MemoryWorkerEvent {
|
||||
worker: self.worker.label().to_string(),
|
||||
status: status.label().to_string(),
|
||||
run_id: self.run_id.to_string(),
|
||||
trigger: self.trigger.label().to_string(),
|
||||
reason: reason.clone(),
|
||||
message: format!(
|
||||
"memory {} {}: {reason}",
|
||||
self.worker.label(),
|
||||
status.label()
|
||||
),
|
||||
timestamp_ms: now_millis() as i64,
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn model_audit_from_manifest(model: &manifest::ModelManifest) -> memory::audit::ModelAudit {
|
||||
memory::audit::ModelAudit {
|
||||
ref_: model.ref_.clone(),
|
||||
scheme: model.scheme.map(|scheme| format!("{scheme:?}")),
|
||||
model_id: model.model_id.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use agen::{HistoryEntry, Item, UsageRecord};
|
||||
|
||||
use super::*;
|
||||
use crate::feature::background::{BackgroundTaskRewritePolicy, BackgroundTaskShutdownPolicy};
|
||||
use crate::session_history::SessionHistoryMetadata;
|
||||
|
||||
fn internal_result(
|
||||
lifecycle: WorkerRunResult,
|
||||
) -> Result<InternalWorkerResult, InternalWorkerError> {
|
||||
Ok(InternalWorkerResult {
|
||||
usage: None,
|
||||
identity: InternalWorkerIdentity {
|
||||
kind: "memory-extract",
|
||||
run_id: uuid::Uuid::now_v7(),
|
||||
},
|
||||
lifecycle,
|
||||
history_entries: 1,
|
||||
})
|
||||
}
|
||||
|
||||
fn capture(history_len: usize, input_total_tokens: u64) -> CommittedSessionCapture {
|
||||
CommittedSessionCapture {
|
||||
session_id: "session-1".to_string(),
|
||||
segment_id: "segment-1".to_string(),
|
||||
session_revision: history_len.try_into().unwrap(),
|
||||
entry_count: history_len,
|
||||
history: (0..history_len)
|
||||
.map(|index| HistoryEntry {
|
||||
item: Item::user_message(format!("message-{index}")),
|
||||
annotation: SessionHistoryMetadata::legacy_unknown(),
|
||||
})
|
||||
.collect(),
|
||||
usage_history: vec![UsageRecord {
|
||||
history_len,
|
||||
input_total_tokens,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 0,
|
||||
}],
|
||||
extensions: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normal_and_empty_extraction_require_explicit_finish() {
|
||||
let result = internal_result(WorkerRunResult::Finished);
|
||||
assert_eq!(
|
||||
extraction_disposition(&result, true),
|
||||
ExtractionDisposition::Completed
|
||||
);
|
||||
// `finish_called = true` with no staged ids is the explicit empty
|
||||
// extraction outcome. Missing Finish is a failed extraction.
|
||||
assert!(matches!(
|
||||
extraction_disposition(&result, false),
|
||||
ExtractionDisposition::Failed(reason)
|
||||
if reason.contains("FinishMemoryExtraction")
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_and_pre_ai_cancelled_extraction_never_reach_pointer_commit() {
|
||||
let failed = internal_result(WorkerRunResult::Interrupted {
|
||||
code: crate::ErrorCode::Internal,
|
||||
message: "provider failed".to_string(),
|
||||
});
|
||||
assert!(matches!(
|
||||
extraction_disposition(&failed, true),
|
||||
ExtractionDisposition::Failed(_)
|
||||
));
|
||||
let cancelled = internal_result(WorkerRunResult::RolledBack);
|
||||
assert_eq!(
|
||||
extraction_disposition(&cancelled, true),
|
||||
ExtractionDisposition::Cancelled("memory-extract cancelled".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_scope_cancels_and_joins_before_rewrite_and_shutdown() {
|
||||
let spec = memory_extraction_task_spec();
|
||||
assert_eq!(spec.trigger, BackgroundTaskTrigger::RunCommitted);
|
||||
assert_eq!(spec.max_concurrency, 1);
|
||||
assert_eq!(spec.rewrite, BackgroundTaskRewritePolicy::CancelAndWait);
|
||||
assert_eq!(spec.shutdown, BackgroundTaskShutdownPolicy::CancelAndWait);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn threshold_uses_committed_usage_after_pointer() {
|
||||
let capture = capture(2, 250);
|
||||
let mut config = manifest::MemoryConfig::default();
|
||||
config.extract_threshold = Some(1);
|
||||
assert!(extraction_threshold_reached(
|
||||
&capture,
|
||||
Some(&memory::ExtractPointerPayload {
|
||||
processed_through_entry: 0,
|
||||
processed_through_history_len: 1,
|
||||
staging_id: "staging-1".to_string(),
|
||||
}),
|
||||
&config
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pointer_folds_latest_committed_extraction_extension() {
|
||||
let mut capture = capture(2, 250);
|
||||
let first = memory::ExtractPointerPayload {
|
||||
processed_through_entry: 1,
|
||||
processed_through_history_len: 1,
|
||||
staging_id: "staging-1".to_string(),
|
||||
};
|
||||
let latest = memory::ExtractPointerPayload {
|
||||
processed_through_entry: 3,
|
||||
processed_through_history_len: 2,
|
||||
staging_id: "staging-2".to_string(),
|
||||
};
|
||||
capture.extensions = vec![
|
||||
(
|
||||
extract::EXTRACT_DOMAIN.to_string(),
|
||||
serde_json::to_value(&first).unwrap(),
|
||||
),
|
||||
("other.feature".to_string(), serde_json::json!({})),
|
||||
(
|
||||
extract::EXTRACT_DOMAIN.to_string(),
|
||||
serde_json::to_value(&latest).unwrap(),
|
||||
),
|
||||
];
|
||||
assert_eq!(extract_pointer(&capture).unwrap(), Some(latest));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_latest_pointer_fails_closed_instead_of_using_older_pointer() {
|
||||
let mut capture = capture(2, 250);
|
||||
capture.extensions = vec![
|
||||
(
|
||||
extract::EXTRACT_DOMAIN.to_string(),
|
||||
serde_json::to_value(memory::ExtractPointerPayload {
|
||||
processed_through_entry: 1,
|
||||
processed_through_history_len: 1,
|
||||
staging_id: "staging-1".to_string(),
|
||||
})
|
||||
.unwrap(),
|
||||
),
|
||||
(
|
||||
extract::EXTRACT_DOMAIN.to_string(),
|
||||
serde_json::json!({"invalid": true}),
|
||||
),
|
||||
];
|
||||
assert!(extract_pointer(&capture).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_core_no_longer_owns_memory_extraction_scheduler() {
|
||||
let worker_source = include_str!("../../worker.rs");
|
||||
for removed in [
|
||||
"spawn_post_run_memory_jobs",
|
||||
"run_extract_once_with_cancel_observer",
|
||||
"consolidation_in_flight",
|
||||
"extract_in_flight",
|
||||
"memory_task:",
|
||||
] {
|
||||
assert!(
|
||||
!worker_source.contains(removed),
|
||||
"Worker core still contains removed extraction scheduler symbol {removed}"
|
||||
);
|
||||
}
|
||||
let controller_source = include_str!("../../controller.rs");
|
||||
assert!(controller_source.contains("if feature_config.memory.enabled"));
|
||||
assert!(controller_source.contains("MemoryExtractionLifecycleFeature::new"));
|
||||
let internal_worker_source = include_str!("../../internal_worker.rs");
|
||||
assert!(!internal_worker_source.contains("manifest.memory = None"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_capture_never_schedules_extraction() {
|
||||
let capture = capture(0, 500);
|
||||
let mut config = manifest::MemoryConfig::default();
|
||||
config.extract_threshold = Some(1);
|
||||
assert!(!extraction_threshold_reached(&capture, None, &config));
|
||||
}
|
||||
}
|
||||
+17
-19
@@ -28,7 +28,7 @@ const FINISH_DESCRIPTION: &str =
|
||||
"Finish Memory extraction after validating the number of candidates staged during this run.";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MemoryExtractState {
|
||||
pub(crate) struct MemoryStagingOutputState {
|
||||
view: Arc<SessionCapture>,
|
||||
workspace_client: Arc<dyn WorkspaceClient>,
|
||||
source: SourceRef,
|
||||
@@ -37,7 +37,7 @@ pub(crate) struct MemoryExtractState {
|
||||
finished: Arc<Mutex<Option<FinishMemoryExtractionParams>>>,
|
||||
}
|
||||
|
||||
impl MemoryExtractState {
|
||||
impl MemoryStagingOutputState {
|
||||
pub(crate) fn new(
|
||||
view: SessionCapture,
|
||||
workspace_client: Arc<dyn WorkspaceClient>,
|
||||
@@ -70,22 +70,20 @@ impl MemoryExtractState {
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct MemoryExtractFeature {
|
||||
state: MemoryExtractState,
|
||||
pub(crate) struct MemoryStagingOutputFeature {
|
||||
state: MemoryStagingOutputState,
|
||||
}
|
||||
|
||||
impl MemoryExtractFeature {
|
||||
pub(crate) fn new(state: MemoryExtractState) -> Self {
|
||||
impl MemoryStagingOutputFeature {
|
||||
pub(crate) fn new(state: MemoryStagingOutputState) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
impl FeatureModule for MemoryExtractFeature {
|
||||
impl FeatureModule for MemoryStagingOutputFeature {
|
||||
fn descriptor(&self) -> FeatureDescriptor {
|
||||
FeatureDescriptor::builtin("memory-extract", "Memory Extract")
|
||||
.with_description(
|
||||
"Memory staging and extraction completion, independent from session exploration.",
|
||||
)
|
||||
FeatureDescriptor::builtin("memory-staging-output", "Memory Staging Output")
|
||||
.with_description("Restricted Memory staging output for an extraction Internal Worker.")
|
||||
.with_tool(ToolDeclaration::new(
|
||||
"StageMemoryCandidate",
|
||||
STAGE_DESCRIPTION,
|
||||
@@ -109,7 +107,7 @@ impl FeatureModule for MemoryExtractFeature {
|
||||
}
|
||||
}
|
||||
|
||||
fn stage_definition(state: MemoryExtractState) -> ToolDefinition {
|
||||
fn stage_definition(state: MemoryStagingOutputState) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(StageMemoryCandidateParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
@@ -123,7 +121,7 @@ fn stage_definition(state: MemoryExtractState) -> ToolDefinition {
|
||||
})
|
||||
}
|
||||
|
||||
fn finish_definition(state: MemoryExtractState) -> ToolDefinition {
|
||||
fn finish_definition(state: MemoryStagingOutputState) -> ToolDefinition {
|
||||
Arc::new(move || {
|
||||
let schema = serde_json::to_value(schemars::schema_for!(FinishMemoryExtractionParams))
|
||||
.unwrap_or_else(|_| serde_json::json!({}));
|
||||
@@ -157,7 +155,7 @@ struct FinishMemoryExtractionParams {
|
||||
}
|
||||
|
||||
struct StageMemoryCandidateTool {
|
||||
state: MemoryExtractState,
|
||||
state: MemoryStagingOutputState,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -252,7 +250,7 @@ impl Tool for StageMemoryCandidateTool {
|
||||
}
|
||||
|
||||
struct FinishMemoryExtractionTool {
|
||||
state: MemoryExtractState,
|
||||
state: MemoryStagingOutputState,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -431,8 +429,8 @@ mod tests {
|
||||
|
||||
use super::*;
|
||||
|
||||
fn state() -> MemoryExtractState {
|
||||
MemoryExtractState::new(
|
||||
fn state() -> MemoryStagingOutputState {
|
||||
MemoryStagingOutputState::new(
|
||||
SessionCapture::new("segment-1", vec![Item::user_message("durable decision")]),
|
||||
crate::worker::marker_workspace_client(None, "test-backend"),
|
||||
SourceRef {
|
||||
@@ -445,8 +443,8 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn memory_extract_declares_only_memory_mutation_tools() {
|
||||
let descriptor = MemoryExtractFeature::new(state()).descriptor();
|
||||
assert_eq!(descriptor.id.as_str(), "builtin:memory-extract");
|
||||
let descriptor = MemoryStagingOutputFeature::new(state()).descriptor();
|
||||
assert_eq!(descriptor.id.as_str(), "builtin:memory-staging-output");
|
||||
assert_eq!(
|
||||
descriptor
|
||||
.tools
|
||||
@@ -0,0 +1,111 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use agen::{HistoryEntry, UsageRecord};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::session_history::SessionHistoryMetadata;
|
||||
|
||||
/// Immutable projection of one durably committed session-log location.
|
||||
///
|
||||
/// Feature code receives this value only after the host has committed the
|
||||
/// terminal run record. The projection deliberately carries annotated history
|
||||
/// rather than the public flattened transcript so provenance-sensitive
|
||||
/// features can construct their own bounded views.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CommittedSessionCapture {
|
||||
pub(crate) session_id: String,
|
||||
pub(crate) segment_id: String,
|
||||
/// Monotonic committed-log revision for the captured Segment.
|
||||
pub(crate) session_revision: u64,
|
||||
pub(crate) entry_count: usize,
|
||||
pub(crate) history: Vec<HistoryEntry<SessionHistoryMetadata>>,
|
||||
pub(crate) usage_history: Vec<UsageRecord>,
|
||||
pub(crate) extensions: Vec<(String, Value)>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub(crate) struct CommittedSessionLocation {
|
||||
pub(crate) session_id: String,
|
||||
pub(crate) segment_id: String,
|
||||
/// Monotonic committed-log revision for the captured Segment.
|
||||
pub(crate) session_revision: u64,
|
||||
pub(crate) entry_count: usize,
|
||||
}
|
||||
|
||||
impl CommittedSessionCapture {
|
||||
pub(crate) fn location(&self) -> CommittedSessionLocation {
|
||||
CommittedSessionLocation {
|
||||
session_id: self.session_id.clone(),
|
||||
segment_id: self.segment_id.clone(),
|
||||
session_revision: self.session_revision,
|
||||
entry_count: self.entry_count,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum FeatureSessionError {
|
||||
#[error("read committed session failed: {0}")]
|
||||
Capture(String),
|
||||
#[error("append session extension failed: {0}")]
|
||||
Extension(String),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CommittedSessionCaptureHandle {
|
||||
capture: Arc<
|
||||
dyn Fn() -> Result<CommittedSessionCapture, FeatureSessionError> + Send + Sync + 'static,
|
||||
>,
|
||||
}
|
||||
|
||||
impl CommittedSessionCaptureHandle {
|
||||
pub(crate) fn new(
|
||||
capture: impl Fn() -> Result<CommittedSessionCapture, FeatureSessionError>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
capture: Arc::new(capture),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn capture(&self) -> Result<CommittedSessionCapture, FeatureSessionError> {
|
||||
(self.capture)()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct SessionExtensionHandle {
|
||||
append: Arc<
|
||||
dyn Fn(&CommittedSessionLocation, &str, Value) -> Result<bool, FeatureSessionError>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
>,
|
||||
}
|
||||
|
||||
impl SessionExtensionHandle {
|
||||
pub(crate) fn new(
|
||||
append: impl Fn(&CommittedSessionLocation, &str, Value) -> Result<bool, FeatureSessionError>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
append: Arc::new(append),
|
||||
}
|
||||
}
|
||||
|
||||
/// Appends an extension only while the committed session is still at the
|
||||
/// exact location captured by the feature. `Ok(false)` is a stale-write
|
||||
/// fence, not an I/O failure.
|
||||
pub(crate) fn append_if_current(
|
||||
&self,
|
||||
expected: &CommittedSessionLocation,
|
||||
domain: &str,
|
||||
payload: Value,
|
||||
) -> Result<bool, FeatureSessionError> {
|
||||
(self.append)(expected, domain, payload)
|
||||
}
|
||||
}
|
||||
@@ -123,14 +123,14 @@ where
|
||||
|
||||
// Internal identities are run-scoped and never enter the public Runtime Worker catalog.
|
||||
manifest.worker.name = format!("internal-{}-{}", identity.kind, identity.run_id);
|
||||
// Internal jobs only receive features supplied below. A parent manifest must not accidentally
|
||||
// grant its normal public tool surface or recursively schedule memory work.
|
||||
// Internal jobs only receive the explicitly supplied Feature set below. A
|
||||
// parent manifest cannot accidentally grant its normal public tool surface
|
||||
// or recursively schedule Feature-owned background work.
|
||||
manifest.feature = Default::default();
|
||||
manifest.plugins = Default::default();
|
||||
manifest.mcp = Default::default();
|
||||
manifest.skills = None;
|
||||
manifest.compaction = None;
|
||||
manifest.memory = None;
|
||||
|
||||
let last_usage = Arc::new(Mutex::new(None::<UsageEvent>));
|
||||
let usage_slot = last_usage.clone();
|
||||
@@ -548,7 +548,6 @@ pub(crate) async fn spawn_internal_worker_session(
|
||||
authority,
|
||||
} = spec;
|
||||
manifest.worker.name = format!("internal-{}-{}", identity.kind, identity.run_id);
|
||||
manifest.memory = None;
|
||||
|
||||
let last_usage = Arc::new(Mutex::new(None::<UsageEvent>));
|
||||
let usage_slot = last_usage.clone();
|
||||
@@ -649,7 +648,6 @@ pub(crate) fn prepare_internal_worker_from_spec(
|
||||
manifest.mcp = Default::default();
|
||||
manifest.skills = None;
|
||||
manifest.compaction = None;
|
||||
manifest.memory = None;
|
||||
|
||||
let mut engine =
|
||||
Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client)
|
||||
|
||||
+144
-1183
File diff suppressed because it is too large
Load Diff
@@ -578,138 +578,6 @@ async fn mid_turn_compact_success_broadcasts_start_and_done() {
|
||||
assert_eq!(new_id_in_event, Some(worker.segment_id()));
|
||||
}
|
||||
|
||||
/// Regression: `Worker::compact()` must reset the in-memory
|
||||
/// `extract_pointer` so extract keeps firing on the new compacted
|
||||
/// session.
|
||||
///
|
||||
/// Without the reset, the pointer's `processed_through_history_len`
|
||||
/// holds the old (typically large) item count, while the new compacted
|
||||
/// session starts with a much shorter history (`[summary, ...]`).
|
||||
/// `cumulative_input_tokens_since` would then filter every new
|
||||
/// usage record out (their `history_len` is below the stale pointer)
|
||||
/// and extract would never re-fire for the rest of the process.
|
||||
const EXTRACT_PLUS_COMPACT_MANIFEST: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
pwd = "./"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
workspace_id = "test-workspace"
|
||||
settings_revision = 1
|
||||
language = "English"
|
||||
extract_threshold = 1
|
||||
|
||||
[compaction]
|
||||
compact_threshold = 1
|
||||
compact_retained_tokens = 0
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
fn finish_memory_extraction_tool_use_events(call_id: &str) -> Vec<LlmEvent> {
|
||||
let input = serde_json::json!({
|
||||
"staged_count": 0,
|
||||
"no_candidates_reason": "test run has no durable candidates"
|
||||
})
|
||||
.to_string();
|
||||
vec![
|
||||
LlmEvent::tool_use_start(0, call_id, "FinishMemoryExtraction"),
|
||||
LlmEvent::tool_input_delta(0, input),
|
||||
LlmEvent::tool_use_stop(0),
|
||||
LlmEvent::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
}),
|
||||
]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compact_resets_extract_pointer_so_extract_can_fire_again() {
|
||||
// Mock LLM responses, in call order:
|
||||
// [0] first run with usage(1000) so extract threshold (=1) fires.
|
||||
// [1] extract worker invokes FinishMemoryExtraction with empty output.
|
||||
// [2] extract worker closes after the tool result.
|
||||
// [3] compact worker invokes write_summary.
|
||||
// [4] compact worker closes after the tool result.
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
write_summary_tool_use_events("sc1", "summary"),
|
||||
single_text_events("done"),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_PLUS_COMPACT_MANIFEST, client).await;
|
||||
|
||||
worker.run_text("first").await.unwrap();
|
||||
|
||||
// extract fires; pointer becomes Some.
|
||||
worker.try_post_run_extract().await.unwrap();
|
||||
assert!(
|
||||
worker.extract_pointer().is_some(),
|
||||
"extract_pointer should be Some after a successful extract"
|
||||
);
|
||||
|
||||
// Compact runs. Without the fix the in-memory pointer would still
|
||||
// reference the old Segment's history_len.
|
||||
worker.try_pre_run_compact().await;
|
||||
assert!(
|
||||
worker.extract_pointer().is_none(),
|
||||
"extract_pointer must be reset to None after compact (matches cold-restore on the new Segment)"
|
||||
);
|
||||
}
|
||||
|
||||
/// `extract_threshold = 0` is treated as "disabled" — without this, a
|
||||
/// raw `>=` comparison against `tokens_since` would fire extract on
|
||||
/// every post-run regardless of activity. Mirrors the consolidation
|
||||
/// zero-threshold convention so users have a single way to opt out
|
||||
/// without removing the `[memory]` section.
|
||||
const EXTRACT_THRESHOLD_ZERO_MANIFEST: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
pwd = "./"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
extract_threshold = 0
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn extract_threshold_zero_is_disabled() {
|
||||
// Mock provides exactly one response — the first run. If extract
|
||||
// were treated as "fire on any change" because of `tokens_since >= 0`,
|
||||
// it would call into the extract worker and exhaust the mock.
|
||||
let client = MockClient::new(vec![text_events_with_usage("hi", 1000)]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_THRESHOLD_ZERO_MANIFEST, client).await;
|
||||
|
||||
worker.run_text("first").await.unwrap();
|
||||
worker
|
||||
.try_post_run_extract()
|
||||
.await
|
||||
.expect("extract_threshold=0 must skip silently, not fail");
|
||||
assert!(
|
||||
worker.extract_pointer().is_none(),
|
||||
"no extract should have run — pointer must remain None"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_run_compact_failure_broadcasts_start_and_failed() {
|
||||
// Only the first run has a response. Compaction will run the
|
||||
@@ -746,112 +614,6 @@ async fn pre_run_compact_failure_broadcasts_start_and_failed() {
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Detached post-run memory jobs (`spawn_post_run_memory_jobs` /
|
||||
// `wait_for_memory_jobs`). Covers the detach round-trip and the structural
|
||||
// invariant that the cloned memory-task Worker shares `SegmentState` with the
|
||||
// source Worker, so that `save_extension` from the background extract does not
|
||||
// leave the next turn's `save_user_input` looking at a stale session pointer.
|
||||
|
||||
const EXTRACT_NO_COMPACT_MANIFEST: &str = r#"
|
||||
[worker]
|
||||
name = "test-worker"
|
||||
pwd = "./"
|
||||
|
||||
[model]
|
||||
scheme = "anthropic"
|
||||
model_id = "test-model"
|
||||
|
||||
[engine]
|
||||
max_tokens = 100
|
||||
|
||||
[memory]
|
||||
workspace_id = "test-workspace"
|
||||
settings_revision = 1
|
||||
language = "English"
|
||||
extract_threshold = 1
|
||||
|
||||
[[scope.allow]]
|
||||
target = "./"
|
||||
permission = "write"
|
||||
"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn extract_large_unprocessed_range_does_not_abort_on_input_occupancy() {
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("recorded", 1000),
|
||||
finish_memory_extraction_tool_use_events("ec-large"),
|
||||
single_text_events("done"),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
|
||||
|
||||
let large_request = format!("remember this large slice: {}", "x ".repeat(200_000));
|
||||
worker.run_text(&large_request).await.unwrap();
|
||||
|
||||
worker.try_post_run_extract().await.expect(
|
||||
"large unprocessed extract ranges must reach the extract worker, not abort locally",
|
||||
);
|
||||
assert!(
|
||||
worker.extract_pointer().is_some(),
|
||||
"successful extract should advance the pointer even when the input range is large"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn spawn_and_wait_drives_extract_to_completion() {
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
|
||||
|
||||
worker.run_text("first").await.unwrap();
|
||||
assert!(
|
||||
worker.extract_pointer().is_none(),
|
||||
"extract has not run yet — pointer must be None"
|
||||
);
|
||||
|
||||
worker.spawn_post_run_memory_jobs();
|
||||
worker.wait_for_memory_jobs().await;
|
||||
|
||||
assert!(
|
||||
worker.extract_pointer().is_some(),
|
||||
"spawn + wait must complete extract; pointer should be set"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn detached_extract_does_not_fork_session_log() {
|
||||
// Source worker and the cloned memory-task worker share `SegmentState` via
|
||||
// `Arc<_>`. The detached extract advances the entry tally through
|
||||
// `save_extension`; the next `run` must see that same tally so
|
||||
// `ensure_head_or_fork` does not spawn a new session.
|
||||
let client = MockClient::new(vec![
|
||||
text_events_with_usage("hi", 1000),
|
||||
finish_memory_extraction_tool_use_events("ec1"),
|
||||
single_text_events("done"),
|
||||
text_events_with_usage("ok", 1000),
|
||||
]);
|
||||
let mut worker = make_worker_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
|
||||
|
||||
worker.run_text("first").await.unwrap();
|
||||
let session_before = worker.segment_id();
|
||||
|
||||
worker.spawn_post_run_memory_jobs();
|
||||
worker.wait_for_memory_jobs().await;
|
||||
|
||||
worker.run_text("second").await.unwrap();
|
||||
let session_after = worker.segment_id();
|
||||
|
||||
assert_eq!(
|
||||
session_before, session_after,
|
||||
"detached extract's save_extension and the next turn's save_user_input \
|
||||
must share the entry tally through SegmentState — a fork here means the \
|
||||
clone carried its own counter"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn controller_compact_method_emits_start_and_done() {
|
||||
let client = MockClient::new(vec![
|
||||
|
||||
Reference in New Issue
Block a user