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

feat: Podのステータスを厳密にし、同期漏れを防ぐ

Browse Source
This commit is contained in:
Hare
2026-05-04 12:55:11 +09:00
parent a0771608b1
commit 560c23bc75
24 changed files with 641 additions and 311 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pod/src/compact/prune.rs
+22 -26
View File
@@ -49,33 +49,29 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
let metrics = self.metrics_tracker_handle();
let usage_tracker = self.usage_tracker_handle();
let observer: PruneObserver = Box::new(move |eval| {
match &eval.decision {
PruneDecision::Fired { .. } => {
let correlation_id = uuid::Uuid::now_v7().to_string();
let mut metric = Metric::now("prune.fire")
.with_value(eval.estimated_savings as f64)
.with_correlation_id(&correlation_id)
.with_dimension("candidate_count", eval.candidate_count.to_string());
if let Some(border) = eval.border_turn {
metric = metric.with_dimension("border_turn", border.to_string());
}
metrics.push(metric);
usage_tracker.note_correlation_id(correlation_id);
}
PruneDecision::SkippedNoCandidates => {
metrics.push(
Metric::now("prune.skip").with_dimension("reason", "no_candidates"),
);
}
PruneDecision::SkippedBelowMinSavings => {
metrics.push(
Metric::now("prune.skip")
.with_dimension("reason", "below_min_savings")
.with_dimension("candidate_count", eval.candidate_count.to_string())
.with_value(eval.estimated_savings as f64),
);
let observer: PruneObserver = Box::new(move |eval| match &eval.decision {
PruneDecision::Fired { .. } => {
let correlation_id = uuid::Uuid::now_v7().to_string();
let mut metric = Metric::now("prune.fire")
.with_value(eval.estimated_savings as f64)
.with_correlation_id(&correlation_id)
.with_dimension("candidate_count", eval.candidate_count.to_string());
if let Some(border) = eval.border_turn {
metric = metric.with_dimension("border_turn", border.to_string());
}
metrics.push(metric);
usage_tracker.note_correlation_id(correlation_id);
}
PruneDecision::SkippedNoCandidates => {
metrics.push(Metric::now("prune.skip").with_dimension("reason", "no_candidates"));
}
PruneDecision::SkippedBelowMinSavings => {
metrics.push(
Metric::now("prune.skip")
.with_dimension("reason", "below_min_savings")
.with_dimension("candidate_count", eval.candidate_count.to_string())
.with_value(eval.estimated_savings as f64),
);
}
});
crates/pod/src/controller.rs
+164 -155
View File
@@ -12,13 +12,15 @@ use crate::ipc::notify_buffer::NotifyBuffer;
use crate::ipc::server::SocketServer;
use crate::pod::{Pod, PodError, PodRunResult};
use crate::runtime::dir::RuntimeDir;
use crate::shared_state::{PodSharedState, PodStatus};
use crate::shared_state::PodSharedState;
use crate::spawn::comm_tools::{
list_pods_tool, read_pod_output_tool, send_to_pod_tool, stop_pod_tool,
};
use crate::spawn::registry::SpawnedPodRegistry;
use crate::spawn::tool::spawn_pod_tool;
use protocol::{AlertLevel, AlertSource, ErrorCode, Event, Method, RunResult, TurnResult};
use protocol::{
AlertLevel, AlertSource, ErrorCode, Event, Method, PodStatus, RunResult, TurnResult,
};
fn is_system_message_item(item: &Item) -> bool {
matches!(
@@ -63,6 +65,75 @@ impl PodHandle {
}
}
async fn set_controller_status(
shared_state: &Arc<PodSharedState>,
runtime_dir: &RuntimeDir,
event_tx: &broadcast::Sender<Event>,
status: PodStatus,
) {
shared_state.set_status(status);
let _ = runtime_dir.write_status(shared_state).await;
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,
{
if new_status == PodStatus::Busy {
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;
let _ = runtime_dir.write_history(shared_state).await;
}
// ---------------------------------------------------------------------------
// PodController — actor that owns a Pod
// ---------------------------------------------------------------------------
@@ -414,8 +485,13 @@ impl PodController {
let _ = event_tx.send(Event::UserMessage {
segments: input.clone(),
});
shared_state.set_status(PodStatus::Running);
let _ = runtime_dir.write_status(&shared_state).await;
set_controller_status(
&shared_state,
&runtime_dir,
&event_tx,
PodStatus::Running,
)
.await;
let run_future = async {
if was_paused {
@@ -430,6 +506,7 @@ impl PodController {
&event_tx,
&cancel_tx,
&shared_state,
&runtime_dir,
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
@@ -437,39 +514,15 @@ impl PodController {
)
.await;
if new_status == PodStatus::Idle {
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}"),
);
}
}
let items = pod.worker().history().to_vec();
shared_state.update_history(items);
shared_state.set_user_segments(pod.user_segments().to_vec());
shared_state.set_status(new_status);
let _ = runtime_dir.write_status(&shared_state).await;
let _ = runtime_dir.write_history(&shared_state).await;
finish_controller_run(
&mut pod,
&shared_state,
&runtime_dir,
&event_tx,
&alerter,
new_status,
)
.await;
if shutdown {
let _ = event_tx.send(Event::Shutdown);
@@ -482,17 +535,23 @@ impl PodController {
message: message.clone(),
});
pod.push_notify(message);
if shared_state.get_status() != PodStatus::Idle {
// RUNNING / Paused: the buffer push is the
// entire operation; the in-flight turn (or
// next Resume) will drain the buffer at its
// next pre_llm_request.
let status = shared_state.get_status();
if status != PodStatus::Idle {
// RUNNING / Paused / Busy: the buffer push is the
// entire operation; an in-flight turn (or the
// next Resume/Run after Busy) will drain the buffer
// at its next pre_llm_request.
continue;
}
// IDLE: auto-start a turn so the LLM sees the
// buffered notification(s) without a human Run.
shared_state.set_status(PodStatus::Running);
let _ = runtime_dir.write_status(&shared_state).await;
set_controller_status(
&shared_state,
&runtime_dir,
&event_tx,
PodStatus::Running,
)
.await;
let (new_status, shutdown) = run_with_cancel_support(
pod.run_for_notification(),
@@ -500,6 +559,7 @@ impl PodController {
&event_tx,
&cancel_tx,
&shared_state,
&runtime_dir,
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
@@ -507,39 +567,15 @@ impl PodController {
)
.await;
if new_status == PodStatus::Idle {
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}"),
);
}
}
let items = pod.worker().history().to_vec();
shared_state.update_history(items);
shared_state.set_user_segments(pod.user_segments().to_vec());
shared_state.set_status(new_status);
let _ = runtime_dir.write_status(&shared_state).await;
let _ = runtime_dir.write_history(&shared_state).await;
finish_controller_run(
&mut pod,
&shared_state,
&runtime_dir,
&event_tx,
&alerter,
new_status,
)
.await;
if shutdown {
let _ = event_tx.send(Event::Shutdown);
@@ -555,8 +591,13 @@ impl PodController {
});
continue;
}
shared_state.set_status(PodStatus::Running);
let _ = runtime_dir.write_status(&shared_state).await;
set_controller_status(
&shared_state,
&runtime_dir,
&event_tx,
PodStatus::Running,
)
.await;
let (new_status, shutdown) = run_with_cancel_support(
pod.resume(),
@@ -564,6 +605,7 @@ impl PodController {
&event_tx,
&cancel_tx,
&shared_state,
&runtime_dir,
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
@@ -571,39 +613,15 @@ impl PodController {
)
.await;
if new_status == PodStatus::Idle {
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}"),
);
}
}
let items = pod.worker().history().to_vec();
shared_state.update_history(items);
shared_state.set_user_segments(pod.user_segments().to_vec());
shared_state.set_status(new_status);
let _ = runtime_dir.write_status(&shared_state).await;
let _ = runtime_dir.write_history(&shared_state).await;
finish_controller_run(
&mut pod,
&shared_state,
&runtime_dir,
&event_tx,
&alerter,
new_status,
)
.await;
if shutdown {
let _ = event_tx.send(Event::Shutdown);
@@ -619,11 +637,12 @@ impl PodController {
}
Method::Pause => {
// 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.
if shared_state.get_status() != PodStatus::Paused {
// Already paused or post-run busy → 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) {
let _ = event_tx.send(Event::Error {
code: ErrorCode::NotRunning,
message: "Pod is not running".into(),
@@ -666,8 +685,13 @@ impl PodController {
// notification is not stranded. Matches the
// `Method::Notify` idle path.
if shared_state.get_status() == PodStatus::Idle {
shared_state.set_status(PodStatus::Running);
let _ = runtime_dir.write_status(&shared_state).await;
set_controller_status(
&shared_state,
&runtime_dir,
&event_tx,
PodStatus::Running,
)
.await;
let (new_status, shutdown) = run_with_cancel_support(
pod.run_for_notification(),
@@ -675,6 +699,7 @@ impl PodController {
&event_tx,
&cancel_tx,
&shared_state,
&runtime_dir,
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
@@ -682,39 +707,15 @@ impl PodController {
)
.await;
if new_status == PodStatus::Idle {
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}"),
);
}
}
let items = pod.worker().history().to_vec();
shared_state.update_history(items);
shared_state.set_user_segments(pod.user_segments().to_vec());
shared_state.set_status(new_status);
let _ = runtime_dir.write_status(&shared_state).await;
let _ = runtime_dir.write_history(&shared_state).await;
finish_controller_run(
&mut pod,
&shared_state,
&runtime_dir,
&event_tx,
&alerter,
new_status,
)
.await;
if shutdown {
let _ = event_tx.send(Event::Shutdown);
@@ -766,6 +767,7 @@ async fn run_with_cancel_support<F>(
event_tx: &broadcast::Sender<Event>,
cancel_tx: &mpsc::Sender<()>,
shared_state: &Arc<PodSharedState>,
runtime_dir: &RuntimeDir,
notify_buffer: &NotifyBuffer,
parent_socket: Option<&std::path::PathBuf>,
self_name: &str,
@@ -784,11 +786,18 @@ where
return match result {
Ok(r) => {
let (status, run_result) = match r {
PodRunResult::Finished => (PodStatus::Idle, RunResult::Finished),
PodRunResult::Finished => (PodStatus::Busy, RunResult::Finished),
PodRunResult::Paused => (PodStatus::Paused, RunResult::Paused),
PodRunResult::LimitReached => (PodStatus::Idle, RunResult::LimitReached),
PodRunResult::LimitReached => (PodStatus::Busy, RunResult::LimitReached),
};
let _ = event_tx.send(Event::RunEnd { result: run_result });
if status == PodStatus::Busy {
shared_state.set_status(PodStatus::Busy);
let _ = runtime_dir.write_status(shared_state).await;
let _ = event_tx.send(Event::Status {
status: PodStatus::Busy,
});
}
if matches!(run_result, RunResult::Finished) {
crate::ipc::event::fire_and_forget(
parent_socket.cloned(),
@@ -822,7 +831,7 @@ where
message,
},
);
(PodStatus::Idle, shutdown_requested)
(PodStatus::Busy, shutdown_requested)
}
};
}
crates/pod/src/ipc/interceptor.rs
+1 -1
View File
@@ -17,9 +17,9 @@ use llm_worker::interceptor::{
Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo,
ToolResultInfo, TurnEndAction,
};
use tracing::warn;
use llm_worker::tool::ToolOutput;
use tracing::info;
use tracing::warn;
use crate::compact::state::CompactState;
use crate::hook::{
crates/pod/src/ipc/server.rs
+2
View File
@@ -159,10 +159,12 @@ async fn handle_connection(stream: tokio::net::UnixStream, handle: PodHandle) {
})
.collect();
let greeting = handle.shared_state.greeting.clone();
let status = handle.shared_state.get_status();
if writer
.write(&Event::History {
items: values,
greeting,
status,
})
.await
.is_err()
crates/pod/src/lib.rs
+2 -2
View File
@@ -26,7 +26,7 @@ pub use pod::{Pod, PodError, PodRunResult, apply_worker_manifest};
pub use prompt::catalog::{CatalogError, PodPrompt, PromptCatalog};
pub use prompt::loader::PromptLoader;
pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate};
pub use protocol::{ErrorCode, Event, Method, TurnResult};
pub use protocol::{ErrorCode, Event, Method, PodStatus, TurnResult};
pub use provider::{ProviderError, build_client};
pub use runtime::dir::RuntimeDir;
pub use shared_state::{PodSharedState, PodStatus};
pub use shared_state::PodSharedState;
crates/pod/src/runtime/dir.rs
+2 -1
View File
@@ -131,7 +131,8 @@ pub fn default_base() -> Result<PathBuf, io::Error> {
#[cfg(test)]
mod tests {
use super::*;
use crate::shared_state::{PodSharedState, PodStatus};
use crate::shared_state::PodSharedState;
use protocol::PodStatus;
fn test_state() -> PodSharedState {
PodSharedState::new(
crates/pod/src/shared_state.rs
+3 -11
View File
@@ -1,8 +1,8 @@
use std::sync::{OnceLock, RwLock};
use llm_worker::llm_client::types::Item;
use protocol::Segment;
use serde::{Deserialize, Serialize};
use protocol::{PodStatus, Segment};
use serde_json::json;
use session_store::SessionId;
use crate::fs_view::PodFsView;
@@ -39,14 +39,6 @@ pub struct PodSharedState {
workflows: OnceLock<Vec<WorkflowCandidate>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PodStatus {
Idle,
Running,
Paused,
}
impl PodSharedState {
pub fn new(
pod_name: String,
@@ -138,7 +130,7 @@ impl PodSharedState {
/// Serialize status as JSON.
pub fn status_json(&self) -> String {
let status = self.get_status();
serde_json::json!({
json!({
"state": status,
"session_id": self.session_id.to_string(),
"pod_name": self.pod_name,
crates/pod/tests/consolidation_test.rs
+19 -19
View File
@@ -206,10 +206,7 @@ async fn no_thresholds_is_a_noop() {
.expect("phase 2 disabled when both thresholds are None");
// No staging entries removed.
assert_eq!(
memory::consolidate::list_staging_entries(&layout).len(),
5
);
assert_eq!(memory::consolidate::list_staging_entries(&layout).len(), 5);
}
#[tokio::test]
@@ -256,10 +253,7 @@ async fn below_threshold_skips_and_does_not_take_lock() {
pod.try_post_run_consolidate().await.unwrap();
// Staging untouched.
assert_eq!(
memory::consolidate::list_staging_entries(&layout).len(),
1
);
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");
@@ -285,10 +279,7 @@ async fn fires_on_threshold_and_cleans_up_consumed_entries() {
);
// Lock removed too.
let lock_path = layout.staging_dir().join(".consolidation.lock");
assert!(
!lock_path.exists(),
"lock file must be removed on success"
);
assert!(!lock_path.exists(), "lock file must be removed on success");
}
#[tokio::test]
@@ -300,7 +291,12 @@ async fn in_flight_guard_skips_reentry_without_clearing() {
write_n_staging(&layout, 2);
let client = MockClient::new(vec![]);
let mut pod = make_pod_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client.clone()).await;
let mut pod = make_pod_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
@@ -334,7 +330,9 @@ async fn in_flight_guard_skips_reentry_without_clearing() {
let mut pod2 = make_pod_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client2).await;
pod2.try_post_run_consolidate().await.unwrap();
assert!(
!pod2.consolidation_in_flight_handle().load(Ordering::Acquire),
!pod2
.consolidation_in_flight_handle()
.load(Ordering::Acquire),
"in-flight flag must be cleared after a normal run"
);
}
@@ -356,7 +354,12 @@ async fn coalesce_loop_terminates_with_one_iteration_when_snapshot_drains_stagin
// 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 pod = make_pod_with(FILES_THRESHOLD_TOML, pwd.path().to_path_buf(), client.clone()).await;
let mut pod = make_pod_with(
FILES_THRESHOLD_TOML,
pwd.path().to_path_buf(),
client.clone(),
)
.await;
pod.try_post_run_consolidate().await.unwrap();
assert_eq!(
@@ -393,8 +396,5 @@ async fn live_lock_held_by_other_pod_skips() {
.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
);
assert_eq!(memory::consolidate::list_staging_entries(&layout).len(), 3);
}
crates/pod/tests/controller_test.rs
+142 -14
View File
@@ -10,7 +10,7 @@ use llm_worker::llm_client::{ClientError, LlmClient, Request};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use session_store::FsStore;
use pod::{Event, Method, Pod, PodController, PodManifest, PodStatus};
use pod::{Event, Method, Pod, PodController, PodHandle, PodManifest, PodStatus};
// ---------------------------------------------------------------------------
// Mock LLM Client
@@ -147,8 +147,6 @@ async fn make_pod_with_pwd(client: MockClient) -> (Pod<MockClient, FsStore>, std
(pod, pwd)
}
use pod::PodHandle;
async fn spawn_controller(pod: Pod<MockClient, FsStore>) -> PodHandle {
let tmp = tempfile::tempdir().unwrap();
let runtime_base = tmp.path().to_owned();
@@ -157,9 +155,132 @@ async fn spawn_controller(pod: Pod<MockClient, FsStore>) -> PodHandle {
handle
}
async fn wait_for_status(handle: &PodHandle, status: PodStatus) {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
if handle.shared_state.get_status() == status {
return;
}
assert!(
tokio::time::Instant::now() < deadline,
"timed out waiting for status {status:?}; current={:?}",
handle.shared_state.get_status()
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[tokio::test]
async fn run_end_enters_busy_until_post_run_finishes_and_broadcasts_status() {
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_run_end = false;
let mut saw_busy_status = false;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
tokio::select! {
event = rx.recv() => {
match event {
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;
break;
}
Ok(_) => {}
Err(_) => break,
}
}
_ = tokio::time::sleep_until(deadline) => break,
}
}
assert!(saw_run_end, "expected RunEnd::Finished");
assert!(
saw_busy_status,
"expected busy status immediately after RunEnd"
);
wait_for_status(&handle, PodStatus::Idle).await;
}
#[tokio::test]
async fn attach_history_includes_current_status() {
let client = MockClient::sequential(vec![MockResponse::Hang(simple_text_events())]);
let pod = make_pod(client).await;
let handle = spawn_controller(pod).await;
handle.send(Method::run_text("Hello")).await.unwrap();
wait_for_status(&handle, PodStatus::Running).await;
let stream = tokio::net::UnixStream::connect(handle.runtime_dir.socket_path())
.await
.unwrap();
let (reader, writer) = stream.into_split();
let mut reader = protocol::stream::JsonLineReader::new(reader);
let mut writer = protocol::stream::JsonLineWriter::new(writer);
writer.write(&Method::GetHistory).await.unwrap();
let event = reader.next::<Event>().await.unwrap().unwrap();
match event {
Event::History { status, .. } => assert_eq!(status, PodStatus::Running),
other => panic!("expected History, got {other:?}"),
}
}
#[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() {
@@ -565,10 +686,10 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
// request context for that turn).
let requests = client_for_assert.captured_requests();
assert_eq!(requests.len(), 1, "one LLM call expected");
let notify_in_request = requests[0]
.items
.iter()
.any(|i| i.as_text().is_some_and(|t| t.contains("[Notification]") && t.contains("turn finished")));
let notify_in_request = requests[0].items.iter().any(|i| {
i.as_text()
.is_some_and(|t| t.contains("[Notification]") && t.contains("turn finished"))
});
assert!(
notify_in_request,
"injected system message missing from request, got items: {:?}",
@@ -583,13 +704,17 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
// (and therefore eventually into history.json), per
// tickets/notify-history-persist.md.
let history = handle.shared_state.history();
let notify_in_history = history
.iter()
.any(|i| i.as_text().is_some_and(|t| t.contains("[Notification]") && t.contains("turn finished")));
let notify_in_history = history.iter().any(|i| {
i.as_text()
.is_some_and(|t| t.contains("[Notification]") && t.contains("turn finished"))
});
assert!(
notify_in_history,
"notify must be committed to worker.history, got items: {:?}",
history.iter().filter_map(|i| i.as_text()).collect::<Vec<_>>()
history
.iter()
.filter_map(|i| i.as_text())
.collect::<Vec<_>>()
);
}
@@ -671,7 +796,10 @@ async fn pod_event_turn_ended_while_idle_auto_starts_turn_and_injects_system_mes
assert!(
event_in_history,
"PodEvent must be committed to worker.history, got items: {:?}",
history.iter().filter_map(|i| i.as_text()).collect::<Vec<_>>()
history
.iter()
.filter_map(|i| i.as_text())
.collect::<Vec<_>>()
);
}
crates/pod/tests/pod_comm_tools_test.rs
+1
View File
@@ -174,6 +174,7 @@ fn serve_history(listener: UnixListener, items: Vec<Item>) -> JoinHandle<()> {
scope_summary: String::new(),
tools: Vec::new(),
},
status: protocol::PodStatus::Idle,
};
let _ = writer.write(&event).await;
}
crates/pod/tests/session_metrics_test.rs
+15 -12
View File
@@ -22,9 +22,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use futures::Stream;
use llm_worker::Worker;
use llm_worker::llm_client::event::{
Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent,
};
use llm_worker::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent, UsageEvent};
use llm_worker::llm_client::{ClientError, LlmClient, Request};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use session_metrics::{DOMAIN, Metric, metrics_from_extensions};
@@ -169,7 +167,11 @@ async fn make_pod(
manifest_toml: String,
client: MockClient,
tool_name: &'static str,
) -> (Pod<MockClient, FsStore>, tempfile::TempDir, tempfile::TempDir) {
) -> (
Pod<MockClient, FsStore>,
tempfile::TempDir,
tempfile::TempDir,
) {
let manifest = PodManifest::from_toml(&manifest_toml).unwrap();
let store_tmp = tempfile::tempdir().unwrap();
let store = FsStore::new(store_tmp.path()).await.unwrap();
@@ -199,8 +201,7 @@ async fn prune_metrics_emit_skip_then_fire_with_post_request_join() {
text_response_with_cache("ok", 0, 200),
text_response_with_cache("done", 1234, 50),
]);
let (mut pod, _store_tmp, _pwd_tmp) =
make_pod(manifest_toml(1, 1), client, "big_tool").await;
let (mut pod, _store_tmp, _pwd_tmp) = make_pod(manifest_toml(1, 1), client, "big_tool").await;
let session_id = pod.session_id();
// Cloning the store handle to read the session log back after the
// runs complete — the Pod retains its own copy.
@@ -253,10 +254,7 @@ async fn prune_metrics_emit_skip_then_fire_with_post_request_join() {
fire.dimensions.contains_key("border_turn"),
"fire missing border_turn: {fire:?}"
);
assert!(
fire.value.is_some(),
"fire missing estimated_savings value"
);
assert!(fire.value.is_some(), "fire missing estimated_savings value");
let fire_id = fire
.correlation_id
.as_ref()
@@ -272,7 +270,9 @@ async fn prune_metrics_emit_skip_then_fire_with_post_request_join() {
assert_eq!(post.correlation_id.as_ref(), Some(fire_id));
assert_eq!(post.value, Some(1234.0));
assert_eq!(
post.dimensions.get("cache_write_tokens").map(String::as_str),
post.dimensions
.get("cache_write_tokens")
.map(String::as_str),
Some("50")
);
assert!(post.dimensions.contains_key("history_len"));
@@ -457,7 +457,10 @@ permission = "write"
let state = session_store::restore(&store, session_id).await.unwrap();
let metrics = metrics_from_extensions(&state.extensions);
assert!(metrics.is_empty(), "no metrics should be recorded: {metrics:?}");
assert!(
metrics.is_empty(),
"no metrics should be recorded: {metrics:?}"
);
// And no extension entries at all in the metrics domain.
assert!(state.extensions.iter().all(|(d, _)| d != DOMAIN));