Something went wrong. Try again.
Personal fork of prime agent. Mirror of: https://mnemosyne.sigilzero.dev/sigilzero/prime-agent
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373use std::{ collections::BTreeMap, sync::{ atomic::{AtomicUsize, Ordering}, Arc, },};
use mlua::{Function, HookTriggers, Lua, LuaSerdeExt, Table, VmState};use serde_json::Value;use tokio::sync::mpsc;use tokio_util::sync::CancellationToken;use uuid::Uuid;
use super::{ CapabilityRequest, ContextDataMode, Contribution, ContributionKind, HostCapabilities, InvocationEvent, InvocationMode, InvocationOutcome, PluginError, PluginSource,};
struct Registered { contribution: Contribution, callback: Function,}
pub(super) struct Generation { pub source: String, pub config: Value, pub retirement: CancellationToken, lua: Lua, callbacks: BTreeMap<(ContributionKind, String), Registered>,}
// Result::map_err needs an owned-error adapter; no Lua error escapes the module.#[allow(clippy::needless_pass_by_value)]fn lua_error(error: mlua::Error) -> PluginError { PluginError::Lua(error.to_string())}
// Unknown fields must not silently imply an unimplemented contribution form.fn validate_fields(table: &Table, allowed: &[&str]) -> Result<(), PluginError> { for key in table.clone().pairs::<String, mlua::Value>() { let (key, _) = key.map_err(lua_error)?; if !allowed.contains(&key.as_str()) { return Err(PluginError::Invalid(format!( "unknown declaration field {key}" ))); } } Ok(())}
/// One evaluated declaration chunk plus its private VM. The initialization/// instruction deadline stays installed: callers walk bounded declarations, not/// arbitrary Lua, and remove the hook once their pass completes.pub(super) struct EvaluatedChunk { pub lua: Lua, pub declaration: Table,}
/// Build the plugin VM and evaluate one declarative chunk. Host capabilities are/// call-scoped, so this VM holds only the chunk, its owned `config`, and the/// standard library.pub(super) fn evaluate(source: &PluginSource) -> Result<EvaluatedChunk, PluginError> { if source.source.len() > 1024 * 1024 { return Err(PluginError::Invalid(format!( "plugin {name}: source exceeds 1 MiB", name = source.name ))); } // Trusted local plugins, not a hostile-code sandbox. The full safe standard // library is present (io, os, load, print) short of the debug library; the // limits below are resource bounds, not isolation promises. let lua = Lua::new(); lua.set_memory_limit(16 * 1024 * 1024).map_err(lua_error)?; lua.globals().set("null", lua.null()).map_err(lua_error)?; // Serialize into this VM before activation. Lua mutations cannot change the // desired JSON retained for reconciliation or any other generation's config. let config = lua.to_value(&source.config).map_err(lua_error)?; lua.globals().set("config", config).map_err(lua_error)?; // Initialization is declarative and cannot await. Bound accidental busy loops. let instructions = AtomicUsize::new(0); lua.set_hook( HookTriggers::new().every_nth_instruction(1000), move |_, _| { if instructions.fetch_add(1000, Ordering::Relaxed) >= 1_000_000 { return Err(mlua::Error::RuntimeError( "initialization instruction budget exceeded".into(), )); } Ok(VmState::Continue) }, ) .map_err(lua_error)?; let declaration: Table = lua .load(&source.source) .set_name(&source.name) .eval() .map_err(lua_error)?; validate_fields( &declaration, &[ "exports", "commands", "hooks", "context_strategies", "compactors", "providers", ], )?; Ok(EvaluatedChunk { lua, declaration })}
/// Dense declaration arrays share one shape: 1..=128 table entries.pub(super) fn dense_array(entries: &Table, field: &str) -> Result<Vec<Table>, PluginError> { let mut indexed = Vec::new(); for entry in entries.clone().pairs::<usize, Table>() { let (index, entry) = entry.map_err(lua_error)?; if index == 0 || index > 128 { return Err(PluginError::Invalid( "declarations must be dense arrays of at most 128 callbacks".into(), )); } indexed.push((index, entry)); } indexed.sort_by_key(|(index, _)| *index); if indexed .iter() .map(|(index, _)| *index) .ne(1..=indexed.len()) { return Err(PluginError::Invalid(format!( "{field} declarations must be dense arrays" ))); } Ok(indexed.into_iter().map(|(_, entry)| entry).collect())}
impl Generation { pub fn load(source: &PluginSource) -> Result<Self, PluginError> { let EvaluatedChunk { lua, declaration } = evaluate(source)?;
let id = Uuid::new_v4(); let mut callbacks = BTreeMap::new(); for (field, kind) in [ ("exports", ContributionKind::Export), ("commands", ContributionKind::Command), ("hooks", ContributionKind::Hook), ("context_strategies", ContributionKind::ContextStrategy), ("compactors", ContributionKind::Compactor), ("providers", ContributionKind::Provider), ] { let Some(entries) = declaration.get::<Option<Table>>(field).map_err(lua_error)? else { continue; }; for entry in dense_array(&entries, field)? { validate_fields( &entry, &["name", "description", "mode", "context_data", "run"], )?; let context_data = entry .get::<Option<String>>("context_data") .map_err(lua_error)?; if context_data.is_some() && !matches!( kind, ContributionKind::ContextStrategy | ContributionKind::Compactor ) { return Err(PluginError::Invalid(format!( "{field} callbacks cannot declare context_data" ))); } let context_data = match context_data.as_deref() { None | Some("json") => ContextDataMode::Json, Some("refs_v1") => ContextDataMode::RefsV1, Some(other) => { return Err(PluginError::Invalid(format!( "unknown context data mode {other}" ))) } }; let name: String = entry.get("name").map_err(lua_error)?; super::names::validate_contribution_name(kind, &name)?; let mode = match entry .get::<Option<String>>("mode") .map_err(lua_error)? .as_deref() { None if kind == ContributionKind::Provider => InvocationMode::Streaming, None | Some("unary") => InvocationMode::Unary, Some("streaming") => InvocationMode::Streaming, Some(other) => { return Err(PluginError::Invalid(format!( "unknown invocation mode {other}" ))) } }; if matches!( kind, ContributionKind::ContextStrategy | ContributionKind::Compactor ) && mode != InvocationMode::Unary { return Err(PluginError::Invalid(format!( "{field} callbacks must be unary" ))); } let registered = Registered { contribution: Contribution { plugin: source.name.clone(), generation: id, kind, name: name.clone(), description: entry .get::<Option<String>>("description") .map_err(lua_error)? .unwrap_or_default(), mode, context_data, }, callback: entry.get("run").map_err(lua_error)?, }; if callbacks.insert((kind, name.clone()), registered).is_some() { return Err(PluginError::Invalid(format!( "duplicate {field} callback {name}" ))); } } } lua.remove_hook(); // Every running coroutine cooperatively yields, even pure Lua busy loops. // mlua releases its internal VM lock on Pending; no registry/session lock // crosses an async capability wait. The caller/retirement select then wins. lua.set_global_hook(HookTriggers::new().every_nth_instruction(10_000), |_, _| { Ok(VmState::Yield) }) .map_err(lua_error)?; Ok(Self { source: source.source.clone(), config: source.config.clone(), retirement: CancellationToken::new(), lua, callbacks, }) }
pub fn contributions(&self) -> impl Iterator<Item = &Contribution> { self.callbacks .values() .map(|registered| ®istered.contribution) }}
#[derive(Clone)]struct CallScope { invocation: Uuid, contribution: Contribution, capabilities: Arc<dyn HostCapabilities>, cancellation: CancellationToken, retirement: CancellationToken, closed: CancellationToken,}
impl CallScope { fn context(&self, lua: &Lua, sender: mpsc::Sender<InvocationEvent>) -> mlua::Result<Table> { let context = lua.create_table()?; let call_scope = self.clone(); context.set("call", lua.create_async_function(move |lua, (operation, input): (String, mlua::Value)| { let scope = call_scope.clone(); async move { let request = CapabilityRequest { invocation: scope.invocation, operation, input: lua.from_value(input)?, }; let output = tokio::select! { biased; () = scope.retirement.cancelled() => return Err(mlua::Error::RuntimeError("retired".into())), () = scope.cancellation.cancelled() => return Err(mlua::Error::RuntimeError("cancelled".into())), () = scope.closed.cancelled() => return Err(mlua::Error::RuntimeError("invocation closed".into())), // Construct the adapter future only after the cancellation fences // have been polled. select! eagerly evaluates branch expressions. output = async { scope.capabilities.call(request, scope.cancellation.child_token()).await } => output.map_err(mlua::Error::RuntimeError)?, }; lua.to_value(&output) } })?)?; let emit_scope = self.clone(); context.set("emit", lua.create_async_function(move |lua, value: mlua::Value| { let scope = emit_scope.clone(); let sender = sender.clone(); async move { if scope.contribution.mode != InvocationMode::Streaming { return Err(mlua::Error::RuntimeError("unary callbacks cannot emit".into())); } let value = lua.from_value(value)?; tokio::select! { biased; () = scope.retirement.cancelled() => Err(mlua::Error::RuntimeError("retired".into())), () = scope.cancellation.cancelled() => Err(mlua::Error::RuntimeError("cancelled".into())), () = scope.closed.cancelled() => Err(mlua::Error::RuntimeError("invocation closed".into())), result = sender.send(InvocationEvent::Update(value)) => result.map_err(mlua::Error::external), } } })?)?; Ok(context) }}
#[tracing::instrument(skip_all, fields(plugin = %contribution.plugin, generation = %contribution.generation))]pub(super) async fn run( generation: Arc<Generation>, contribution: Contribution, input: Value, capabilities: Arc<dyn HostCapabilities>, cancellation: CancellationToken, sender: mpsc::Sender<InvocationEvent>,) { let closed = CancellationToken::new(); let _close_on_exit = closed.clone().drop_guard(); // Retained Lua contexts must not keep working after their originating invocation. // Cancel this adapter token on both success and failure, independently of delivery. let work = cancellation.child_token(); let _cancel_work_on_exit = work.clone().drop_guard(); let scope = CallScope { invocation: Uuid::new_v4(), contribution: contribution.clone(), capabilities, cancellation: work, retirement: generation.retirement.clone(), closed, }; let execution = async { let context = scope .context(&generation.lua, sender.clone()) .map_err(lua_error)?; let input = generation.lua.to_value(&input).map_err(lua_error)?; let registered = generation .callbacks .get(&(contribution.kind, contribution.name.clone())) .ok_or(PluginError::Missing)?; // call_async owns a recyclable coroutine. Its drop resets the thread and // drops pending Rust capability futures on cancellation, not just their result. let output: mlua::Value = registered .callback .call_async((input, context)) .await .map_err(lua_error)?; let output = generation.lua.from_value(output).map_err(lua_error)?; Ok(match contribution.mode { InvocationMode::Unary => InvocationOutcome::Unary(output), InvocationMode::Streaming => InvocationOutcome::Streamed(output), }) }; let result = tokio::select! { biased; () = generation.retirement.cancelled() => Err(PluginError::Retired), () = cancellation.cancelled() => Err(PluginError::Cancelled), result = execution => result, }; scope.closed.cancel(); scope.cancellation.cancel(); tokio::select! { biased; () = generation.retirement.cancelled() => {}, () = cancellation.cancelled() => {}, result = sender.send(InvocationEvent::Finished(result)) => { if result.is_err() { tracing::debug!("plugin invocation consumer dropped"); } } }}