podのモジュール分割
This commit is contained in:
@@ -0,0 +1,187 @@
|
||||
//! `PodEvent` send / receive helpers.
|
||||
//!
|
||||
//! This module owns the parent-facing lifecycle-event primitive
|
||||
//! (`PodEvent`) that children fire upward on turn-end / error /
|
||||
//! shutdown / scope-sub-delegation. Three responsibilities live here:
|
||||
//!
|
||||
//! - **Send** a `Method::PodEvent` to the parent socket, fire-and-forget,
|
||||
//! logging failures without blocking the child.
|
||||
//! - **Render** a variant into a human-readable string that the parent's
|
||||
//! LLM sees via the notification buffer.
|
||||
//! - **Apply side effects** on the parent (registry / scope-lock
|
||||
//! updates) so that the receive path is idempotent and tolerant of
|
||||
//! out-of-order delivery.
|
||||
//!
|
||||
//! Transport is fire-and-forget — the ticket's decision is that
|
||||
//! callbacks are an optimisation and `ListPods` + `reclaim_stale` are
|
||||
//! the real fallback. This module is allowed to drop events on the
|
||||
//! floor (with a warn log) rather than retry.
|
||||
//!
|
||||
//! `apply_event_side_effects` takes its dependencies (registry, scope
|
||||
//! lock path, self identity) by reference so the caller owns lifetime
|
||||
//! and locking concerns.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use protocol::{Method, PodEvent, ScopeRule};
|
||||
|
||||
use crate::spawn::comm_tools::connect_and_send;
|
||||
use crate::runtime::dir::SpawnedPodRecord;
|
||||
use crate::runtime::scope_lock::{self, ScopeLockError};
|
||||
use crate::spawn::registry::SpawnedPodRegistry;
|
||||
|
||||
/// Connect to `socket`, send a single `Method::PodEvent(event)`, and
|
||||
/// return. Used by children to report up to their parent.
|
||||
///
|
||||
/// This is a synchronous helper — callers that want fire-and-forget
|
||||
/// semantics should wrap the call in `tokio::spawn` themselves.
|
||||
pub async fn send_pod_event(socket: &Path, event: PodEvent) -> std::io::Result<()> {
|
||||
connect_and_send(socket, &Method::PodEvent(event)).await
|
||||
}
|
||||
|
||||
/// Spawn a fire-and-forget task that sends `event` to `socket`. If
|
||||
/// `socket` is `None`, no send happens (top-level Pods have no parent).
|
||||
/// Any send failure is logged at warn level but otherwise ignored —
|
||||
/// the parent is treated as best-effort.
|
||||
pub fn fire_and_forget(socket: Option<PathBuf>, event: PodEvent) {
|
||||
let Some(socket) = socket else { return };
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = send_pod_event(&socket, event).await {
|
||||
tracing::warn!(error = %e, socket = %socket.display(), "PodEvent send failed");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Render a variant into the one-line human-readable string that will
|
||||
/// be injected into the parent's LLM context as a system message.
|
||||
///
|
||||
/// Kept deliberately short — the LLM can always call `ReadPodOutput`
|
||||
/// to fetch more detail if the event summary is not enough.
|
||||
pub fn render_event(event: &PodEvent) -> String {
|
||||
match event {
|
||||
PodEvent::TurnEnded { pod_name } => {
|
||||
format!("Pod `{pod_name}` finished a turn.")
|
||||
}
|
||||
PodEvent::Errored { pod_name, message } => {
|
||||
format!("Pod `{pod_name}` reported an error: {message}")
|
||||
}
|
||||
PodEvent::ShutDown { pod_name } => {
|
||||
format!("Pod `{pod_name}` has stopped.")
|
||||
}
|
||||
PodEvent::ScopeSubDelegated {
|
||||
parent_pod,
|
||||
sub_pod,
|
||||
..
|
||||
} => {
|
||||
format!("Pod `{parent_pod}` spawned `{sub_pod}` and delegated scope to it.")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply the variant-specific side effect on the parent side.
|
||||
///
|
||||
/// All operations are idempotent so that out-of-order delivery (e.g.
|
||||
/// `TurnEnded` arriving after `ShutDown`) does not produce errors:
|
||||
///
|
||||
/// - `TurnEnded` / `Errored`: no system work; the LLM handles the
|
||||
/// semantic response.
|
||||
/// - `ShutDown`: remove the child from `spawned_pods.json` and release
|
||||
/// its scope allocation. Missing entries are swallowed.
|
||||
/// - `ScopeSubDelegated`: register the grandchild locally and re-emit
|
||||
/// upward to our own parent if we have one. Duplicate grandchild
|
||||
/// entries (re-delivery) are swallowed.
|
||||
pub async fn apply_event_side_effects(
|
||||
event: &PodEvent,
|
||||
registry: &Arc<SpawnedPodRegistry>,
|
||||
self_name: &str,
|
||||
self_parent_socket: &Option<PathBuf>,
|
||||
) {
|
||||
match event {
|
||||
PodEvent::TurnEnded { .. } | PodEvent::Errored { .. } => {}
|
||||
|
||||
PodEvent::ShutDown { pod_name } => {
|
||||
if let Err(e) = registry.remove(pod_name).await {
|
||||
tracing::warn!(error = %e, pod = %pod_name, "registry remove on ShutDown failed");
|
||||
}
|
||||
release_scope_silently(pod_name);
|
||||
}
|
||||
|
||||
PodEvent::ScopeSubDelegated {
|
||||
parent_pod,
|
||||
sub_pod,
|
||||
sub_socket,
|
||||
scope,
|
||||
} => {
|
||||
if registry.get(sub_pod).await.is_some() {
|
||||
return;
|
||||
}
|
||||
let callback_address = registry
|
||||
.get(parent_pod)
|
||||
.await
|
||||
.map(|r| r.socket_path)
|
||||
.unwrap_or_else(PathBuf::new);
|
||||
let record = SpawnedPodRecord {
|
||||
pod_name: sub_pod.clone(),
|
||||
socket_path: sub_socket.clone(),
|
||||
scope_delegated: scope.clone(),
|
||||
callback_address,
|
||||
};
|
||||
if let Err(e) = registry.add(record).await {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
sub_pod = %sub_pod,
|
||||
"registry add on ScopeSubDelegated failed"
|
||||
);
|
||||
}
|
||||
reemit_scope_sub_delegated(
|
||||
self_parent_socket,
|
||||
self_name,
|
||||
sub_pod.clone(),
|
||||
sub_socket.clone(),
|
||||
scope.clone(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn release_scope_silently(pod_name: &str) {
|
||||
let lock_path = match scope_lock::default_lock_path() {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "default_lock_path failed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let mut guard = match scope_lock::LockFileGuard::open(&lock_path) {
|
||||
Ok(g) => g,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "LockFileGuard open failed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
match scope_lock::release_pod(&mut guard, pod_name) {
|
||||
Ok(()) => {}
|
||||
Err(ScopeLockError::UnknownPod(_)) => {}
|
||||
Err(e) => tracing::warn!(error = ?e, pod = %pod_name, "release_pod failed"),
|
||||
}
|
||||
}
|
||||
|
||||
fn reemit_scope_sub_delegated(
|
||||
self_parent_socket: &Option<PathBuf>,
|
||||
self_name: &str,
|
||||
sub_pod: String,
|
||||
sub_socket: PathBuf,
|
||||
scope: Vec<ScopeRule>,
|
||||
) {
|
||||
let Some(parent_socket) = self_parent_socket.clone() else {
|
||||
return;
|
||||
};
|
||||
let event = PodEvent::ScopeSubDelegated {
|
||||
parent_pod: self_name.to_string(),
|
||||
sub_pod,
|
||||
sub_socket,
|
||||
scope,
|
||||
};
|
||||
fire_and_forget(Some(parent_socket), event);
|
||||
}
|
||||
@@ -0,0 +1,468 @@
|
||||
//! Pod-owned `Interceptor` implementation.
|
||||
//!
|
||||
//! Bridges Pod's internal mechanisms (compaction trigger today;
|
||||
//! notification injection / output truncation in the future) and the
|
||||
//! public `HookRegistry`. Internal mechanisms run first and have full
|
||||
//! mutable access via the `Interceptor` trait. Hooks then receive
|
||||
//! read-only summary information and only return control-flow
|
||||
//! decisions (continue / skip / abort / pause).
|
||||
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use llm_worker::Item;
|
||||
use llm_worker::interceptor::{
|
||||
Interceptor, PostToolAction, PreRequestAction, PreToolAction, PromptAction, ToolCallInfo,
|
||||
ToolResultInfo, TurnEndAction,
|
||||
};
|
||||
use llm_worker::tool::ToolOutput;
|
||||
use session_store::UsageRecord;
|
||||
use tracing::info;
|
||||
|
||||
use crate::compact::state::CompactState;
|
||||
use crate::hook::{
|
||||
AbortInfo, HookRegistry, PreRequestInfo, PromptSubmitInfo, ToolCallSummary, ToolResultSummary,
|
||||
TurnEndInfo,
|
||||
};
|
||||
use crate::ipc::notification_buffer::{NotificationBuffer, format_notification};
|
||||
use crate::prompt::catalog::PromptCatalog;
|
||||
use crate::compact::token_counter::total_tokens_impl;
|
||||
use tracing::warn;
|
||||
|
||||
/// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`.
|
||||
const FINAL_TEXT_PREVIEW_LIMIT: usize = 512;
|
||||
|
||||
pub(crate) struct PodInterceptor {
|
||||
registry: Arc<HookRegistry>,
|
||||
compact_state: Option<Arc<CompactState>>,
|
||||
/// Shared view of the cumulative UsageRecord timeline. Used with the
|
||||
/// per-request `context` to estimate current occupancy for threshold
|
||||
/// checks. `None` when compaction is disabled (both thresholds unset).
|
||||
usage_history: Option<Arc<Mutex<Vec<UsageRecord>>>>,
|
||||
/// Pending-notification buffer drained into the per-request
|
||||
/// context at the head of `pre_llm_request`.
|
||||
pending_notifications: NotificationBuffer,
|
||||
/// Prompt catalog used to render the injected notification wrapper.
|
||||
prompts: Arc<PromptCatalog>,
|
||||
/// Next turn index assigned by `on_prompt_submit`.
|
||||
next_turn_index: AtomicUsize,
|
||||
/// Tool calls observed in the current turn (reset on each new prompt).
|
||||
tool_calls_this_turn: AtomicUsize,
|
||||
}
|
||||
|
||||
impl PodInterceptor {
|
||||
pub(crate) fn new(
|
||||
registry: Arc<HookRegistry>,
|
||||
compact_state: Option<Arc<CompactState>>,
|
||||
usage_history: Option<Arc<Mutex<Vec<UsageRecord>>>>,
|
||||
pending_notifications: NotificationBuffer,
|
||||
prompts: Arc<PromptCatalog>,
|
||||
) -> Self {
|
||||
Self {
|
||||
registry,
|
||||
compact_state,
|
||||
usage_history,
|
||||
pending_notifications,
|
||||
prompts,
|
||||
next_turn_index: AtomicUsize::new(0),
|
||||
tool_calls_this_turn: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn current_turn_index(&self) -> usize {
|
||||
self.next_turn_index
|
||||
.load(Ordering::Relaxed)
|
||||
.saturating_sub(1)
|
||||
}
|
||||
|
||||
/// Estimate current input-token occupancy for `context`, projected
|
||||
/// through the shared UsageRecord timeline. Returns `None` when
|
||||
/// `usage_history` is not attached (compaction fully disabled).
|
||||
fn estimated_tokens(&self, context: &[Item]) -> Option<u64> {
|
||||
let handle = self.usage_history.as_ref()?;
|
||||
let records = handle.lock().expect("usage_history poisoned").clone();
|
||||
Some(total_tokens_impl(context, &records).tokens)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Interceptor for PodInterceptor {
|
||||
async fn on_prompt_submit(&self, item: &mut Item) -> PromptAction {
|
||||
let turn_index = self.next_turn_index.fetch_add(1, Ordering::Relaxed);
|
||||
self.tool_calls_this_turn.store(0, Ordering::Relaxed);
|
||||
|
||||
let info = PromptSubmitInfo {
|
||||
input_text: extract_message_text(item).unwrap_or_default(),
|
||||
turn_index,
|
||||
};
|
||||
for hook in &self.registry.on_prompt_submit {
|
||||
let action = hook.call(&info).await;
|
||||
if !matches!(action, PromptAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PromptAction::Continue
|
||||
}
|
||||
|
||||
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
|
||||
let current_tokens = self.estimated_tokens(context);
|
||||
|
||||
// Internal mechanism: between-requests compaction trigger (safety net).
|
||||
if let Some(state) = self.compact_state.as_ref() {
|
||||
if !state.is_disabled() {
|
||||
let current = current_tokens.unwrap_or(0);
|
||||
if state.exceeds_request(current) {
|
||||
info!(
|
||||
input_tokens = current,
|
||||
threshold = state.request_threshold().unwrap_or(0),
|
||||
"Between-requests compaction threshold exceeded, yielding"
|
||||
);
|
||||
return PreRequestAction::Yield;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Internal mechanism: drain pending `Method::Notify` notifications
|
||||
// into the per-request context as transient system messages.
|
||||
// These are not persisted to the Worker history; they exist only
|
||||
// for this single LLM request.
|
||||
for notification in self.pending_notifications.drain() {
|
||||
match format_notification(¬ification, &self.prompts) {
|
||||
Ok(item) => context.push(item),
|
||||
Err(e) => {
|
||||
// A render failure here would starve the LLM of the
|
||||
// notification text. Fall back to the raw message —
|
||||
// it still carries the intent, just without the
|
||||
// wrapper phrasing.
|
||||
warn!(error = %e, "failed to render notify_wrapper; using raw message");
|
||||
context.push(Item::system_message(notification.message.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let info = PreRequestInfo {
|
||||
item_count: context.len(),
|
||||
estimated_tokens: current_tokens,
|
||||
turn_index: self.current_turn_index(),
|
||||
tool_calls_this_turn: self.tool_calls_this_turn.load(Ordering::Relaxed),
|
||||
};
|
||||
for hook in &self.registry.pre_llm_request {
|
||||
let action = hook.call(&info).await;
|
||||
if !matches!(action, PreRequestAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
|
||||
async fn pre_tool_call(&self, info: &mut ToolCallInfo) -> PreToolAction {
|
||||
let summary = ToolCallSummary {
|
||||
call_id: info.call.id.clone(),
|
||||
tool_name: info.call.name.clone(),
|
||||
arguments: info.call.input.clone(),
|
||||
};
|
||||
for hook in &self.registry.pre_tool_call {
|
||||
let action = hook.call(&summary).await;
|
||||
if !matches!(action, PreToolAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
self.tool_calls_this_turn.fetch_add(1, Ordering::Relaxed);
|
||||
PreToolAction::Continue
|
||||
}
|
||||
|
||||
async fn post_tool_call(&self, info: &mut ToolResultInfo) -> PostToolAction {
|
||||
let summary = ToolResultSummary {
|
||||
call_id: info.result.tool_use_id.clone(),
|
||||
tool_name: info.call.name.clone(),
|
||||
is_error: info.result.is_error,
|
||||
output: ToolOutput {
|
||||
summary: info.result.summary.clone(),
|
||||
content: info.result.content.clone(),
|
||||
},
|
||||
};
|
||||
for hook in &self.registry.post_tool_call {
|
||||
let action = hook.call(&summary).await;
|
||||
if !matches!(action, PostToolAction::Continue) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
PostToolAction::Continue
|
||||
}
|
||||
|
||||
async fn on_turn_end(&self, history: &[Item]) -> TurnEndAction {
|
||||
let final_text_preview = history
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|i| i.is_assistant_message())
|
||||
.and_then(extract_message_text)
|
||||
.map(|t| preview(&t, FINAL_TEXT_PREVIEW_LIMIT))
|
||||
.unwrap_or_default();
|
||||
let info = TurnEndInfo {
|
||||
turn_index: self.current_turn_index(),
|
||||
tool_calls_count: self.tool_calls_this_turn.load(Ordering::Relaxed),
|
||||
final_text_preview,
|
||||
};
|
||||
for hook in &self.registry.on_turn_end {
|
||||
let action = hook.call(&info).await;
|
||||
if !matches!(action, TurnEndAction::Finish) {
|
||||
return action;
|
||||
}
|
||||
}
|
||||
TurnEndAction::Finish
|
||||
}
|
||||
|
||||
async fn on_abort(&self, reason: &str) {
|
||||
let info = AbortInfo {
|
||||
reason: reason.to_string(),
|
||||
};
|
||||
for hook in &self.registry.on_abort {
|
||||
hook.call(&info).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_message_text(item: &Item) -> Option<String> {
|
||||
match item {
|
||||
Item::Message { content, .. } => Some(
|
||||
content
|
||||
.iter()
|
||||
.map(|p| p.as_text())
|
||||
.collect::<Vec<_>>()
|
||||
.join(""),
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn preview(text: &str, limit: usize) -> String {
|
||||
if text.len() <= limit {
|
||||
return text.to_string();
|
||||
}
|
||||
let mut end = limit;
|
||||
while end > 0 && !text.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
text[..end].to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize};
|
||||
|
||||
use super::*;
|
||||
use crate::hook::{Hook, HookRegistryBuilder, PreLlmRequest};
|
||||
|
||||
struct CountingHook(Arc<AtomicUsize>);
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for CountingHook {
|
||||
async fn call(&self, _info: &PreRequestInfo) -> PreRequestAction {
|
||||
self.0.fetch_add(1, Ordering::Relaxed);
|
||||
PreRequestAction::Continue
|
||||
}
|
||||
}
|
||||
|
||||
fn registry_with_pre_llm_hook(counter: Arc<AtomicUsize>) -> Arc<HookRegistry> {
|
||||
let mut builder = HookRegistryBuilder::new();
|
||||
builder.add_pre_llm_request(CountingHook(counter));
|
||||
Arc::new(builder.build())
|
||||
}
|
||||
|
||||
/// Build a usage_history handle with a single record pinned at the
|
||||
/// current `context_len` so that `total_tokens_impl` returns exactly
|
||||
/// `tokens` (Measured, no interpolation or byte-based fallback).
|
||||
fn usage_handle_with(context_len: usize, tokens: u64) -> Arc<Mutex<Vec<UsageRecord>>> {
|
||||
Arc::new(Mutex::new(vec![UsageRecord {
|
||||
history_len: context_len,
|
||||
input_total_tokens: tokens,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
output_tokens: 0,
|
||||
}]))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_yields_and_skips_hooks_when_request_threshold_exceeded() {
|
||||
let count = Arc::new(AtomicUsize::new(0));
|
||||
let registry = registry_with_pre_llm_hook(count.clone());
|
||||
|
||||
let state = Arc::new(CompactState::new(None, Some(100), 2));
|
||||
let ctx_items = vec![Item::user_message("hi")];
|
||||
let history = usage_handle_with(ctx_items.len(), 200);
|
||||
|
||||
let interceptor = PodInterceptor::new(
|
||||
registry,
|
||||
Some(state),
|
||||
Some(history),
|
||||
NotificationBuffer::new(),
|
||||
PromptCatalog::builtins_only().unwrap(),
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Yield));
|
||||
// Hook must not run when an internal mechanism short-circuits first.
|
||||
assert_eq!(count.load(Ordering::Relaxed), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_runs_hooks_when_under_threshold() {
|
||||
let count = Arc::new(AtomicUsize::new(0));
|
||||
let registry = registry_with_pre_llm_hook(count.clone());
|
||||
|
||||
let state = Arc::new(CompactState::new(None, Some(100), 2));
|
||||
let ctx_items = vec![Item::user_message("hi")];
|
||||
let history = usage_handle_with(ctx_items.len(), 50);
|
||||
|
||||
let interceptor = PodInterceptor::new(
|
||||
registry,
|
||||
Some(state),
|
||||
Some(history),
|
||||
NotificationBuffer::new(),
|
||||
PromptCatalog::builtins_only().unwrap(),
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_does_not_yield_when_only_post_run_threshold_set() {
|
||||
// request_threshold = None → safety-net check is inert inside the turn
|
||||
// even if current occupancy is huge. Post-run check runs elsewhere.
|
||||
let count = Arc::new(AtomicUsize::new(0));
|
||||
let registry = registry_with_pre_llm_hook(count.clone());
|
||||
|
||||
let state = Arc::new(CompactState::new(Some(100), None, 2));
|
||||
let ctx_items = vec![Item::user_message("hi")];
|
||||
let history = usage_handle_with(ctx_items.len(), 10_000);
|
||||
|
||||
let interceptor = PodInterceptor::new(
|
||||
registry,
|
||||
Some(state),
|
||||
Some(history),
|
||||
NotificationBuffer::new(),
|
||||
PromptCatalog::builtins_only().unwrap(),
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_runs_hooks_when_no_compact_state() {
|
||||
let count = Arc::new(AtomicUsize::new(0));
|
||||
let registry = registry_with_pre_llm_hook(count.clone());
|
||||
|
||||
let interceptor = PodInterceptor::new(
|
||||
registry,
|
||||
None,
|
||||
None,
|
||||
NotificationBuffer::new(),
|
||||
PromptCatalog::builtins_only().unwrap(),
|
||||
);
|
||||
let mut ctx: Vec<Item> = Vec::new();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
assert_eq!(count.load(Ordering::Relaxed), 1);
|
||||
}
|
||||
|
||||
struct AbortingHook(Arc<AtomicBool>);
|
||||
|
||||
#[async_trait]
|
||||
impl Hook<PreLlmRequest> for AbortingHook {
|
||||
async fn call(&self, _info: &PreRequestInfo) -> PreRequestAction {
|
||||
self.0.store(true, Ordering::Relaxed);
|
||||
PreRequestAction::Cancel("nope".into())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_drains_pending_notifications_into_context() {
|
||||
let registry = Arc::new(HookRegistryBuilder::new().build());
|
||||
let buffer = NotificationBuffer::new();
|
||||
buffer.push("first".into());
|
||||
buffer.push("second".into());
|
||||
|
||||
let interceptor = PodInterceptor::new(
|
||||
registry,
|
||||
None,
|
||||
None,
|
||||
buffer.clone(),
|
||||
PromptCatalog::builtins_only().unwrap(),
|
||||
);
|
||||
let mut ctx: Vec<Item> = vec![Item::user_message("hi")];
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Continue));
|
||||
// Original user message preserved, two notifications appended in order.
|
||||
assert_eq!(ctx.len(), 3);
|
||||
let second = ctx[1].as_text().unwrap_or_default();
|
||||
let third = ctx[2].as_text().unwrap_or_default();
|
||||
assert!(second.contains("[Notification]"));
|
||||
assert!(second.contains("first"));
|
||||
assert!(third.contains("[Notification]"));
|
||||
assert!(third.contains("second"));
|
||||
// Buffer is drained after a single pre_llm_request call.
|
||||
assert!(buffer.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_skips_notification_injection_when_yielding() {
|
||||
// When compaction yields, notifications remain in the buffer for
|
||||
// the next pre_llm_request (after compaction + resume).
|
||||
let registry = Arc::new(HookRegistryBuilder::new().build());
|
||||
let buffer = NotificationBuffer::new();
|
||||
buffer.push("msg".into());
|
||||
|
||||
let state = Arc::new(CompactState::new(None, Some(100), 2));
|
||||
let ctx_items = vec![Item::user_message("hi")];
|
||||
let history = usage_handle_with(ctx_items.len(), 200);
|
||||
|
||||
let interceptor = PodInterceptor::new(
|
||||
registry,
|
||||
Some(state),
|
||||
Some(history),
|
||||
buffer.clone(),
|
||||
PromptCatalog::builtins_only().unwrap(),
|
||||
);
|
||||
let mut ctx = ctx_items;
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Yield));
|
||||
// Notifications were not drained (still held for post-compact resume).
|
||||
assert_eq!(ctx.len(), 1);
|
||||
assert_eq!(buffer.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pre_llm_request_short_circuits_on_first_non_continue() {
|
||||
let first_called = Arc::new(AtomicBool::new(false));
|
||||
let second_count = Arc::new(AtomicUsize::new(0));
|
||||
let mut builder = HookRegistryBuilder::new();
|
||||
builder.add_pre_llm_request(AbortingHook(first_called.clone()));
|
||||
builder.add_pre_llm_request(CountingHook(second_count.clone()));
|
||||
let registry = Arc::new(builder.build());
|
||||
|
||||
let interceptor = PodInterceptor::new(
|
||||
registry,
|
||||
None,
|
||||
None,
|
||||
NotificationBuffer::new(),
|
||||
PromptCatalog::builtins_only().unwrap(),
|
||||
);
|
||||
let mut ctx: Vec<Item> = Vec::new();
|
||||
let action = interceptor.pre_llm_request(&mut ctx).await;
|
||||
|
||||
assert!(matches!(action, PreRequestAction::Cancel(_)));
|
||||
assert!(first_called.load(Ordering::Relaxed));
|
||||
assert_eq!(second_count.load(Ordering::Relaxed), 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
pub mod event;
|
||||
pub mod notifier;
|
||||
pub mod server;
|
||||
|
||||
pub(crate) mod interceptor;
|
||||
pub(crate) mod notification_buffer;
|
||||
@@ -0,0 +1,126 @@
|
||||
//! Pending-notification buffer for `Method::Notify`.
|
||||
//!
|
||||
//! Notifications are queued here by the Controller and drained by
|
||||
//! `PodInterceptor::pre_llm_request` into the per-request context
|
||||
//! (never into the Worker's persistent history). Each queued entry
|
||||
//! becomes one `Item::system_message` in the outgoing request.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use llm_worker::Item;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::prompt::catalog::{CatalogError, PromptCatalog};
|
||||
|
||||
/// Maximum queued notifications. Oldest entries are dropped beyond this.
|
||||
const CAPACITY: usize = 128;
|
||||
|
||||
/// One pending notification awaiting injection into the next LLM request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PendingNotification {
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
/// Shared, mutex-guarded buffer of pending notifications.
|
||||
///
|
||||
/// Cloned between the Pod (producer) and PodInterceptor (consumer).
|
||||
#[derive(Clone, Default)]
|
||||
pub struct NotificationBuffer {
|
||||
inner: Arc<Mutex<VecDeque<PendingNotification>>>,
|
||||
}
|
||||
|
||||
impl NotificationBuffer {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Push a notification onto the queue. If the queue is full, the
|
||||
/// oldest entry is dropped and a `tracing::warn` is emitted — the
|
||||
/// caller should never hit this in normal operation.
|
||||
pub fn push(&self, message: String) {
|
||||
let mut q = self.inner.lock().expect("notification buffer poisoned");
|
||||
if q.len() >= CAPACITY {
|
||||
let dropped = q.pop_front();
|
||||
warn!(
|
||||
capacity = CAPACITY,
|
||||
dropped_message = dropped.as_ref().map(|n| n.message.as_str()),
|
||||
"notification buffer overflow; dropped oldest"
|
||||
);
|
||||
}
|
||||
q.push_back(PendingNotification { message });
|
||||
}
|
||||
|
||||
/// Remove and return all pending notifications in FIFO order.
|
||||
pub fn drain(&self) -> Vec<PendingNotification> {
|
||||
let mut q = self.inner.lock().expect("notification buffer poisoned");
|
||||
q.drain(..).collect()
|
||||
}
|
||||
|
||||
/// Number of pending notifications. Primarily for tests.
|
||||
pub fn len(&self) -> usize {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("notification buffer poisoned")
|
||||
.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
/// Format a single pending notification into the `Item::system_message`
|
||||
/// that gets injected into the per-request context. The wrapper body
|
||||
/// comes from `PodPrompt::NotifyWrapper` so the surrounding phrasing
|
||||
/// can be customised via a prompt pack (translation, tone, ...).
|
||||
pub(crate) fn format_notification(
|
||||
n: &PendingNotification,
|
||||
prompts: &PromptCatalog,
|
||||
) -> Result<Item, CatalogError> {
|
||||
let text = prompts.notify_wrapper(&n.message)?;
|
||||
Ok(Item::system_message(text))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn push_then_drain_preserves_order() {
|
||||
let buf = NotificationBuffer::new();
|
||||
buf.push("one".into());
|
||||
buf.push("two".into());
|
||||
let drained = buf.drain();
|
||||
assert_eq!(drained.len(), 2);
|
||||
assert_eq!(drained[0].message, "one");
|
||||
assert_eq!(drained[1].message, "two");
|
||||
assert!(buf.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capacity_drops_oldest() {
|
||||
let buf = NotificationBuffer::new();
|
||||
for i in 0..(CAPACITY + 5) {
|
||||
buf.push(format!("msg{i}"));
|
||||
}
|
||||
let drained = buf.drain();
|
||||
assert_eq!(drained.len(), CAPACITY);
|
||||
// Oldest 5 were dropped; first retained is msg5.
|
||||
assert_eq!(drained[0].message, "msg5");
|
||||
assert_eq!(drained[CAPACITY - 1].message, format!("msg{}", CAPACITY + 4));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_notification_includes_message_and_nonblocking_hint() {
|
||||
let n = PendingNotification {
|
||||
message: "hello".into(),
|
||||
};
|
||||
let catalog = PromptCatalog::builtins_only().unwrap();
|
||||
let item = format_notification(&n, &catalog).unwrap();
|
||||
let text = item.as_text().unwrap_or_default().to_string();
|
||||
assert!(text.contains("[Notification]"));
|
||||
assert!(text.contains("hello"));
|
||||
assert!(text.contains("not a blocking request"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
//! User-facing notification channel for Pod → client.
|
||||
//!
|
||||
//! Separate from `tracing` (which is for developer logs). Notifications
|
||||
//! are short human-readable messages the Pod layer wants a client to
|
||||
//! see — for example "compaction failed", "tool output truncated".
|
||||
//!
|
||||
//! Each notification is broadcast on the shared `Event` channel and
|
||||
//! also appended to an in-memory buffer so that clients connecting
|
||||
//! after the fact still see everything emitted during the session.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use protocol::{Event, Notification, NotificationLevel, NotificationSource};
|
||||
|
||||
/// Upper bound on buffered notifications. When exceeded, the oldest
|
||||
/// entries are discarded so a long-running session cannot leak
|
||||
/// memory through a pathological loop of recurring notifications
|
||||
/// (e.g. compaction failing every turn).
|
||||
const MAX_BUFFERED_NOTIFICATIONS: usize = 512;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Notifier {
|
||||
inner: Arc<Inner>,
|
||||
}
|
||||
|
||||
struct Inner {
|
||||
event_tx: broadcast::Sender<Event>,
|
||||
buffer: Mutex<VecDeque<Notification>>,
|
||||
}
|
||||
|
||||
impl Notifier {
|
||||
pub fn new(event_tx: broadcast::Sender<Event>) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Inner {
|
||||
event_tx,
|
||||
buffer: Mutex::new(VecDeque::with_capacity(MAX_BUFFERED_NOTIFICATIONS)),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record and broadcast a notification.
|
||||
///
|
||||
/// The broadcast may have no subscribers (e.g. during Pod
|
||||
/// construction before any client has connected); the buffer
|
||||
/// guarantees the message is still delivered once a client
|
||||
/// attaches.
|
||||
///
|
||||
/// The buffer mutex is held across `broadcast::send` to make
|
||||
/// `subscribe_with_snapshot` race-free — a client that snapshots
|
||||
/// the buffer while holding the same lock sees every notification
|
||||
/// exactly once: older ones from the snapshot, newer ones from
|
||||
/// the freshly-subscribed receiver.
|
||||
pub fn notify(&self, level: NotificationLevel, source: NotificationSource, message: String) {
|
||||
let notification = Notification {
|
||||
level,
|
||||
source,
|
||||
message,
|
||||
timestamp_ms: now_ms(),
|
||||
};
|
||||
if let Ok(mut buf) = self.inner.buffer.lock() {
|
||||
if buf.len() >= MAX_BUFFERED_NOTIFICATIONS {
|
||||
buf.pop_front();
|
||||
}
|
||||
buf.push_back(notification.clone());
|
||||
let _ = self
|
||||
.inner
|
||||
.event_tx
|
||||
.send(Event::Notification(notification));
|
||||
}
|
||||
}
|
||||
|
||||
/// Subscribe and atomically snapshot the current buffer.
|
||||
///
|
||||
/// The returned snapshot contains notifications emitted before
|
||||
/// this call; the receiver will deliver notifications emitted
|
||||
/// after. A notification cannot appear in both.
|
||||
pub fn subscribe_with_snapshot(&self) -> (Vec<Notification>, broadcast::Receiver<Event>) {
|
||||
let buf = self
|
||||
.inner
|
||||
.buffer
|
||||
.lock()
|
||||
.expect("notifier buffer mutex poisoned");
|
||||
let rx = self.inner.event_tx.subscribe();
|
||||
let snapshot: Vec<Notification> = buf.iter().cloned().collect();
|
||||
(snapshot, rx)
|
||||
}
|
||||
}
|
||||
|
||||
fn now_ms() -> i64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn notify_broadcasts_to_existing_subscriber() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let notifier = Notifier::new(tx);
|
||||
let (_snapshot, mut rx) = notifier.subscribe_with_snapshot();
|
||||
|
||||
notifier.notify(
|
||||
NotificationLevel::Warn,
|
||||
NotificationSource::Compactor,
|
||||
"test message".into(),
|
||||
);
|
||||
|
||||
match rx.try_recv() {
|
||||
Ok(Event::Notification(n)) => assert_eq!(n.message, "test message"),
|
||||
other => panic!("unexpected event: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn late_subscriber_sees_earlier_notifications_via_snapshot() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let notifier = Notifier::new(tx);
|
||||
|
||||
notifier.notify(
|
||||
NotificationLevel::Error,
|
||||
NotificationSource::Pod,
|
||||
"first".into(),
|
||||
);
|
||||
notifier.notify(
|
||||
NotificationLevel::Warn,
|
||||
NotificationSource::AgentsMd,
|
||||
"second".into(),
|
||||
);
|
||||
|
||||
let (snapshot, mut rx) = notifier.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), 2);
|
||||
assert_eq!(snapshot[0].message, "first");
|
||||
assert_eq!(snapshot[1].message, "second");
|
||||
assert!(rx.try_recv().is_err()); // nothing pending on the receiver
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn buffer_discards_oldest_past_cap() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(1024);
|
||||
let notifier = Notifier::new(tx);
|
||||
|
||||
for i in 0..(MAX_BUFFERED_NOTIFICATIONS + 50) {
|
||||
notifier.notify(
|
||||
NotificationLevel::Warn,
|
||||
NotificationSource::Worker,
|
||||
format!("msg-{i}"),
|
||||
);
|
||||
}
|
||||
|
||||
let (snapshot, _rx) = notifier.subscribe_with_snapshot();
|
||||
assert_eq!(snapshot.len(), MAX_BUFFERED_NOTIFICATIONS);
|
||||
// First 50 were evicted; the oldest remaining is msg-50.
|
||||
assert_eq!(snapshot.first().unwrap().message, "msg-50");
|
||||
let last = format!("msg-{}", MAX_BUFFERED_NOTIFICATIONS + 49);
|
||||
assert_eq!(snapshot.last().unwrap().message, last);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subscribe_snapshot_and_live_do_not_overlap() {
|
||||
let (tx, _keep) = broadcast::channel::<Event>(8);
|
||||
let notifier = Notifier::new(tx);
|
||||
|
||||
notifier.notify(
|
||||
NotificationLevel::Warn,
|
||||
NotificationSource::Worker,
|
||||
"historic".into(),
|
||||
);
|
||||
let (snapshot, mut rx) = notifier.subscribe_with_snapshot();
|
||||
notifier.notify(
|
||||
NotificationLevel::Error,
|
||||
NotificationSource::Worker,
|
||||
"live".into(),
|
||||
);
|
||||
|
||||
assert_eq!(snapshot.len(), 1);
|
||||
assert_eq!(snapshot[0].message, "historic");
|
||||
match rx.try_recv() {
|
||||
Ok(Event::Notification(n)) => assert_eq!(n.message, "live"),
|
||||
other => panic!("unexpected: {other:?}"),
|
||||
}
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
use std::io;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use protocol::stream::{JsonLineReader, JsonLineWriter};
|
||||
use tokio::net::UnixListener;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
use crate::controller::PodHandle;
|
||||
use protocol::{Event, Method};
|
||||
|
||||
/// Unix socket server for Pod Protocol.
|
||||
///
|
||||
/// Listens on the Pod's runtime directory socket path.
|
||||
/// Each client connection gets bidirectional JSONL:
|
||||
/// - Client writes Method lines → forwarded to PodController
|
||||
/// - Pod events → written as Event lines to all connected clients
|
||||
pub struct SocketServer {
|
||||
_accept_task: JoinHandle<()>,
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl SocketServer {
|
||||
/// Start listening on the PodHandle's socket path.
|
||||
pub async fn start(handle: &PodHandle) -> Result<Self, io::Error> {
|
||||
let path = handle.runtime_dir.socket_path();
|
||||
|
||||
// Remove stale socket file if it exists
|
||||
let _ = tokio::fs::remove_file(&path).await;
|
||||
|
||||
let listener = UnixListener::bind(&path)?;
|
||||
let handle = handle.clone();
|
||||
|
||||
let _accept_task = tokio::spawn(async move {
|
||||
loop {
|
||||
match listener.accept().await {
|
||||
Ok((stream, _)) => {
|
||||
let handle = handle.clone();
|
||||
tokio::spawn(handle_connection(stream, handle));
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Self { _accept_task, path })
|
||||
}
|
||||
|
||||
/// The socket file path.
|
||||
pub fn path(&self) -> &std::path::Path {
|
||||
&self.path
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SocketServer {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_file(&self.path);
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_connection(stream: tokio::net::UnixStream, handle: PodHandle) {
|
||||
let (reader, writer) = stream.into_split();
|
||||
let mut reader = JsonLineReader::new(reader);
|
||||
let mut writer = JsonLineWriter::new(writer);
|
||||
|
||||
// Atomically subscribe and snapshot buffered notifications so that
|
||||
// warnings emitted before this client connected are replayed
|
||||
// exactly once — they appear in the snapshot, and any notification
|
||||
// arriving afterwards reaches us through `rx`.
|
||||
let (notification_snapshot, mut rx) = handle.notifier.subscribe_with_snapshot();
|
||||
for notification in notification_snapshot {
|
||||
if writer
|
||||
.write(&Event::Notification(notification))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// Broadcast events → this client
|
||||
event = rx.recv() => {
|
||||
match event {
|
||||
Ok(event) => {
|
||||
if writer.write(&event).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
// Client methods → handle or forward to controller
|
||||
method = reader.next::<Method>() => {
|
||||
match method {
|
||||
Ok(Some(Method::GetHistory)) => {
|
||||
let items = handle.shared_state.history();
|
||||
let values = items
|
||||
.iter()
|
||||
.map(|item| serde_json::to_value(item).expect("Item is Serialize"))
|
||||
.collect();
|
||||
let greeting = handle.shared_state.greeting.clone();
|
||||
if writer
|
||||
.write(&Event::History {
|
||||
items: values,
|
||||
greeting,
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Some(method)) => {
|
||||
let _ = handle.send(method).await;
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(e) => {
|
||||
let _ = handle.send_event(Event::Error {
|
||||
code: protocol::ErrorCode::Internal,
|
||||
message: format!("invalid method: {e}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user