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/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));