Something went wrong. Try again.
Personal fork of prime agent. Mirror of: https://mnemosyne.sigilzero.dev/sigilzero/prime-agent
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429//! The session runtime: shared goal-driver + cron-store state with kernel//! host-handler registration — the wiring layer binding the//! goal/rlm-heartbeat host requests to a live session (the Rust//! equivalent of the `AgentSession` host-request controllers).
use std::sync::Arc;
use serde_json::Value;use tokio::sync::Mutex;
use crate::cron::store::AgentCronJobStore;use crate::goals::GoalHostResponse;use crate::kernel::shared::{host_handler, HostRequestHandlers};use crate::session::manager::SessionManager;
use super::goal_driver::GoalDriver;use super::host_requests::{ handle_goal_host_request, handle_rlm_heartbeat_host_request, RlmHeartbeatMutationHook, SessionBinding,};
/// The host-side purge of queued goal-context turns: the queue lanes live in the daemon/// worker, so the completing kernel host request invokes this seam — dropping a/// continuation queued while the goal was completing.pub type QueuedGoalContextPurge = Arc<dyn Fn() + Send + Sync>;
/// Session-scoped runtime state the kernel bridge reaches.pub struct SessionRuntime { goal_driver: Arc<Mutex<GoalDriver>>, cron_store: Arc<AgentCronJobStore>, active_session_id: String, binding: SessionBinding, /// Invoked after a kernel `goal.complete` settles the goal; `None` /// when the embedding owns no queued goal contexts. goal_complete_purge: Option<QueuedGoalContextPurge>, /// Invoked after a kernel `rlm_heartbeat.*` mutation; `None` when /// the embedding owns no scheduler to re-arm. cron_mutation_hook: Option<RlmHeartbeatMutationHook>,}
impl SessionRuntime { /// Build a runtime around a cron store, rehydrating any persisted goal. pub fn new( session: &SessionManager, cron_store: Arc<AgentCronJobStore>, active_session_id: String, binding: SessionBinding, ) -> Self { Self { goal_driver: Arc::new(Mutex::new(GoalDriver::load_persisted(session))), cron_store, active_session_id, binding, goal_complete_purge: None, cron_mutation_hook: None, } }
/// Set the post-completion purge seam (the daemon worker's queue purge). pub fn set_goal_complete_purge(&mut self, purge: QueuedGoalContextPurge) { self.goal_complete_purge = Some(purge); }
/// Set the rlm heartbeat mutation hook (queued-fire withdrawal + /// scheduler re-arm). pub fn set_cron_mutation_hook(&mut self, hook: RlmHeartbeatMutationHook) { self.cron_mutation_hook = Some(hook); }
#[must_use] pub fn goal_driver(&self) -> &Arc<Mutex<GoalDriver>> { &self.goal_driver }
#[must_use] pub fn cron_store(&self) -> &Arc<AgentCronJobStore> { &self.cron_store }
/// Register the `goal.*` and `rlm_heartbeat.*` handlers onto a handler map. pub fn register_host_handlers( &self, session: Arc<Mutex<SessionManager>>, handlers: &mut HostRequestHandlers, ) { let driver = self.goal_driver.clone(); let goal_session = session; // TS `_completeGoalFromHost` clears the queued goal contexts: a continuation // queued behind the completing turn never runs post-completion. let goal_complete_purge = self.goal_complete_purge.clone(); for request_type in ["goal.get", "goal.create", "goal.complete"] { let driver = driver.clone(); let session = goal_session.clone(); let purge = if request_type == "goal.complete" { goal_complete_purge.clone() } else { None }; handlers.register( request_type, host_handler(move |payload| { let driver = driver.clone(); let session = session.clone(); let purge = purge.clone(); Box::pin(async move { let response = { super::goal_driver::persist_goal_mutation( &driver, &session, |driver, drafts| { handle_goal_host_request( request_type, &payload.data, driver, drafts, ) }, ) .await? }; if let Some(purge) = purge { purge(); } host_ok(&response) }) }), ); } let store = self.cron_store.clone(); let active_session_id = self.active_session_id.clone(); let binding_session_id = self.binding.session_id.clone(); let binding_session_file = self.binding.session_file.clone(); let binding_cwd = self.binding.cwd.clone(); let cron_mutation_hook = self.cron_mutation_hook.clone(); for request_type in [ "rlm_heartbeat.list", "rlm_heartbeat.create", "rlm_heartbeat.update", "rlm_heartbeat.delete", ] { let store = store.clone(); let active_session_id = active_session_id.clone(); let session_id = binding_session_id.clone(); let session_file = binding_session_file.clone(); let cwd = binding_cwd.clone(); let mutation_hook = cron_mutation_hook.clone(); handlers.register( request_type, host_handler(move |payload| { let store = store.clone(); let active_session_id = active_session_id.clone(); let binding = SessionBinding { session_id: session_id.clone(), session_file: session_file.clone(), cwd: cwd.clone(), }; let mutation_hook = mutation_hook.clone(); Box::pin(async move { let outcome = handle_rlm_heartbeat_host_request( payload .data .get("type") .and_then(Value::as_str) .unwrap_or("rlm_heartbeat.list"), &payload.data, &store, &active_session_id, &binding, )?; // The embedding's post-mutation work: withdraw the // queued fire and re-arm the cron scheduler. if let Some(mutation) = outcome.mutation { if let Some(hook) = &mutation_hook { hook(mutation).await; } } Ok(outcome.response) }) }), ); } }}
fn host_ok(response: &GoalHostResponse) -> anyhow::Result<Value> { serde_json::to_value(response).map_err(anyhow::Error::new)}
#[cfg(test)]mod tests { use super::*; use crate::kernel::shared::HostRequestPayload;
fn persisted_session() -> (tempfile::TempDir, SessionManager) { let dir = tempfile::TempDir::new().unwrap(); let session_dir = dir.path().join("session"); std::fs::create_dir_all(&session_dir).unwrap(); let mut session = SessionManager::in_memory(dir.path()); session.materialize_session_file(Some(session_dir)).unwrap(); (dir, session) }
fn payload(data: Value) -> HostRequestPayload { HostRequestPayload { data } }
/// The purge seam fires only for `goal.complete`, never on the way in. #[tokio::test] async fn goal_complete_purge_runs_only_after_goal_complete() { let (_session_dir, session) = persisted_session(); let session = Arc::new(Mutex::new(session)); let dir = tempfile::TempDir::new().unwrap(); let mut runtime = SessionRuntime::new( &*session.lock().await, Arc::new(AgentCronJobStore::new(dir.path().join("jobs.json"))), "live-1".to_string(), SessionBinding { session_id: "session-1".to_string(), session_file: "/w/s.jsonl".to_string(), cwd: "/w".to_string(), }, ); let purges = Arc::new(std::sync::atomic::AtomicUsize::new(0)); let sink = Arc::clone(&purges); runtime.set_goal_complete_purge(Arc::new(move || { sink.fetch_add(1, std::sync::atomic::Ordering::SeqCst); })); let mut handlers = HostRequestHandlers::default(); runtime.register_host_handlers(session, &mut handlers);
let create = handlers.get("goal.create").unwrap().clone(); create(payload(serde_json::json!({ "type": "goal.create", "objective": "pin the purge ordering" }))) .await .unwrap(); let get = handlers.get("goal.get").unwrap().clone(); get(payload(serde_json::json!({ "type": "goal.get" }))) .await .unwrap(); assert_eq!(purges.load(std::sync::atomic::Ordering::SeqCst), 0);
let complete = handlers.get("goal.complete").unwrap().clone(); complete(payload(serde_json::json!({ "type": "goal.complete" }))) .await .unwrap(); assert_eq!(purges.load(std::sync::atomic::Ordering::SeqCst), 1); }
#[tokio::test] async fn goal_handlers_round_trip_through_the_registry() { let (_session_dir, session) = persisted_session(); let session = Arc::new(Mutex::new(session)); let dir = tempfile::TempDir::new().unwrap(); let runtime = SessionRuntime::new( &*session.lock().await, Arc::new(AgentCronJobStore::new(dir.path().join("jobs.json"))), "live-1".to_string(), SessionBinding { session_id: "session-1".to_string(), session_file: "/w/s.jsonl".to_string(), cwd: "/w".to_string(), }, ); let mut handlers = HostRequestHandlers::default(); runtime.register_host_handlers(session.clone(), &mut handlers);
let create = handlers.get("goal.create").unwrap().clone(); let response = create(payload(serde_json::json!({ "type": "goal.create", "objective": "finish the port", "token_budget": 1000 }))) .await .unwrap(); assert_eq!(response["goal"]["status"], "active"); assert_eq!(response["goal"]["objective"], "finish the port"); let persisted = session.lock().await; let goal_entries: Vec<&pa_types::session::FileEntry> = persisted .get_all_entries() .iter() .filter(|entry| { matches!(entry, pa_types::session::FileEntry::Custom { payload, .. } if payload.custom_type == crate::goals::GOAL_STATE_CUSTOM_TYPE) }) .collect(); assert_eq!(goal_entries.len(), 1); drop(persisted);
let get = handlers.get("goal.get").unwrap().clone(); let response = get(payload(serde_json::json!({ "type": "goal.get" }))) .await .unwrap(); assert_eq!(response["goal"]["objective"], "finish the port"); assert_eq!(response["remaining_tokens"], 1000);
let complete = handlers.get("goal.complete").unwrap().clone(); let response = complete(payload(serde_json::json!({ "type": "goal.complete" }))) .await .unwrap(); assert_eq!(response["goal"]["status"], "complete"); assert!(response["completion_budget_report"].is_string()); }
#[tokio::test] async fn heartbeat_handlers_round_trip_through_the_registry() { let (_session_dir, session) = persisted_session(); let session = Arc::new(Mutex::new(session)); let dir = tempfile::TempDir::new().unwrap(); let runtime = SessionRuntime::new( &*session.lock().await, Arc::new(AgentCronJobStore::new(dir.path().join("jobs.json"))), "live-1".to_string(), SessionBinding { session_id: "session-1".to_string(), session_file: "/w/s.jsonl".to_string(), cwd: "/w".to_string(), }, ); let mut handlers = HostRequestHandlers::default(); runtime.register_host_handlers(session, &mut handlers);
let create = handlers.get("rlm_heartbeat.create").unwrap().clone(); let response = create(payload(serde_json::json!({ "type": "rlm_heartbeat.create", "instruction": "watch the mission", "interval": "every 15m" }))) .await .unwrap(); assert_eq!(response["heartbeat"]["status"], "active"); assert_eq!(response["heartbeat"]["instruction"], "watch the mission"); let id = response["heartbeat"]["id"].as_str().unwrap().to_string();
let update = handlers.get("rlm_heartbeat.update").unwrap().clone(); let response = update(payload(serde_json::json!({ "type": "rlm_heartbeat.update", "id": id, "status": "pause" }))) .await .unwrap(); assert_eq!(response["heartbeat"]["status"], "paused"); let jobs = runtime.cron_store().list_rlm_heartbeats("live-1", true); assert_eq!(jobs.len(), 1); }
/// Catalog reads announce nothing. #[tokio::test] async fn rlm_heartbeat_mutations_invoke_the_cron_mutation_hook() { let (_session_dir, session) = persisted_session(); let session = Arc::new(Mutex::new(session)); let dir = tempfile::TempDir::new().unwrap(); let mut runtime = SessionRuntime::new( &*session.lock().await, Arc::new(AgentCronJobStore::new(dir.path().join("jobs.json"))), "live-1".to_string(), SessionBinding { session_id: "session-1".to_string(), session_file: "/w/s.jsonl".to_string(), cwd: "/w".to_string(), }, ); let mutations: Arc<Mutex<Vec<super::super::host_requests::RlmHeartbeatMutation>>> = Arc::new(Mutex::new(Vec::new())); let sink = mutations.clone(); runtime.set_cron_mutation_hook(Arc::new(move |mutation| { let sink = sink.clone(); Box::pin(async move { sink.lock().await.push(mutation); }) })); let mut handlers = HostRequestHandlers::default(); runtime.register_host_handlers(session, &mut handlers);
// Create announces the job (never withdrawing a queued fire). let create = handlers.get("rlm_heartbeat.create").unwrap().clone(); let response = create(payload(serde_json::json!({ "type": "rlm_heartbeat.create", "instruction": "watch the mission", "interval": "every 15m" }))) .await .unwrap(); let id = response["heartbeat"]["id"].as_str().unwrap().to_string(); let seen = mutations.lock().await; assert_eq!(seen.len(), 1); assert_eq!(seen[0].job.id, id); assert_eq!(seen[0].job.source.as_deref(), Some("rlm_heartbeat")); assert!(!seen[0].drop_queued); drop(seen);
// A catalog read announces nothing. let list = handlers.get("rlm_heartbeat.list").unwrap().clone(); let response = list(payload(serde_json::json!({ "type": "rlm_heartbeat.list" }))) .await .unwrap(); assert_eq!(response["heartbeats"].as_array().unwrap().len(), 1); assert_eq!(mutations.lock().await.len(), 1);
// A pause withdraws the queued fire (TS `updateRlmHeartbeatForState`). let update = handlers.get("rlm_heartbeat.update").unwrap().clone(); update(payload(serde_json::json!({ "type": "rlm_heartbeat.update", "id": id, "status": "pause" }))) .await .unwrap(); let seen = mutations.lock().await; assert_eq!(seen.len(), 2); assert!(seen[1].drop_queued); drop(seen);
// A delete always withdraws it. let delete = handlers.get("rlm_heartbeat.delete").unwrap().clone(); delete(payload(serde_json::json!({ "type": "rlm_heartbeat.delete", "id": id }))) .await .unwrap(); let seen = mutations.lock().await; assert_eq!(seen.len(), 3); assert!(seen[2].drop_queued); }}