Something went wrong. Try again.
Personal fork of prime agent. Mirror of: https://mnemosyne.sigilzero.dev/sigilzero/prime-agent
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298//! Execution bookkeeping: request writes, interrupts, abort forcing, and//! active-execution resolution plus late-sent agent message handlers.
use super::{ anyhow, json, lock, parse_sent_agent_message, AbortSignal, ActiveExecution, Arc, AsyncWriteExt, Duration, ExecuteResult, ExecuteStatus, Inner, Instant, InternalExecuteResult, KernelState, LateSentAgentMessageCallback, Value, AGENT_MESSAGE_DISPLAY_MIME, KERNEL_BUSY_INTERRUPT_INTERVAL_MS, KERNEL_BUSY_REUSE_WAIT_MS, MAX_BACKGROUND_OUTPUT_CHARS, MAX_LATE_SENT_AGENT_MESSAGE_HANDLERS,};use std::fmt::Write as _;
use crate::kernel::shared::KernelBusyAfterInterruptError;
impl Inner { /// Write one JSON-lines request frame; completes when the OS accepted the bytes. pub(crate) async fn write_line(&self, frame: &Value) -> anyhow::Result<()> { let stdin = { let child = lock(&self.child); child .as_ref() .map(|c| c.stdin.clone()) .ok_or_else(|| anyhow!("Kernel stdin is not connected"))? }; let mut guard = stdin.lock().await; let Some(stdin) = guard.as_mut() else { return Err(anyhow!("Kernel stdin is not connected")); }; let mut line = frame.to_string(); line.push('\n'); stdin.write_all(line.as_bytes()).await?; stdin.flush().await?; Ok(()) }
pub(crate) async fn interrupt(&self, id: Option<&str>) -> anyhow::Result<()> { let mut frame = json!({ "type": "interrupt" }); if let Some(id) = id { frame["id"] = json!(id); } self.write_line(&frame).await }
pub(crate) fn force_abort(&self, execution: &Arc<ActiveExecution>) { { let g = lock(&self.guarded); let is_active = matches!(g.active_execution.as_ref(), Some(active) if Arc::ptr_eq(active, execution)); if !is_active { return; } } lock(&execution.buffers).status = ExecuteStatus::Aborted; self.resolve_execution(execution, false); }
/// Wait until no execution is active, interrupting a busy cell, bounded by /// the busy-reuse window. pub(crate) async fn wait_for_active_execution_to_clear_for_reuse( self: &Arc<Self>, signal: Option<&AbortSignal>, ) -> anyhow::Result<()> { let started = Instant::now(); loop { let notified = self.busy_notify.notified(); { let g = lock(&self.guarded); if g.active_execution.is_none() { return Ok(()); } if g.state == KernelState::Shutdown { return Err(anyhow!("Kernel has been shut down")); } } let elapsed = started.elapsed(); if elapsed >= Duration::from_millis(KERNEL_BUSY_REUSE_WAIT_MS) { break; } let wait_ms = (KERNEL_BUSY_REUSE_WAIT_MS - elapsed.as_millis() as u64) .min(KERNEL_BUSY_INTERRUPT_INTERVAL_MS); let wait = tokio::time::sleep(Duration::from_millis(wait_ms.max(1))); let active_id = lock(&self.guarded) .active_execution .as_ref() .map(|a| a.request_id.clone()); let interrupt = { let inner = Arc::clone(self); tokio::spawn(async move { let _ = inner.interrupt(active_id.as_deref()).await; }) }; tokio::select! { () = notified => {} () = wait => {} () = async { match signal { Some(signal) => signal.cancelled().await, None => std::future::pending().await, } } => { interrupt.abort(); return Ok(()); } } } if lock(&self.guarded).active_execution.is_some() { // Typed on purpose: the ipython tool's wait/kill recovery branches on // this refusal, so it must survive the trip to the tool layer. return Err(KernelBusyAfterInterruptError.into()); } Ok(()) }
fn notify_active_execution_idle(&self) { self.busy_notify.notify_waiters(); }
pub(crate) fn reject_active_execution(&self, message: &str) { let execution = { let mut g = lock(&self.guarded); g.active_execution.take() }; if let Some(execution) = execution { if let Some(tx) = lock(&execution.result_tx).take() { let _ = tx.send(Err(anyhow!("{message}"))); } self.notify_active_execution_idle(); } }
pub(crate) fn finish_active_execution(&self, execution: &Arc<ActiveExecution>) { let is_active = { let g = lock(&self.guarded); matches!(g.active_execution.as_ref(), Some(active) if Arc::ptr_eq(active, execution)) }; if !is_active { return; } self.resolve_execution(execution, true); }
fn resolve_execution(&self, execution: &Arc<ActiveExecution>, clear_active: bool) { let did_clear_active = if clear_active { let mut g = lock(&self.guarded); matches!(g.active_execution.as_ref(), Some(active) if Arc::ptr_eq(active, execution)) .then(|| { g.active_execution = None; true }) .unwrap_or(false) } else { false }; let mut buffers = lock(&execution.buffers); if !buffers.settled { buffers.settled = true; { let mut g = lock(&self.guarded); g.completed_executions += 1; if !execution.opts.internal { g.user_executions += 1; } // The freshness memo describes the namespace as of its commit. A settled request // that runs user-namespace code or replaces the namespace ends that description. if execution.namespace_code || execution.restores_namespace { g.capture_freshness = None; g.freshness_epoch += 1; } } if let Some(callback) = execution.opts.on_late_sent_agent_message.clone() { self.register_late_sent_agent_message_handler(&execution.request_id, callback); }
let mut stdout = std::mem::take(&mut buffers.stdout); buffers.stdout_chars = 0; let mut stderr = std::mem::take(&mut buffers.stderr); buffers.stderr_chars = 0; let mut result = buffers.result.take(); let mut status = buffers.status; if buffers.stdout_truncated { let _ = write!( stdout, "\n[... output truncated at {} chars ...]", execution.max_chars ); } if buffers.stderr_truncated { let _ = write!( stderr, "\n[... output truncated at {} chars ...]", execution.max_chars ); } if let Some(text) = &result { if let Some((end, _)) = text.char_indices().nth(execution.max_chars) { // Match the Unicode-scalar limit used for stdout and stderr. let mut clipped = text[..end].to_string(); let _ = write!( clipped, "\n[... output truncated at {} chars ...]", execution.max_chars ); result = Some(clipped); } } if execution .opts .signal .as_ref() .is_some_and(AbortSignal::is_aborted) { status = ExecuteStatus::Aborted; }
let mut background_output = std::mem::take(&mut buffers.background_output); buffers.background_output_chars = 0; if buffers.background_output_truncated { let _ = write!(background_output, "\n[... background output truncated at {MAX_BACKGROUND_OUTPUT_CHARS} chars ...]", ); } let done_fields = buffers.done_fields.take(); let result = ExecuteResult { stdout, stderr, result, diffs: (!buffers.diffs.is_empty()).then(|| std::mem::take(&mut buffers.diffs)), attachments: (!buffers.attachments.is_empty()) .then(|| std::mem::take(&mut buffers.attachments)), sent_agent_messages: (!buffers.sent_agent_messages.is_empty()) .then(|| std::mem::take(&mut buffers.sent_agent_messages)), bash_commands: buffers.bash_commands.take(), background_output: (!background_output.is_empty()).then_some(background_output), status, error: buffers.error.take(), duration_ms: execution.started.elapsed().as_millis() as u64, }; drop(buffers); if let Some(tx) = lock(&execution.result_tx).take() { let _ = tx.send(Ok(InternalExecuteResult { result, done_fields, })); } } if did_clear_active { self.notify_active_execution_idle(); } }
fn register_late_sent_agent_message_handler( &self, request_id: &str, callback: LateSentAgentMessageCallback, ) { let mut g = lock(&self.guarded); g.late_handlers .retain(|(existing, _)| existing != request_id); g.late_handlers .push_back((request_id.to_string(), callback)); while g.late_handlers.len() > MAX_LATE_SENT_AGENT_MESSAGE_HANDLERS { g.late_handlers.pop_front(); } }
pub(crate) fn dispatch_late_sent_agent_message( &self, request_id: Option<&str>, data: &Value, ) -> bool { let Some(request_id) = request_id else { return false; }; let Some(payload) = data.get(AGENT_MESSAGE_DISPLAY_MIME) else { return false; }; let Some(message) = parse_sent_agent_message(payload) else { return false; }; let callback = { let mut g = lock(&self.guarded); let Some(position) = g.late_handlers.iter().position(|(id, _)| id == request_id) else { return false; }; let callback = g .late_handlers .remove(position) .expect("position checked") .1; // Refresh recency, matching the TS map delete+set. g.late_handlers .push_back((request_id.to_string(), callback.clone())); callback }; callback(message); true }}