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

feat: Pos処理の非同期化・Busy状態の削除

Browse Source
This commit is contained in:
Hare
2026-05-04 15:52:27 +09:00
parent 5b38aa6a87
commit 954cf200e2
15 changed files with 537 additions and 256 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pod/src/compact/state.rs
+4 -4
View File
@@ -4,8 +4,8 @@
//! flags shared between:
//! - `PodInterceptor` (reads `request_threshold` — the *safety net* for
//! between-requests yielding)
//! - `Pod::try_post_run_compact` (reads `post_run_threshold` — the
//! *proactive* check between turns)
//! - `Pod::try_pre_run_compact` (reads `post_run_threshold` — the
//! *proactive* check before the next turn starts)
//! - `Pod::run()` / `resume()` (circuit breaker, thrash detection)
//!
//! Current occupancy (input-token count) is **not** stored here. The single
@@ -19,8 +19,8 @@ const MAX_COMPACT_FAILURES: usize = 3;
/// Shared mutable state for compaction decisions.
pub(crate) struct CompactState {
/// Between-turns threshold (proactive). Checked by the Controller
/// after a run completes. `None` disables the post-run check.
/// Between-turns threshold (proactive). Checked before the next turn
/// starts. `None` disables the pre-run check.
post_run_threshold: Option<u64>,
/// Between-requests threshold (safety net). Checked inside a turn
/// before each LLM request. `None` disables the request check.
crates/pod/src/controller.rs
+19 -64
View File
@@ -76,68 +76,22 @@ async fn set_controller_status(
let _ = event_tx.send(Event::Status { status });
}
async fn run_post_run_jobs<C, St>(pod: &mut Pod<C, St>, alerter: &Alerter)
where
C: LlmClient,
St: Store,
{
if let Err(e) = pod.try_post_run_extract().await {
tracing::warn!(error = %e, "Post-run memory extract error");
alerter.alert(
AlertLevel::Warn,
AlertSource::Pod,
format!("post-run memory extract error: {e}"),
);
}
if let Err(e) = pod.try_post_run_consolidate().await {
tracing::warn!(error = %e, "Post-run memory consolidate error");
alerter.alert(
AlertLevel::Warn,
AlertSource::Pod,
format!("post-run memory consolidate error: {e}"),
);
}
if let Err(e) = pod.try_post_run_compact().await {
tracing::warn!(error = %e, "Post-run compaction error");
alerter.alert(
AlertLevel::Warn,
AlertSource::Compactor,
format!("post-run compaction error: {e}"),
);
}
}
async fn finish_controller_run<C, St>(
pod: &mut Pod<C, St>,
shared_state: &Arc<PodSharedState>,
runtime_dir: &RuntimeDir,
event_tx: &broadcast::Sender<Event>,
alerter: &Alerter,
new_status: PodStatus,
) where
C: LlmClient,
St: Store,
C: LlmClient + Clone + 'static,
St: Store + Clone + 'static,
{
if new_status == PodStatus::Busy {
// Surface the post-run busy window before kicking off the jobs so
// TUI / external observers see Busy regardless of whether the
// worker turn ended via success or error. Both branches in
// `run_with_cancel_support` return `PodStatus::Busy` for this
// path; emitting here keeps the two unified.
set_controller_status(shared_state, runtime_dir, event_tx, PodStatus::Busy).await;
run_post_run_jobs(pod, alerter).await;
}
let items = pod.worker().history().to_vec();
shared_state.update_history(items);
shared_state.set_user_segments(pod.user_segments().to_vec());
let final_status = if new_status == PodStatus::Busy {
PodStatus::Idle
} else {
new_status
};
set_controller_status(shared_state, runtime_dir, event_tx, final_status).await;
set_controller_status(shared_state, runtime_dir, event_tx, new_status).await;
let _ = runtime_dir.write_history(shared_state).await;
pod.spawn_post_run_memory_jobs();
}
// ---------------------------------------------------------------------------
@@ -154,8 +108,8 @@ impl PodController {
runtime_base: &Path,
) -> Result<(PodHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + 'static,
St: Store + 'static,
C: LlmClient + Clone + 'static,
St: Store + Clone + 'static,
{
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let (method_tx, mut method_rx) = mpsc::channel::<Method>(32);
@@ -524,7 +478,6 @@ impl PodController {
&shared_state,
&runtime_dir,
&event_tx,
&alerter,
new_status,
)
.await;
@@ -542,9 +495,9 @@ impl PodController {
pod.push_notify(message);
let status = shared_state.get_status();
if status != PodStatus::Idle {
// RUNNING / Paused / Busy: the buffer push is the
// RUNNING / Paused: the buffer push is the
// entire operation; an in-flight turn (or the
// next Resume/Run after Busy) will drain the buffer
// next Resume/Run) will drain the buffer
// at its next pre_llm_request.
continue;
}
@@ -576,7 +529,6 @@ impl PodController {
&shared_state,
&runtime_dir,
&event_tx,
&alerter,
new_status,
)
.await;
@@ -621,7 +573,6 @@ impl PodController {
&shared_state,
&runtime_dir,
&event_tx,
&alerter,
new_status,
)
.await;
@@ -640,12 +591,12 @@ impl PodController {
}
Method::Pause => {
// Already paused or post-run busy → idempotent no-op.
// Otherwise the Pod is Idle (Running turns go through
// Already paused → idempotent no-op. Otherwise the
// Pod is Idle (Running turns go through
// `run_with_cancel_support`, not this outer match), so
// there is nothing to pause.
let status = shared_state.get_status();
if !matches!(status, PodStatus::Paused | PodStatus::Busy) {
if status != PodStatus::Paused {
let _ = event_tx.send(Event::Error {
code: ErrorCode::NotRunning,
message: "Pod is not running".into(),
@@ -714,7 +665,6 @@ impl PodController {
&shared_state,
&runtime_dir,
&event_tx,
&alerter,
new_status,
)
.await;
@@ -728,6 +678,11 @@ impl PodController {
}
}
// Background memory jobs own extract/consolidate workers after a
// turn completes. Join them before the controller task exits so
// staging writes and consolidation cleanups are not abandoned.
pod.wait_for_memory_jobs().await;
// Report upward that this Pod is stopping before the
// controller task exits. Awaited (not fire-and-forget):
// after `shutdown_tx.send` the process may exit quickly,
@@ -787,9 +742,9 @@ where
return match result {
Ok(r) => {
let (status, run_result) = match r {
PodRunResult::Finished => (PodStatus::Busy, RunResult::Finished),
PodRunResult::Finished => (PodStatus::Idle, RunResult::Finished),
PodRunResult::Paused => (PodStatus::Paused, RunResult::Paused),
PodRunResult::LimitReached => (PodStatus::Busy, RunResult::LimitReached),
PodRunResult::LimitReached => (PodStatus::Idle, RunResult::LimitReached),
};
let _ = event_tx.send(Event::RunEnd { result: run_result });
if matches!(run_result, RunResult::Finished) {
@@ -825,7 +780,7 @@ where
message,
},
);
(PodStatus::Busy, shutdown_requested)
(PodStatus::Idle, shutdown_requested)
}
};
}
crates/pod/src/pod.rs
+259 -85
View File
@@ -1,6 +1,7 @@
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::Mutex as AsyncMutex;
use llm_worker::Item;
use llm_worker::llm_client::RequestConfig;
@@ -35,6 +36,12 @@ use async_trait::async_trait;
use llm_worker::interceptor::PreRequestAction;
use protocol::{AlertLevel, AlertSource, Event, Segment};
use tokio::sync::broadcast;
use tokio::task::JoinHandle;
struct SessionHead {
session_id: SessionId,
head_hash: Option<EntryHash>,
}
/// Pre-LLM-request hook that records `history.len()` at send time into a
/// shared `UsageTracker`. The on_usage callback later pairs this with the
@@ -61,7 +68,7 @@ pub struct Pod<C: LlmClient, St: Store> {
worker: Option<Worker<C, Mutable>>,
store: St,
session_id: SessionId,
head_hash: Option<EntryHash>,
session_head: Arc<AsyncMutex<SessionHead>>,
/// Absolute working directory of the Pod.
pwd: PathBuf,
/// Shared, atomically-swappable view of the Pod's resolved scope.
@@ -172,7 +179,13 @@ pub struct Pod<C: LlmClient, St: Store> {
/// run yet on this session — next extract starts from entry 0.
/// Restored from `RestoredState.extensions` on `restore`, updated
/// after each successful extract via `save_extension`.
extract_pointer: Mutex<Option<memory::ExtractPointerPayload>>,
extract_pointer: Arc<Mutex<Option<memory::ExtractPointerPayload>>>,
/// Phase 1/2 memory job running outside the controller method loop.
/// The task owns the extract/consolidate worker execution and is joined
/// at shutdown. A single slot is enough: Phase 1/2 implementations loop
/// until thresholds fall below their trigger points, and concurrent
/// triggers are coalesced by skipping when this handle is still active.
memory_task: Option<JoinHandle<()>>,
/// Typed user submissions in submit order. K-th entry corresponds to
/// the K-th `Item::user_message` in `worker.history()` (modulo seed
/// history loaded via `SessionStart.history`, whose original segments
@@ -182,6 +195,84 @@ pub struct Pod<C: LlmClient, St: Store> {
user_segments: Vec<Vec<Segment>>,
}
impl<C: LlmClient + 'static, St: Store + 'static> Pod<C, St> {
pub async fn wait_for_memory_jobs(&mut self) {
if let Some(handle) = self.memory_task.take()
&& let Err(e) = handle.await
{
tracing::warn!(error = %e, "Post-run memory task join failed");
}
}
}
impl<C: LlmClient + Clone + 'static, St: Store + Clone + 'static> Pod<C, St> {
fn clone_for_memory_task(&self) -> Self {
// The cloned Pod's worker exists only as a snapshot for the memory
// task: `run_extract_once` reads `worker.history()`, and the
// extract/consolidate workers are built fresh inside their own
// methods using `worker.client()` as fallback when no override
// model is configured. system_prompt / request_config / cache_key
// are unused on this path, so we deliberately skip copying them.
let source_worker = self.worker.as_ref().expect("worker present");
let mut worker = Worker::new(source_worker.client().clone());
worker.set_history(source_worker.history().to_vec());
Self {
manifest: self.manifest.clone(),
worker: Some(worker),
store: self.store.clone(),
session_id: self.session_id,
session_head: self.session_head.clone(),
pwd: self.pwd.clone(),
scope: self.scope.clone(),
hook_builder: HookRegistryBuilder::new(),
interceptor_installed: false,
compact_state: None,
usage_tracker: Arc::new(UsageTracker::new()),
metrics_tracker: Arc::new(crate::compact::metrics_tracker::MetricsTracker::new()),
usage_history: self.usage_history.clone(),
tracker: None,
task_store: self.task_store.clone(),
system_prompt_template: None,
alerter: self.alerter.clone(),
event_tx: self.event_tx.clone(),
pending_notifies: NotifyBuffer::new(),
pending_attachments: Arc::new(Mutex::new(Vec::new())),
scope_allocation: None,
callback_socket: None,
prompts: self.prompts.clone(),
workflow_registry: self.workflow_registry.clone(),
memory_layout: self.memory_layout.clone(),
inject_resident_knowledge: self.inject_resident_knowledge,
pending_scope_snapshot: self.pending_scope_snapshot.clone(),
extract_in_flight: self.extract_in_flight.clone(),
consolidation_in_flight: self.consolidation_in_flight.clone(),
extract_pointer: self.extract_pointer.clone(),
memory_task: None,
user_segments: self.user_segments.clone(),
}
}
pub fn spawn_post_run_memory_jobs(&mut self) {
// Drop a finished prior handle so we can spawn a fresh task.
// If the prior task is still running, coalesce by skipping —
// Phase 1/2 implementations re-evaluate thresholds on completion.
self.cleanup_finished_memory_task();
if self.memory_task.is_some() {
return;
}
let mut pod = self.clone_for_memory_task();
self.memory_task = Some(tokio::spawn(async move {
if let Err(e) = pod.try_post_run_extract().await {
tracing::warn!(error = %e, "Post-run memory extract task error");
}
if let Err(e) = pod.try_post_run_consolidate().await {
tracing::warn!(error = %e, "Post-run memory consolidate task error");
}
}));
}
}
impl<C: LlmClient, St: Store> Pod<C, St> {
/// Create a new Pod from a pre-built Worker and store.
///
@@ -210,7 +301,10 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
worker: Some(worker),
store,
session_id,
head_hash: None,
session_head: Arc::new(AsyncMutex::new(SessionHead {
session_id,
head_hash: None,
})),
pwd,
scope: SharedScope::new(scope),
hook_builder: HookRegistryBuilder::new(),
@@ -235,7 +329,8 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
pending_scope_snapshot: Arc::new(Mutex::new(None)),
extract_in_flight: Arc::new(AtomicBool::new(false)),
consolidation_in_flight: Arc::new(AtomicBool::new(false)),
extract_pointer: Mutex::new(None),
extract_pointer: Arc::new(Mutex::new(None)),
memory_task: None,
user_segments: Vec::new(),
};
pod.apply_prune_from_manifest();
@@ -330,7 +425,8 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// can restore the narrowed scope instead of reclaiming delegated
/// writes.
pub async fn persist_scope_snapshot(&mut self) -> Result<(), StoreError> {
if self.head_hash.is_none() {
let mut head = self.session_head.lock().await;
if head.head_hash.is_none() {
return Ok(());
}
let snapshot = {
@@ -340,7 +436,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
deny: scope.deny_rules(),
}
};
session_store::save_pod_scope(&self.store, self.session_id, &mut self.head_hash, &snapshot)
session_store::save_pod_scope(&self.store, head.session_id, &mut head.head_hash, &snapshot)
.await
}
@@ -362,10 +458,11 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
.expect("pending_scope_snapshot poisoned")
.take();
if let Some(snapshot) = snapshot {
let mut head = self.session_head.lock().await;
session_store::save_pod_scope(
&self.store,
self.session_id,
&mut self.head_hash,
head.session_id,
&mut head.head_hash,
&snapshot,
)
.await?;
@@ -531,10 +628,11 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// (the entry is dropped) and a `Warn` alert + `tracing::warn!` are
/// emitted so the failure isn't completely silent.
async fn try_record_metric(&mut self, metric: &session_metrics::Metric) {
let mut head = self.session_head.lock().await;
if let Err(err) = session_metrics::record_metric(
&self.store,
self.session_id,
&mut self.head_hash,
head.session_id,
&mut head.head_hash,
metric,
)
.await
@@ -803,6 +901,52 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
self.run(vec![Segment::text(s)]).await
}
/// Drop the prior memory_task handle if it has finished. Keep it if
/// still running so callers can decide whether to wait or coalesce.
fn cleanup_finished_memory_task(&mut self) {
if self.memory_task.as_ref().is_some_and(|h| h.is_finished()) {
self.memory_task = None;
}
}
/// Wait for the in-flight memory task (if any) to finish. Used before
/// compact rewrites history (extract reads the same history).
async fn join_memory_task(&mut self) {
if let Some(handle) = self.memory_task.take()
&& let Err(e) = handle.await
{
tracing::warn!(error = %e, "Memory task join failed");
}
}
/// Whether `try_pre_run_compact` would actually compact. The same
/// check is duplicated inside `try_pre_run_compact` itself for
/// defensive reasons; this is the gate for joining the memory task
/// before the compact runs.
fn should_pre_run_compact(&self) -> bool {
self.compact_state.as_ref().is_some_and(|s| {
!s.is_disabled()
&& !s.just_compacted()
&& s.exceeds_post_run(self.total_tokens().tokens)
})
}
/// Prelude shared by `run` / `run_for_notification` / `resume`.
/// Wires up worker hooks, ensures the session is materialized on the
/// store, and runs pre-run compact (joining any in-flight memory task
/// first so extract sees a stable history range).
async fn prepare_for_run(&mut self) -> Result<(), PodError> {
self.ensure_interceptor_installed();
self.ensure_system_prompt_materialized()?;
self.cleanup_finished_memory_task();
self.ensure_session_head().await?;
if self.should_pre_run_compact() {
self.join_memory_task().await;
}
self.try_pre_run_compact().await;
Ok(())
}
/// Send user input and run until the LLM turn completes.
///
/// `input` is a typed segment list (see [`protocol::Segment`]). The
@@ -816,20 +960,22 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// the Worker is aborted, history is compacted, and execution resumes
/// automatically.
pub async fn run(&mut self, input: Vec<Segment>) -> Result<PodRunResult, PodError> {
self.ensure_interceptor_installed();
self.ensure_system_prompt_materialized()?;
self.ensure_session_head().await?;
self.prepare_for_run().await?;
// Persist the user input as typed segments before the worker
// pushes its flattened copy into history. save_delta deliberately
// skips the resulting `is_user_message()` item to avoid double-write.
session_store::save_user_input(
&self.store,
self.session_id,
&mut self.head_hash,
input.clone(),
)
.await?;
{
let mut head = self.session_head.lock().await;
self.session_id = head.session_id;
session_store::save_user_input(
&self.store,
head.session_id,
&mut head.head_hash,
input.clone(),
)
.await?;
}
self.user_segments.push(input.clone());
// Resolve `@<path>` refs and `/<slug>` workflow invocations to
@@ -989,9 +1135,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// Worker's resume path issues the LLM request without a new
/// user turn.
pub async fn run_for_notification(&mut self) -> Result<PodRunResult, PodError> {
self.ensure_interceptor_installed();
self.ensure_system_prompt_materialized()?;
self.ensure_session_head().await?;
self.prepare_for_run().await?;
let history_before = self.worker.as_ref().unwrap().history().len();
@@ -1005,9 +1149,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// Resume from a paused state.
pub async fn resume(&mut self) -> Result<PodRunResult, PodError> {
self.ensure_interceptor_installed();
self.ensure_system_prompt_materialized()?;
self.ensure_session_head().await?;
self.prepare_for_run().await?;
let history_before = self.worker.as_ref().unwrap().history().len();
@@ -1035,27 +1177,29 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
config: w.request_config(),
history: w.history(),
};
if self.head_hash.is_none() {
let mut head = self.session_head.lock().await;
if head.head_hash.is_none() {
let hash =
session_store::create_session_with_id(&self.store, self.session_id, state).await?;
self.head_hash = Some(hash);
session_store::create_session_with_id(&self.store, head.session_id, state).await?;
head.head_hash = Some(hash);
drop(head);
self.persist_scope_snapshot().await?;
return Ok(());
}
let prev_session_id = self.session_id;
session_store::ensure_head_or_fork(
&self.store,
&mut self.session_id,
&mut self.head_hash,
state,
)
.await?;
let prev_session_id = head.session_id;
let mut session_id = head.session_id;
let mut head_hash = head.head_hash.clone();
session_store::ensure_head_or_fork(&self.store, &mut session_id, &mut head_hash, state)
.await?;
head.session_id = session_id;
head.head_hash = head_hash;
self.session_id = session_id;
// ensure_head_or_fork mints a fresh session_id when it auto-
// forks. Sync that to pods.json so a concurrent
// restore_from_manifest can't see "no live writer" for the new
// session and grab it.
if self.session_id != prev_session_id && self.scope_allocation.is_some() {
pod_registry::update_session(&self.manifest.pod.name, self.session_id)?;
if session_id != prev_session_id && self.scope_allocation.is_some() {
pod_registry::update_session(&self.manifest.pod.name, session_id)?;
}
Ok(())
}
@@ -1142,17 +1286,21 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
})
}
/// Attempt proactive compaction (called by Controller after run).
/// Attempt proactive compaction at the beginning of a controller Run.
///
/// Best-effort: failures are logged but do not propagate.
pub async fn try_post_run_compact(&mut self) -> Result<(), PodError> {
/// This used to run in the controller's post-run path. Keeping it here
/// preserves the ordering requirement that the next turn starts with a
/// compacted history, without introducing a separate Busy controller state.
/// Best-effort: failures are logged and surfaced, but do not abort the
/// user turn that triggered the check.
pub async fn try_pre_run_compact(&mut self) {
let state = match self.compact_state.as_ref() {
Some(s) if !s.is_disabled() && !s.just_compacted() => s.clone(),
_ => return Ok(()),
_ => return,
};
let current_tokens = self.total_tokens().tokens;
if !state.exceeds_post_run(current_tokens) {
return Ok(());
return;
}
let retained = state.retained_tokens();
@@ -1161,24 +1309,22 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
Ok(new_session_id) => {
info!(
new_session_id = %new_session_id,
"Proactive post-run compaction succeeded"
"Proactive pre-run compaction succeeded"
);
self.send_event(Event::CompactDone { new_session_id });
state.record_compact_success();
Ok(())
}
Err(e) => {
warn!(error = %e, "Proactive post-run compaction failed");
warn!(error = %e, "Proactive pre-run compaction failed");
self.send_event(Event::CompactFailed {
error: e.to_string(),
});
self.alert(
AlertLevel::Warn,
AlertSource::Compactor,
format!("post-run compaction failed: {e}"),
format!("pre-run compaction failed: {e}"),
);
state.record_compact_failure();
Ok(())
}
}
}
@@ -1193,19 +1339,24 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
// head_hash mutable).
let w = self.worker.as_ref().unwrap();
let new_items = &w.history()[history_before..];
session_store::save_delta(&self.store, self.session_id, &mut self.head_hash, new_items)
let mut head = self.session_head.lock().await;
self.session_id = head.session_id;
session_store::save_delta(&self.store, head.session_id, &mut head.head_hash, new_items)
.await?;
drop(head);
self.flush_pending_scope_snapshot().await?;
let turn_count = self.worker.as_ref().unwrap().turn_count();
let mut head = self.session_head.lock().await;
session_store::save_turn_end(
&self.store,
self.session_id,
&mut self.head_hash,
head.session_id,
&mut head.head_hash,
turn_count,
)
.await?;
drop(head);
// Flush any sync-buffered metrics from this run first
// (currently `prune.fire` / `prune.skip` from the prune observer).
@@ -1238,10 +1389,11 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
record,
correlation_id,
} = recorded;
let mut head = self.session_head.lock().await;
session_store::save_usage(
&self.store,
self.session_id,
&mut self.head_hash,
head.session_id,
&mut head.head_hash,
record.history_len,
record.input_total_tokens,
record.cache_read_tokens,
@@ -1249,6 +1401,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
record.output_tokens,
)
.await?;
drop(head);
if let Some(id) = correlation_id {
let metric = session_metrics::Metric::now("prune.post_request")
.with_correlation_id(&id)
@@ -1266,20 +1419,22 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
let interrupted = self.worker.as_ref().unwrap().last_run_interrupted();
match result {
Ok(r) => {
let mut head = self.session_head.lock().await;
session_store::save_run_completed(
&self.store,
self.session_id,
&mut self.head_hash,
head.session_id,
&mut head.head_hash,
r.clone(),
interrupted,
)
.await?;
}
Err(e) => {
let mut head = self.session_head.lock().await;
session_store::save_run_errored(
&self.store,
self.session_id,
&mut self.head_hash,
head.session_id,
&mut head.head_hash,
e.to_string(),
interrupted,
)
@@ -1511,8 +1666,9 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
));
// Persist as a new compacted session.
let old_session_id = self.session_id;
let old_head_hash = self
let mut head = self.session_head.lock().await;
let old_session_id = head.session_id;
let old_head_hash = head
.head_hash
.clone()
.expect("head_hash should be set after at least one entry");
@@ -1535,7 +1691,8 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
// session — the new compacted session starts with no measurements
// until its first LLM call.
self.session_id = new_session_id;
self.head_hash = Some(new_head_hash);
head.session_id = new_session_id;
head.head_hash = Some(new_head_hash);
// Keep pods.json pointing at the live session_id. Without this
// a concurrent `restore_from_manifest(new_session_id)` would
// see no live writer and grab the session this Pod just moved
@@ -1545,6 +1702,7 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
if self.scope_allocation.is_some() {
pod_registry::update_session(&self.manifest.pod.name, new_session_id)?;
}
drop(head);
// Align user_segments with the post-compaction history. Items
// before `retain_from` (now folded into the summary) lose their
// segments; only the user_messages surviving in retained_items
@@ -1641,10 +1799,10 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// Phase 1 (memory.extract) post-run trigger.
///
/// Called by the Controller **before** [`try_post_run_compact`] so
/// the extract worker sees a stable session-log entry range
/// (compact rewrites history). Best-effort: failures are logged but
/// not propagated.
/// Called by the Controller before spawning the background memory task so
/// the extract worker sees a stable session-log entry range while compact
/// is deferred until the next turn starts. Best-effort: failures are
/// logged but not propagated.
///
/// Behaviour follows `docs/plan/memory.md` §Phase 1 並走防止:
/// in-flight 中の trigger は skip し、完了時点で閾値再評価する
@@ -1798,11 +1956,12 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
extract::ExtractedPayload::default()
});
let source_session_id = self.session_head.lock().await.session_id;
let staging_id = if payload.is_empty() {
String::new()
} else {
let source = memory::schema::SourceRef {
session_id: self.session_id.to_string(),
session_id: source_session_id.to_string(),
range: [start_entry as u64, end_entry as u64],
};
let (id, _) = extract::write_staging(&layout, source, payload)
@@ -1817,14 +1976,18 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
};
let payload_value = serde_json::to_value(&pointer_payload)
.expect("ExtractPointerPayload is always JSON-serializable");
session_store::save_extension(
&self.store,
self.session_id,
&mut self.head_hash,
extract::EXTRACT_DOMAIN,
payload_value,
)
.await?;
{
let mut head = self.session_head.lock().await;
session_store::save_extension(
&self.store,
head.session_id,
&mut head.head_hash,
extract::EXTRACT_DOMAIN,
payload_value,
)
.await?;
self.session_id = head.session_id;
}
*self
.extract_pointer
@@ -1850,12 +2013,11 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
Ok(worker.client().clone_boxed())
}
/// Phase 2 (memory.consolidation) post-run trigger.
/// Phase 2 (memory.consolidation) trigger.
///
/// Called by the Controller **after** [`try_post_run_extract`] and
/// **before** [`try_post_run_compact`]: extract feeds staging, compact
/// rewrites history. Phase 2 must consume staging before compact
/// reshapes the session.
/// Intended to run from a background memory task after Phase 1 may have
/// added staging entries. Compact is deferred until the next turn starts,
/// so consolidation no longer blocks the controller's post-run path.
///
/// Behaviour follows `docs/plan/memory.md` §Phase 2 / §並走防止:
/// the staging-side `StagingLock` enforces cross-process exclusion;
@@ -2096,7 +2258,10 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
worker: Some(worker),
store,
session_id,
head_hash: None,
session_head: Arc::new(AsyncMutex::new(SessionHead {
session_id,
head_hash: None,
})),
pwd: common.pwd,
scope: SharedScope::new(common.scope),
hook_builder: HookRegistryBuilder::new(),
@@ -2121,7 +2286,8 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
pending_scope_snapshot: Arc::new(Mutex::new(None)),
extract_in_flight: Arc::new(AtomicBool::new(false)),
consolidation_in_flight: Arc::new(AtomicBool::new(false)),
extract_pointer: Mutex::new(None),
extract_pointer: Arc::new(Mutex::new(None)),
memory_task: None,
user_segments: Vec::new(),
};
pod.apply_prune_from_manifest();
@@ -2160,7 +2326,10 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
worker: Some(worker),
store,
session_id,
head_hash: None,
session_head: Arc::new(AsyncMutex::new(SessionHead {
session_id,
head_hash: None,
})),
pwd: common.pwd,
scope: SharedScope::new(common.scope),
hook_builder: HookRegistryBuilder::new(),
@@ -2185,7 +2354,8 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
pending_scope_snapshot: Arc::new(Mutex::new(None)),
extract_in_flight: Arc::new(AtomicBool::new(false)),
consolidation_in_flight: Arc::new(AtomicBool::new(false)),
extract_pointer: Mutex::new(None),
extract_pointer: Arc::new(Mutex::new(None)),
memory_task: None,
user_segments: Vec::new(),
};
pod.apply_prune_from_manifest();
@@ -2288,7 +2458,10 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
worker: Some(worker),
store,
session_id,
head_hash: state.head_hash,
session_head: Arc::new(AsyncMutex::new(SessionHead {
session_id,
head_hash: state.head_hash,
})),
pwd: common.pwd,
scope: SharedScope::new(common.scope),
hook_builder: HookRegistryBuilder::new(),
@@ -2315,7 +2488,8 @@ impl<St: Store> Pod<Box<dyn LlmClient>, St> {
pending_scope_snapshot: Arc::new(Mutex::new(None)),
extract_in_flight: Arc::new(AtomicBool::new(false)),
consolidation_in_flight: Arc::new(AtomicBool::new(false)),
extract_pointer: Mutex::new(extract_pointer),
extract_pointer: Arc::new(Mutex::new(extract_pointer)),
memory_task: None,
user_segments: state.user_segments,
};
pod.apply_prune_from_manifest();
crates/pod/tests/compact_events_test.rs
+90 -8
View File
@@ -1,8 +1,8 @@
//! Compact lifecycle `Event` broadcasting.
//!
//! Covers three paths:
//! - `try_post_run_compact` success → `CompactStart + CompactDone`
//! - `try_post_run_compact` failure → `CompactStart + CompactFailed`
//! - `try_pre_run_compact` success → `CompactStart + CompactDone`
//! - `try_pre_run_compact` failure → `CompactStart + CompactFailed`
//! - mid-turn `do_compact_and_resume` success → `CompactStart + CompactDone`
//! (driven by `compact_request_threshold` → `PreRequestAction::Yield`)
@@ -96,7 +96,7 @@ fn write_summary_tool_use_events(call_id: &str, text: &str) -> Vec<LlmEvent> {
]
}
// A low compact_threshold guarantees `try_post_run_compact` will fire
// A low compact_threshold guarantees `try_pre_run_compact` will fire
// the first time we check after a run.
const POST_RUN_MANIFEST_TOML: &str = r#"
[pod]
@@ -228,7 +228,7 @@ async fn compact_broadcasts_only_new_system_messages_not_retained_ones() {
}
#[tokio::test]
async fn post_run_compact_success_broadcasts_start_and_done() {
async fn pre_run_compact_success_broadcasts_start_and_done() {
// Responses: (1) first run returns short text, (2) compact worker
// emits write_summary then closes (two LLM calls inside the compact
// worker: one for write_summary, one that the compact loop consumes
@@ -247,7 +247,7 @@ async fn post_run_compact_success_broadcasts_start_and_done() {
// Drain run events so only compact events remain in `rx`.
let _ = drain(&mut rx);
pod.try_post_run_compact().await.unwrap();
pod.try_pre_run_compact().await;
let events = drain(&mut rx);
let kinds: Vec<&str> = events
@@ -412,7 +412,7 @@ async fn compact_resets_extract_pointer_so_phase1_can_fire_again() {
// Compact runs. Without the fix the in-memory pointer would still
// reference the old session's history_len.
pod.try_post_run_compact().await.unwrap();
pod.try_pre_run_compact().await;
assert!(
pod.extract_pointer().is_none(),
"extract_pointer must be reset to None after compact (matches cold-restore on the new session)"
@@ -463,7 +463,7 @@ async fn extract_threshold_zero_is_disabled() {
}
#[tokio::test]
async fn post_run_compact_failure_broadcasts_start_and_failed() {
async fn pre_run_compact_failure_broadcasts_start_and_failed() {
// Only the first run has a response. Compaction will run the
// compact worker which immediately exhausts the mock → failure.
let client = MockClient::new(vec![single_text_events("hi")]);
@@ -476,7 +476,7 @@ async fn post_run_compact_failure_broadcasts_start_and_failed() {
let _ = drain(&mut rx);
// Best-effort: returns Ok(()) even on failure, but emits CompactFailed.
pod.try_post_run_compact().await.unwrap();
pod.try_pre_run_compact().await;
let events = drain(&mut rx);
let kinds: Vec<&str> = events
@@ -497,3 +497,85 @@ async fn post_run_compact_failure_broadcasts_start_and_failed() {
"unexpected CompactDone in {kinds:?}"
);
}
// ---------------------------------------------------------------------------
// 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 Pod shares `SessionHead` with the
// source Pod, so that `save_extension` from the background extract does not
// leave the next turn's `save_user_input` looking at a stale head_hash.
const EXTRACT_NO_COMPACT_MANIFEST: &str = r#"
[pod]
name = "test-pod"
pwd = "./"
[model]
scheme = "anthropic"
model_id = "test-model"
[worker]
max_tokens = 100
[memory]
extract_threshold = 1
[[scope.allow]]
target = "./"
permission = "write"
"#;
#[tokio::test]
async fn spawn_and_wait_drives_extract_to_completion() {
let client = MockClient::new(vec![
text_events_with_usage("hi", 1000),
write_extracted_tool_use_events("ec1"),
single_text_events("done"),
]);
let mut pod = make_pod_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
pod.run_text("first").await.unwrap();
assert!(
pod.extract_pointer().is_none(),
"extract has not run yet — pointer must be None"
);
pod.spawn_post_run_memory_jobs();
pod.wait_for_memory_jobs().await;
assert!(
pod.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 pod and the cloned memory-task pod share `SessionHead` via
// `Arc<AsyncMutex<_>>`. The detached extract advances head_hash through
// `save_extension`; the next `run` must see that same head_hash so
// `ensure_head_or_fork` does not spawn a new session.
let client = MockClient::new(vec![
text_events_with_usage("hi", 1000),
write_extracted_tool_use_events("ec1"),
single_text_events("done"),
text_events_with_usage("ok", 1000),
]);
let mut pod = make_pod_with_manifest(EXTRACT_NO_COMPACT_MANIFEST, client).await;
pod.run_text("first").await.unwrap();
let session_before = pod.session_id();
pod.spawn_post_run_memory_jobs();
pod.wait_for_memory_jobs().await;
pod.run_text("second").await.unwrap();
let session_after = pod.session_id();
assert_eq!(
session_before, session_after,
"detached extract's save_extension and the next turn's save_user_input \
must share head_hash through SessionHead — a fork here means the clone \
carried its own head_hash"
);
}
crates/pod/tests/controller_test.rs
+7 -54
View File
@@ -173,7 +173,7 @@ async fn wait_for_status(handle: &PodHandle, status: PodStatus) {
// ---------------------------------------------------------------------------
#[tokio::test]
async fn run_end_enters_busy_until_post_run_finishes_and_broadcasts_status() {
async fn run_end_returns_to_idle_without_busy_status() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
@@ -182,7 +182,7 @@ async fn run_end_enters_busy_until_post_run_finishes_and_broadcasts_status() {
handle.send(Method::run_text("Hello")).await.unwrap();
let mut saw_run_end = false;
let mut saw_busy_status = false;
let mut saw_idle_status = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
tokio::select! {
@@ -191,10 +191,8 @@ async fn run_end_enters_busy_until_post_run_finishes_and_broadcasts_status() {
Ok(Event::RunEnd { result: protocol::RunResult::Finished }) => {
saw_run_end = true;
}
Ok(Event::Status {
status: PodStatus::Busy,
}) if saw_run_end => {
saw_busy_status = true;
Ok(Event::Status { status: PodStatus::Idle }) if saw_run_end => {
saw_idle_status = true;
break;
}
Ok(_) => {}
@@ -207,10 +205,10 @@ async fn run_end_enters_busy_until_post_run_finishes_and_broadcasts_status() {
assert!(saw_run_end, "expected RunEnd::Finished");
assert!(
saw_busy_status,
"expected busy status immediately after RunEnd"
saw_idle_status,
"expected idle status immediately after RunEnd"
);
wait_for_status(&handle, PodStatus::Idle).await;
assert_eq!(handle.shared_state.get_status(), PodStatus::Idle);
}
#[tokio::test]
@@ -237,51 +235,6 @@ async fn attach_history_includes_current_status() {
}
}
#[tokio::test]
async fn pause_while_busy_is_idempotent_not_not_running() {
let client = MockClient::new(simple_text_events());
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
let mut rx = handle.subscribe();
handle.send(Method::run_text("Hello")).await.unwrap();
let mut saw_busy = false;
let mut saw_idle = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
tokio::select! {
event = rx.recv() => {
match event {
Ok(Event::RunEnd { .. }) => {
handle.send(Method::Pause).await.unwrap();
}
Ok(Event::Status { status: PodStatus::Busy }) => {
saw_busy = true;
}
Ok(Event::Status { status: PodStatus::Idle }) if saw_busy => {
saw_idle = true;
break;
}
Ok(Event::Error {
code: protocol::ErrorCode::NotRunning,
..
}) if saw_busy && !saw_idle => {
panic!("Pause while Busy should be an idempotent no-op");
}
Ok(_) => {}
Err(_) => break,
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
assert!(saw_busy, "expected Busy status");
assert!(saw_idle, "expected final Idle status");
assert_eq!(handle.shared_state.get_status(), PodStatus::Idle);
}
#[tokio::test]
async fn shared_state_starts_idle() {
let client = MockClient::new(simple_text_events());