refactor: Podのメインループのリファクタリング
This commit is contained in:
@@ -177,17 +177,39 @@ fn drain(rx: &mut broadcast::Receiver<Event>) -> Vec<Event> {
|
||||
out
|
||||
}
|
||||
|
||||
fn system_event_text(event: &Event) -> Option<&str> {
|
||||
match event {
|
||||
Event::SystemMessage { item } => item["content"]
|
||||
.as_array()
|
||||
.and_then(|parts| parts.iter().filter_map(|p| p["text"].as_str()).next()),
|
||||
_ => None,
|
||||
/// Collect every system-message text that the post-compaction
|
||||
/// `SessionStart.history` carries, by reading the sink mirror directly.
|
||||
fn system_texts_in_sink_session_start(pod: &pod::Pod<impl llm_worker::llm_client::client::LlmClient + Clone + 'static, impl session_store::Store + Clone + 'static>) -> Vec<String> {
|
||||
let (entries, _rx) = pod.sink().subscribe_with_snapshot();
|
||||
for entry in entries.into_iter().rev() {
|
||||
if let session_store::LogEntry::SessionStart { history, .. } = entry {
|
||||
return history
|
||||
.into_iter()
|
||||
.filter_map(|logged| {
|
||||
let item: Item = logged.into();
|
||||
match item {
|
||||
Item::Message {
|
||||
role: llm_worker::Role::System,
|
||||
content,
|
||||
..
|
||||
} => Some(
|
||||
content
|
||||
.iter()
|
||||
.map(|p| p.as_text().to_owned())
|
||||
.collect::<Vec<_>>()
|
||||
.join(""),
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
}
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compact_broadcasts_only_new_system_messages_not_retained_ones() {
|
||||
async fn compact_emits_session_start_carrying_summary_and_task_snapshot() {
|
||||
let client = MockClient::new(vec![
|
||||
single_text_events("hi"),
|
||||
write_summary_tool_use_events("call-1", "summary"),
|
||||
@@ -195,18 +217,16 @@ async fn compact_broadcasts_only_new_system_messages_not_retained_ones() {
|
||||
]);
|
||||
let mut pod = make_pod(client).await;
|
||||
|
||||
let (tx, mut rx) = broadcast::channel::<Event>(64);
|
||||
let (tx, _rx_keep) = broadcast::channel::<Event>(64);
|
||||
pod.attach_event_tx(tx);
|
||||
|
||||
pod.run_text("first").await.unwrap();
|
||||
let retained_message = Item::system_message("[Retained system]\nold");
|
||||
pod.worker_mut().push_item(retained_message);
|
||||
let _ = drain(&mut rx);
|
||||
|
||||
pod.compact(10_000).await.unwrap();
|
||||
|
||||
let events = drain(&mut rx);
|
||||
let system_texts: Vec<&str> = events.iter().filter_map(system_event_text).collect();
|
||||
let system_texts = system_texts_in_sink_session_start(&pod);
|
||||
// The post-compaction `SessionStart.history` carries the new system
|
||||
// messages introduced by the compactor. Clients re-seed their view
|
||||
// from this entry alone, so it is the load-bearing payload.
|
||||
assert!(
|
||||
system_texts
|
||||
.iter()
|
||||
@@ -219,12 +239,6 @@ async fn compact_broadcasts_only_new_system_messages_not_retained_ones() {
|
||||
.any(|text| text.starts_with("[Session TaskStore snapshot]")),
|
||||
"task snapshot system message missing from {system_texts:?}"
|
||||
);
|
||||
assert!(
|
||||
!system_texts
|
||||
.iter()
|
||||
.any(|text| text.starts_with("[Retained system]")),
|
||||
"retained system message should not be rebroadcast: {system_texts:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -6,12 +6,40 @@ use async_trait::async_trait;
|
||||
use futures::{Stream, StreamExt};
|
||||
use llm_worker::Worker;
|
||||
use llm_worker::llm_client::event::{Event as LlmEvent, ResponseStatus, StatusEvent};
|
||||
use llm_worker::llm_client::types::Item;
|
||||
use llm_worker::llm_client::{ClientError, LlmClient, Request};
|
||||
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
||||
use session_store::FsStore;
|
||||
use session_store::{FsStore, LogEntry};
|
||||
|
||||
use pod::{Event, Method, Pod, PodController, PodHandle, PodManifest, PodStatus};
|
||||
|
||||
/// Reconstruct a worker-history-like `Vec<Item>` from the live session
|
||||
/// log mirror held by the Pod's broadcast sink. Replaces the previous
|
||||
/// `PodSharedState.history()` test helper now that the mirror lives in
|
||||
/// the sink.
|
||||
fn history_from_sink(handle: &PodHandle) -> Vec<Item> {
|
||||
let (entries, _rx) = handle.sink.subscribe_with_snapshot();
|
||||
let mut items = Vec::new();
|
||||
for entry in entries {
|
||||
match entry {
|
||||
LogEntry::SessionStart { history, .. } => {
|
||||
items.extend(history.into_iter().map(Item::from));
|
||||
}
|
||||
LogEntry::UserInput { segments, .. } => {
|
||||
let text = protocol::Segment::flatten_to_text(&segments);
|
||||
items.push(Item::user_message(text));
|
||||
}
|
||||
LogEntry::AssistantItems { items: i, .. }
|
||||
| LogEntry::ToolResults { items: i, .. }
|
||||
| LogEntry::HookInjectedItems { items: i, .. } => {
|
||||
items.extend(i.into_iter().map(Item::from));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
items
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Mock LLM Client
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -218,8 +246,61 @@ async fn run_end_returns_to_idle_without_busy_status() {
|
||||
assert_eq!(handle.shared_state.get_status(), PodStatus::Idle);
|
||||
}
|
||||
|
||||
/// Mid-turn re-attach: a client connecting while the worker is still
|
||||
/// running observes the in-flight `UserInput` entry in the connect-time
|
||||
/// `Event::Snapshot`. This is the load-bearing property of the new
|
||||
/// session-log-driven IPC: a late attacher reconstructs the running
|
||||
/// view without needing the prior client's diff.
|
||||
#[tokio::test]
|
||||
async fn attach_history_includes_current_status() {
|
||||
async fn snapshot_includes_user_input_for_in_flight_turn() {
|
||||
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 in-flight"))
|
||||
.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);
|
||||
|
||||
loop {
|
||||
let event = reader.next::<Event>().await.unwrap().unwrap();
|
||||
match event {
|
||||
Event::Snapshot { entries, .. } => {
|
||||
// Walk the entries, find a `LogEntry::UserInput` and
|
||||
// confirm its segments flatten to our submitted text.
|
||||
let mut found = false;
|
||||
for value in entries {
|
||||
let entry: session_store::LogEntry =
|
||||
serde_json::from_value(value).expect("LogEntry deserialise");
|
||||
if let session_store::LogEntry::UserInput { segments, .. } = entry {
|
||||
let text = protocol::Segment::flatten_to_text(&segments);
|
||||
if text == "hello in-flight" {
|
||||
found = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
found,
|
||||
"snapshot must carry the in-flight UserInput entry"
|
||||
);
|
||||
return;
|
||||
}
|
||||
Event::Alert(_) => continue,
|
||||
other => panic!("expected Snapshot first, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn attach_snapshot_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;
|
||||
@@ -230,15 +311,20 @@ async fn attach_history_includes_current_status() {
|
||||
let stream = tokio::net::UnixStream::connect(handle.runtime_dir.socket_path())
|
||||
.await
|
||||
.unwrap();
|
||||
let (reader, writer) = stream.into_split();
|
||||
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:?}"),
|
||||
// First event after connect is the snapshot — it carries the current status.
|
||||
loop {
|
||||
let event = reader.next::<Event>().await.unwrap().unwrap();
|
||||
match event {
|
||||
Event::Snapshot { status, .. } => {
|
||||
assert_eq!(status, PodStatus::Running);
|
||||
return;
|
||||
}
|
||||
Event::Alert(_) => continue,
|
||||
other => panic!("expected Snapshot, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -275,11 +361,11 @@ async fn run_populates_history() {
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
|
||||
let history = handle.shared_state.history_json();
|
||||
assert_ne!(history, "[]");
|
||||
let parsed: serde_json::Value = serde_json::from_str(&history).unwrap();
|
||||
assert!(parsed.is_array());
|
||||
assert!(parsed.as_array().unwrap().len() >= 2); // user + assistant
|
||||
let history = history_from_sink(&handle);
|
||||
assert!(
|
||||
history.len() >= 2,
|
||||
"history must include user + assistant items, got {history:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -712,7 +798,7 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
|
||||
// The notification must also be persisted into the Worker history
|
||||
// (and therefore eventually into history.json), per
|
||||
// tickets/notify-history-persist.md.
|
||||
let history = handle.shared_state.history();
|
||||
let history = history_from_sink(&handle);
|
||||
let notify_in_history = history.iter().any(|i| {
|
||||
i.as_text()
|
||||
.is_some_and(|t| t.contains("[Notification]") && t.contains("turn finished"))
|
||||
@@ -796,7 +882,7 @@ async fn pod_event_turn_ended_while_idle_auto_starts_turn_and_injects_system_mes
|
||||
|
||||
// Same item must be present in worker.history (persisted lane),
|
||||
// not just the per-request clone — see tickets/notify-history-persist.md.
|
||||
let history = handle.shared_state.history();
|
||||
let history = history_from_sink(&handle);
|
||||
let event_in_history = history.iter().any(|i| {
|
||||
i.as_text().is_some_and(|t| {
|
||||
t.contains("[Notification]") && t.contains("child") && t.contains("finished a turn")
|
||||
@@ -1149,22 +1235,42 @@ async fn pause_then_resume_transitions_and_preserves_history_consistency() {
|
||||
// History consistency: exactly [user "hello", assistant
|
||||
// "resumed output"]. No artifacts from the aborted stream
|
||||
// (partial text is not committed), no orphan tool_use.
|
||||
let history_json = handle.shared_state.history_json();
|
||||
let items: Vec<serde_json::Value> = serde_json::from_str(&history_json).unwrap();
|
||||
let roles: Vec<&str> = items.iter().filter_map(|i| i["role"].as_str()).collect();
|
||||
let history = history_from_sink(&handle);
|
||||
let roles: Vec<&str> = history
|
||||
.iter()
|
||||
.filter_map(|i| match i {
|
||||
Item::Message { role, .. } => match role {
|
||||
llm_worker::Role::User => Some("user"),
|
||||
llm_worker::Role::Assistant => Some("assistant"),
|
||||
llm_worker::Role::System => Some("system"),
|
||||
},
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(
|
||||
roles,
|
||||
vec!["user", "assistant"],
|
||||
"history = user + assistant only; got {items:?}"
|
||||
"history = user + assistant only; got {history:?}"
|
||||
);
|
||||
let assistant_text = items[1]["content"]
|
||||
.as_array()
|
||||
.and_then(|parts| parts.iter().filter_map(|p| p["text"].as_str()).next())
|
||||
.unwrap_or("");
|
||||
assert_eq!(assistant_text, "resumed output");
|
||||
let has_tool_call = items
|
||||
let assistant_text = history
|
||||
.iter()
|
||||
.any(|i| i["type"].as_str() == Some("tool_call"));
|
||||
.find_map(|i| match i {
|
||||
Item::Message {
|
||||
role: llm_worker::Role::Assistant,
|
||||
content,
|
||||
..
|
||||
} => Some(
|
||||
content
|
||||
.iter()
|
||||
.map(|p: &llm_worker::ContentPart| p.as_text().to_owned())
|
||||
.collect::<Vec<_>>()
|
||||
.join(""),
|
||||
),
|
||||
_ => None,
|
||||
})
|
||||
.unwrap_or_default();
|
||||
assert_eq!(assistant_text, "resumed output");
|
||||
let has_tool_call = history.iter().any(|i| i.is_tool_call());
|
||||
assert!(!has_tool_call, "no orphan tool_call in history");
|
||||
}
|
||||
|
||||
|
||||
@@ -144,44 +144,41 @@ fn accept_method_and_respond(
|
||||
})
|
||||
}
|
||||
|
||||
/// Pretend to be a spawned Pod that responds to `GetHistory` with a
|
||||
/// fixed set of items. Accepts connections until the first one that
|
||||
/// delivers a `GetHistory` method; earlier probes (empty accepts) and
|
||||
/// non-history methods are ignored. Returns nothing — tests await the
|
||||
/// handle only to keep the listener alive until shutdown.
|
||||
/// Pretend to be a spawned Pod whose connect-time snapshot carries a
|
||||
/// fixed set of assistant items. Sends `Event::Snapshot` immediately on
|
||||
/// every accept — the real Pod does the same, so `ReadPodOutput`'s
|
||||
/// `fetch_history` just consumes the first non-Alert event.
|
||||
fn serve_history(listener: UnixListener, items: Vec<Item>) -> JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let Ok((stream, _)) = listener.accept().await else {
|
||||
return;
|
||||
};
|
||||
let (r, w) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(r);
|
||||
let (_r, w) = stream.into_split();
|
||||
let mut writer = JsonLineWriter::new(w);
|
||||
match reader.next::<Method>().await {
|
||||
Ok(Some(Method::GetHistory)) => {
|
||||
let values: Vec<serde_json::Value> = items
|
||||
.iter()
|
||||
.map(|i| serde_json::to_value(i).unwrap())
|
||||
.collect();
|
||||
let event = Event::History {
|
||||
items: values,
|
||||
greeting: Greeting {
|
||||
pod_name: "child".into(),
|
||||
cwd: "/tmp".into(),
|
||||
provider: "anthropic".into(),
|
||||
model: "x".into(),
|
||||
scope_summary: String::new(),
|
||||
tools: Vec::new(),
|
||||
},
|
||||
status: protocol::PodStatus::Idle,
|
||||
};
|
||||
let _ = writer.write(&event).await;
|
||||
}
|
||||
Ok(Some(_)) | Ok(None) | Err(_) => {
|
||||
// Ignore: loop accepts another connection.
|
||||
}
|
||||
}
|
||||
// Wrap the assistant items in a single
|
||||
// `LogEntry::AssistantItems` entry — that's the only kind
|
||||
// that contributes assistant text via `extract_assistant_text`.
|
||||
let logged: Vec<session_store::LoggedItem> =
|
||||
items.iter().map(session_store::LoggedItem::from).collect();
|
||||
let entry = session_store::LogEntry::AssistantItems {
|
||||
ts: 0,
|
||||
items: logged,
|
||||
};
|
||||
let entry_value = serde_json::to_value(&entry).unwrap();
|
||||
let event = Event::Snapshot {
|
||||
entries: vec![entry_value],
|
||||
greeting: Greeting {
|
||||
pod_name: "child".into(),
|
||||
cwd: "/tmp".into(),
|
||||
provider: "anthropic".into(),
|
||||
model: "x".into(),
|
||||
scope_summary: String::new(),
|
||||
tools: Vec::new(),
|
||||
},
|
||||
status: protocol::PodStatus::Idle,
|
||||
};
|
||||
let _ = writer.write(&event).await;
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user