From 67e79c7d9ea0f645522f5745cfeeb46b15a51c64 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 3 Aug 2026 14:02:05 -0400 Subject: [PATCH 01/24] fix(runtime): bound tool and agent resource fanout --- changelog.d/bounded-tool-batching.fixed.md | 1 + n00n-agent/src/agent/run.rs | 9 + n00n-agent/src/agent/streaming.rs | 17 +- n00n-agent/src/agent/tool_dispatch.rs | 112 +++++- n00n-agent/src/mcp/transport.rs | 17 +- n00n-agent/src/tools/admission.rs | 443 +++++++++++++++++++++ n00n-agent/src/tools/mod.rs | 56 ++- n00n-agent/src/tools/registry.rs | 16 + n00n-agent/src/types.rs | 149 +++++++ n00n-interpreter/src/runner.rs | 97 ++++- n00n-lua/src/api/agent.rs | 32 +- n00n-lua/src/api/tool.rs | 21 +- n00n-lua/src/api/util/ctx.rs | 2 +- n00n-lua/src/runtime.rs | 3 +- n00n-providers/src/admission.rs | 139 +++++++ n00n-providers/src/lib.rs | 1 + n00n-ui/src/app/btw.rs | 6 +- plugins/batch/init.lua | 1 + plugins/glob/init.lua | 1 + plugins/grep/init.lua | 1 + plugins/read/init.lua | 1 + plugins/task/init.lua | 1 + plugins/team/init.lua | 1 + plugins/workflow/init.lua | 68 +++- 24 files changed, 1149 insertions(+), 46 deletions(-) create mode 100644 changelog.d/bounded-tool-batching.fixed.md create mode 100644 n00n-agent/src/tools/admission.rs create mode 100644 n00n-providers/src/admission.rs diff --git a/changelog.d/bounded-tool-batching.fixed.md b/changelog.d/bounded-tool-batching.fixed.md new file mode 100644 index 000000000..e9dde11c5 --- /dev/null +++ b/changelog.d/bounded-tool-batching.fixed.md @@ -0,0 +1 @@ +Bounded model-emitted tool calls so excess work queues instead of starting every command at once. Cheap reads use a wider lane, while process-backed tools share an eight-call limit and nested agents share a four-call limit. diff --git a/n00n-agent/src/agent/run.rs b/n00n-agent/src/agent/run.rs index 9a44c2bab..1acdae91c 100644 --- a/n00n-agent/src/agent/run.rs +++ b/n00n-agent/src/agent/run.rs @@ -221,6 +221,7 @@ pub struct Agent<'h> { prompt_slots: Arc, subagent_cancels: Arc>, registry: Arc, + admission_scope: Arc, audience: ToolAudience, workflow: bool, local_tools: LocalTools, @@ -237,6 +238,12 @@ impl<'h> Agent<'h> { pub fn new(params: AgentParams, run: AgentRunParams<'h>) -> Self { let supports_tool_examples = params.model.supports_tool_examples(); let fusion_enabled = params.config.fusion.enabled; + let admission_scope = params + .session_id + .as_ref() + .map_or_else(crate::tools::ToolAdmission::new_scope, |id| { + Arc::::from(id.to_string()) + }); let fusion_state = if fusion_enabled { Some(FusionState::new()) } else { @@ -278,6 +285,7 @@ impl<'h> Agent<'h> { prompt_slots: params.prompt_slots, subagent_cancels: params.subagent_cancels, registry: params.registry, + admission_scope, audience: params.audience, workflow: false, local_tools: LocalTools::default(), @@ -891,6 +899,7 @@ impl<'h> Agent<'h> { opts: self.opts.clone(), subagent_cancels: Arc::clone(&self.subagent_cancels), registry: Arc::clone(&self.registry), + admission_scope: Arc::clone(&self.admission_scope), workflow: self.workflow, audience: self.audience, tool_filter: self.effective_tool_filter(), diff --git a/n00n-agent/src/agent/streaming.rs b/n00n-agent/src/agent/streaming.rs index cd4f31455..57f17bd7b 100644 --- a/n00n-agent/src/agent/streaming.rs +++ b/n00n-agent/src/agent/streaming.rs @@ -8,6 +8,8 @@ use tracing::{info, warn}; use crate::cancel::CancelToken; use crate::{AgentError, AgentEvent, EventSender}; +const PROVIDER_EVENT_QUEUE_CAPACITY: usize = 256; + pub(crate) struct StreamContext<'a> { pub provider: &'a dyn Provider, pub model: &'a Model, @@ -63,11 +65,23 @@ pub(crate) async fn stream_with_retry( let messages = &*messages; let mut retry = RetryState::new(); loop { - let (ptx, prx) = flume::unbounded(); + let (ptx, prx) = flume::bounded(PROVIDER_EVENT_QUEUE_CAPACITY); let forwarder = smol::spawn({ let event_tx = ctx.event_tx.clone(); async move { forward_provider_events(prx, &event_tx).await } }); + let Ok(permit) = ctx + .cancel + .race( + n00n_providers::admission::ProviderAdmission::global() + .acquire(ctx.model.provider.as_ref()), + ) + .await + else { + drop(ptx); + let _ = forwarder.await; + return Err(AgentError::Cancelled); + }; let result = futures_lite::future::race( ctx.provider.stream_message( ctx.model, @@ -84,6 +98,7 @@ pub(crate) async fn stream_with_retry( }, ) .await; + drop(permit); drop(ptx); let emitted_output = forwarder.await; match result { diff --git a/n00n-agent/src/agent/tool_dispatch.rs b/n00n-agent/src/agent/tool_dispatch.rs index 88757f1c7..5ec6a2ae2 100644 --- a/n00n-agent/src/agent/tool_dispatch.rs +++ b/n00n-agent/src/agent/tool_dispatch.rs @@ -13,7 +13,7 @@ use crate::permissions::PermissionCheckContext; use crate::skill_policy::SKILL_POLICY_DENIED_PREFIX; use crate::task_set::TaskSet; use crate::tools::registry::{ToolInvocation, ToolRegistry, ToolSource}; -use crate::tools::{LocalToolFn, ToolContext}; +use crate::tools::{LocalToolFn, ToolAdmissionClass, ToolContext}; use crate::{AgentError, AgentEvent, ToolDoneEvent, ToolOutput, ToolStartEvent}; use n00n_config::ToolKey; @@ -400,6 +400,16 @@ async fn run_authorized( return tool_done_error(id.clone(), Arc::from(name), reason); } if let Some(local) = ctx.local_tools.get(name) { + let class = ToolAdmissionClass::for_tool(name, None); + let _admission = match ctx + .registry + .admission() + .acquire(&ctx.admission_scope, class, &ctx.cancel) + .await + { + Ok(guard) => guard, + Err(error) => return tool_done_error(id, Arc::from(name), error.to_string()), + }; return run_local_tool(local, id, name, input, ctx, emit); } let entry = registry.get(name); @@ -502,6 +512,19 @@ async fn run_authorized( return tool_done_error(id.clone(), Arc::clone(&tool_id), e); } + let _admission = match ctx + .registry + .admission() + .acquire( + &ctx.admission_scope, + entry.tool.admission_class(), + &ctx.cancel, + ) + .await + { + Ok(guard) => guard, + Err(error) => return tool_done_error(id, Arc::clone(&tool_id), error.to_string()), + }; let result = invocation.execute(ctx).await; let elapsed = started.elapsed(); @@ -516,7 +539,8 @@ async fn run_authorized( let output = match result.telemetry { Some(telemetry) => output.with_telemetry(Some(telemetry)), None => output, - }; + } + .bounded(ctx.config.max_output_lines, ctx.config.max_output_bytes); ToolDoneEvent { id, tool: tool_id, @@ -538,7 +562,11 @@ async fn run_authorized( id, tool: tool_id, output: ToolOutput::Plain(crate::TextOutput { - text: message, + text: crate::tools::truncate_output( + &message, + ctx.config.max_output_lines, + ctx.config.max_output_bytes, + ), instructions: None, state: None, telemetry: result.telemetry, @@ -550,8 +578,30 @@ async fn run_authorized( } } } else if let Some(mcp) = mcp.filter(|_| name == TOOL_SEARCH_TOOL_NAME) { + let _admission = match ctx + .registry + .admission() + .acquire(&ctx.admission_scope, ToolAdmissionClass::Cheap, &ctx.cancel) + .await + { + Ok(guard) => guard, + Err(error) => return tool_done_error(id, tool_id, error.to_string()), + }; run_tool_search(mcp, id, input, ctx, emit) } else if mcp.is_some_and(|m| m.has_tool(&mcp_lookup)) { + let _admission = match ctx + .registry + .admission() + .acquire( + &ctx.admission_scope, + ToolAdmissionClass::Standard, + &ctx.cancel, + ) + .await + { + Ok(guard) => guard, + Err(error) => return tool_done_error(id, tool_id, error.to_string()), + }; execute_mcp_tool(ctx, &id, tool_id, &mcp_lookup, input, emit).await } else { let msg = format!("{UNKNOWN_TOOL_PREFIX}: {mcp_lookup}"); @@ -599,8 +649,22 @@ fn run_tool_search( let query = input["query"].as_str().unwrap_or_else(Default::default); emit_raw_start(ctx, emit, &id, &tool_id, query.to_owned(), input); let (output, is_error) = match mcp.search_tools(query) { - Ok(out) => (out, false), - Err(e) => (e, true), + Ok(out) => ( + crate::tools::truncate_output( + &out, + ctx.config.max_output_lines, + ctx.config.max_output_bytes, + ), + false, + ), + Err(e) => ( + crate::tools::truncate_output( + &e, + ctx.config.max_output_lines, + ctx.config.max_output_bytes, + ), + true, + ), }; ToolDoneEvent { id, @@ -623,10 +687,24 @@ fn run_local_tool( let tool_id: Arc = Arc::from(name); emit_raw_start(ctx, emit, &id, &tool_id, name.to_owned(), input); let (output, is_error) = match local(input) { - Ok(output) => (output, false), + Ok(output) => ( + crate::tools::truncate_output( + &output, + ctx.config.max_output_lines, + ctx.config.max_output_bytes, + ), + false, + ), Err(e) => { warn!(tool = %name, error = %e, "local tool failed"); - (e, true) + ( + crate::tools::truncate_output( + &e, + ctx.config.max_output_lines, + ctx.config.max_output_bytes, + ), + true, + ) } }; ToolDoneEvent { @@ -787,8 +865,24 @@ async fn execute_mcp_tool( // definition joins the next request; a denied call must not load anything. mcp.mark_loaded(tool_name); match mcp.call_tool(tool_name, input).await { - Ok(text) => tool_done_plain(id.to_owned(), tool_id, text), - Err(e) => tool_done_error(id.to_owned(), tool_id, e.to_string()), + Ok(text) => tool_done_plain( + id.to_owned(), + tool_id, + crate::tools::truncate_output( + &text, + ctx.config.max_output_lines, + ctx.config.max_output_bytes, + ), + ), + Err(e) => tool_done_error( + id.to_owned(), + tool_id, + crate::tools::truncate_output( + &e.to_string(), + ctx.config.max_output_lines, + ctx.config.max_output_bytes, + ), + ), } } diff --git a/n00n-agent/src/mcp/transport.rs b/n00n-agent/src/mcp/transport.rs index 5f90d81bd..dc153c3f2 100644 --- a/n00n-agent/src/mcp/transport.rs +++ b/n00n-agent/src/mcp/transport.rs @@ -85,6 +85,9 @@ pub async fn list_tools(transport: &dyn McpTransport) -> Result, M } const METHOD_NOT_FOUND: i64 = -32601; +const MAX_MCP_ARGUMENT_BYTES: usize = 64 * 1024; +const MAX_MCP_OUTPUT_BYTES: usize = 256 * 1024; +const MAX_MCP_OUTPUT_LINES: usize = 8192; /// List prompts available on the MCP server. /// @@ -132,6 +135,14 @@ pub async fn call_tool( args: &Value, ) -> Result { let server = &**transport.server_name(); + let argument_bytes = serde_json::to_vec(args).map_err(|error| { + McpError::Config(format!("MCP arguments could not be serialized: {error}")) + })?; + if argument_bytes.len() > MAX_MCP_ARGUMENT_BYTES { + return Err(McpError::Config(format!( + "MCP arguments exceed the {MAX_MCP_ARGUMENT_BYTES} byte limit" + ))); + } let start = Instant::now(); let params = serde_json::json!({ "name": tool_name, @@ -141,7 +152,11 @@ pub async fn call_tool( let call_result: CallToolResult = serde_json::from_value(result).map_err(|e| invalid_response(transport.server_name(), e))?; - let text = call_result.joined_text(); + let text = crate::tools::truncate_output( + &call_result.joined_text(), + MAX_MCP_OUTPUT_LINES, + MAX_MCP_OUTPUT_BYTES, + ); if call_result.is_error { return Err(McpError::RpcError { diff --git a/n00n-agent/src/tools/admission.rs b/n00n-agent/src/tools/admission.rs new file mode 100644 index 000000000..e97bec7da --- /dev/null +++ b/n00n-agent/src/tools/admission.rs @@ -0,0 +1,443 @@ +//! Cancellation-safe admission for work that can fan out or hold process resources. + +use async_lock::{Semaphore, SemaphoreGuardArc}; +use std::collections::HashMap; +use std::fmt; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use crate::cancel::CancelToken; + +pub const DEFAULT_MAX_CONCURRENT_TOOLS: usize = 8; +pub const DEFAULT_MAX_CONCURRENT_CHEAP_TOOLS: usize = 32; +pub const DEFAULT_MAX_CONCURRENT_AGENT_TOOLS: usize = 4; + +const ORCHESTRATOR_TOOLS: &[&str] = &["agent_control", "batch", "task", "team", "workflow"]; +static NEXT_SCOPE: AtomicU64 = AtomicU64::new(1); +const CHEAP_TOOL_KINDS: &[&str] = &["cheap", "read", "metadata", "search"]; +const ORCHESTRATOR_TOOL_KINDS: &[&str] = &["orchestrator", "fanout"]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ToolAdmissionClass { + /// A bounded, low-cost read. Cheap calls have their own wider lane and do + /// not consume a slot from the expensive work budget. + Cheap, + /// Work that may invoke a process, network request, interpreter, or MCP. + Standard, + /// A wrapper that is expected to acquire admission for its children. + /// Wrappers themselves do not hold a slot, which avoids parent-child + /// deadlocks when the process budget is full. + Orchestrator, +} + +impl ToolAdmissionClass { + #[must_use] + pub fn from_workload(value: &str) -> Option { + match value { + "cheap" => Some(Self::Cheap), + "standard" | "expensive" => Some(Self::Standard), + "orchestrator" | "fanout" => Some(Self::Orchestrator), + _ => None, + } + } + + #[must_use] + pub fn for_tool(name: &str, kind: Option<&str>) -> Self { + if ORCHESTRATOR_TOOLS.contains(&name) + || kind.is_some_and(|k| ORCHESTRATOR_TOOL_KINDS.contains(&k)) + { + return Self::Orchestrator; + } + if matches!( + name, + "read" | "glob" | "grep" | "view_image" | "tool_search" + ) || kind.is_some_and(|k| CHEAP_TOOL_KINDS.contains(&k)) + { + return Self::Cheap; + } + Self::Standard + } + + #[must_use] + pub const fn is_orchestrator(self) -> bool { + matches!(self, Self::Orchestrator) + } +} + +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum AdmissionError { + #[error("tool admission cancelled")] + Cancelled, +} + +struct AgentSlot { + semaphore: Arc, + users: AtomicUsize, +} + +struct AdmissionState { + agents: Mutex>>, +} + +/// A registry-owned admission controller. The process and per-agent lanes are +/// independent from Lua's user-facing semaphores, so every dispatch path is +/// covered even when a plugin forgets to opt into a Lua semaphore. +pub struct ToolAdmission { + process: Arc, + cheap: Arc, + agent_limit: usize, + state: AdmissionState, + process_active: AtomicUsize, + cheap_active: AtomicUsize, +} + +impl fmt::Debug for ToolAdmission { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ToolAdmission") + .field("agent_limit", &self.agent_limit) + .field( + "process_active", + &self.process_active.load(Ordering::Relaxed), + ) + .field("cheap_active", &self.cheap_active.load(Ordering::Relaxed)) + .finish_non_exhaustive() + } +} + +impl Default for ToolAdmission { + fn default() -> Self { + Self::new() + } +} + +impl ToolAdmission { + #[must_use] + pub fn new() -> Self { + Self::with_limits( + DEFAULT_MAX_CONCURRENT_TOOLS, + DEFAULT_MAX_CONCURRENT_AGENT_TOOLS, + DEFAULT_MAX_CONCURRENT_CHEAP_TOOLS, + ) + } + + #[must_use] + pub fn new_scope() -> Arc { + Arc::from(format!( + "agent-{}", + NEXT_SCOPE.fetch_add(1, Ordering::Relaxed) + )) + } + + #[must_use] + pub fn with_limits(process_limit: usize, agent_limit: usize, cheap_limit: usize) -> Self { + Self { + process: Arc::new(Semaphore::new(process_limit.max(1))), + cheap: Arc::new(Semaphore::new(cheap_limit.max(1))), + agent_limit: agent_limit.max(1), + state: AdmissionState { + agents: Mutex::new(HashMap::new()), + }, + process_active: AtomicUsize::new(0), + cheap_active: AtomicUsize::new(0), + } + } + + /// Wait for admission while allowing cancellation to remove the waiter. + /// All guards release on drop, including when the tool future panics. + /// + /// # Errors + /// + /// Returns [`AdmissionError::Cancelled`] when cancellation wins the race + /// while waiting for a permit. + pub async fn acquire( + &self, + scope: &str, + class: ToolAdmissionClass, + cancel: &CancelToken, + ) -> Result, AdmissionError> { + if class.is_orchestrator() { + return Ok(ToolAdmissionGuard::empty()); + } + + if matches!(class, ToolAdmissionClass::Cheap) { + let permit = cancel + .race(self.cheap.acquire_arc()) + .await + .map_err(|_| AdmissionError::Cancelled)?; + self.cheap_active.fetch_add(1, Ordering::Relaxed); + return Ok(ToolAdmissionGuard { + _process: None, + _cheap: Some(ActivePermit { + guard: permit, + active: &self.cheap_active, + }), + agent: None, + state: None, + scope: None, + }); + } + + let process = ActivePermit { + guard: cancel + .race(self.process.acquire_arc()) + .await + .map_err(|_| AdmissionError::Cancelled)?, + active: &self.process_active, + }; + self.process_active.fetch_add(1, Ordering::Relaxed); + + let scope = scope.to_owned(); + let agent = { + let mut agents = self + .state + .agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let slot = Arc::clone(agents.entry(scope.clone()).or_insert_with(|| { + Arc::new(AgentSlot { + semaphore: Arc::new(Semaphore::new(self.agent_limit)), + users: AtomicUsize::new(0), + }) + })); + slot.users.fetch_add(1, Ordering::Relaxed); + slot + }; + + let Ok(agent_guard) = cancel.race(agent.semaphore.acquire_arc()).await else { + self.release_agent(&scope, &agent); + return Err(AdmissionError::Cancelled); + }; + + Ok(ToolAdmissionGuard { + _process: Some(process), + _cheap: None, + agent: Some(agent_guard), + state: Some(&self.state), + scope: Some(scope), + }) + } + + fn release_agent(&self, scope: &str, slot: &Arc) { + let mut agents = self + .state + .agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if !agents + .get(scope) + .is_some_and(|current| Arc::ptr_eq(current, slot)) + { + return; + } + if slot.users.fetch_sub(1, Ordering::Relaxed) == 1 { + agents.remove(scope); + } + } + + #[must_use] + pub fn process_active(&self) -> usize { + self.process_active.load(Ordering::Relaxed) + } + + #[must_use] + pub fn cheap_active(&self) -> usize { + self.cheap_active.load(Ordering::Relaxed) + } +} + +struct ActivePermit<'a> { + guard: SemaphoreGuardArc, + active: &'a AtomicUsize, +} + +impl Drop for ActivePermit<'_> { + fn drop(&mut self) { + let _ = &self.guard; + self.active.fetch_sub(1, Ordering::Relaxed); + } +} + +pub struct ToolAdmissionGuard<'a> { + _process: Option>, + _cheap: Option>, + agent: Option, + state: Option<&'a AdmissionState>, + scope: Option, +} + +impl ToolAdmissionGuard<'_> { + fn empty() -> Self { + Self { + _process: None, + _cheap: None, + agent: None, + state: None, + scope: None, + } + } +} + +impl Drop for ToolAdmissionGuard<'_> { + fn drop(&mut self) { + self.agent.take(); + let (Some(state), Some(scope)) = (self.state, self.scope.take()) else { + return; + }; + let mut agents = state + .agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let Some(slot) = agents.get(&scope).cloned() else { + return; + }; + if slot.users.fetch_sub(1, Ordering::Relaxed) == 1 + && agents + .get(&scope) + .is_some_and(|current| Arc::ptr_eq(current, &slot)) + { + agents.remove(&scope); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cancel::CancelToken; + + #[test] + fn classifies_wrappers_without_consuming_expensive_lane() { + assert_eq!( + ToolAdmissionClass::for_tool("batch", Some("execute")), + ToolAdmissionClass::Orchestrator + ); + assert_eq!( + ToolAdmissionClass::for_tool("read", Some("filesystem")), + ToolAdmissionClass::Cheap + ); + assert_eq!( + ToolAdmissionClass::for_tool("webfetch", Some("network")), + ToolAdmissionClass::Standard + ); + } + + #[test] + fn cancelled_wait_does_not_leak_process_or_agent_slot() { + smol::block_on(async { + let admission = Arc::new(ToolAdmission::with_limits(1, 1, 1)); + let cancel = CancelToken::none(); + let first = admission + .acquire("agent", ToolAdmissionClass::Standard, &cancel) + .await + .expect("first permit"); + let (trigger, waiter_cancel) = CancelToken::new(); + let (started_tx, started_rx) = flume::bounded(1); + let admission_for_waiter = Arc::clone(&admission); + let waiter = smol::spawn(async move { + started_tx + .send_async(()) + .await + .expect("test receiver remains available"); + assert!(matches!( + admission_for_waiter + .acquire("agent", ToolAdmissionClass::Standard, &waiter_cancel) + .await, + Err(AdmissionError::Cancelled) + )); + }); + started_rx + .recv_async() + .await + .expect("waiter reached admission"); + trigger.cancel(); + waiter.await; + drop(first); + assert_eq!(admission.process_active(), 0); + }); + } + + #[test] + fn cancelled_agent_wait_releases_process_permit() { + smol::block_on(async { + let admission = Arc::new(ToolAdmission::with_limits(2, 1, 1)); + let first = admission + .acquire("agent", ToolAdmissionClass::Standard, &CancelToken::none()) + .await + .expect("first permit"); + let (trigger, waiter_cancel) = CancelToken::new(); + let (started_tx, started_rx) = flume::bounded(1); + let admission_for_waiter = Arc::clone(&admission); + let waiter = smol::spawn(async move { + started_tx + .send_async(()) + .await + .expect("test receiver remains available"); + assert!(matches!( + admission_for_waiter + .acquire("agent", ToolAdmissionClass::Standard, &waiter_cancel) + .await, + Err(AdmissionError::Cancelled) + )); + }); + started_rx + .recv_async() + .await + .expect("waiter reached admission"); + trigger.cancel(); + waiter.await; + assert_eq!(admission.process_active(), 1); + drop(first); + assert_eq!(admission.process_active(), 0); + }); + } + + #[test] + fn permit_releases_when_work_returns_error() { + smol::block_on(async { + let admission = ToolAdmission::with_limits(1, 1, 1); + let result: Result<(), AdmissionError> = async { + let _permit = admission + .acquire("agent", ToolAdmissionClass::Standard, &CancelToken::none()) + .await?; + Err(AdmissionError::Cancelled) + } + .await; + assert_eq!(result, Err(AdmissionError::Cancelled)); + assert_eq!(admission.process_active(), 0); + }); + } + + #[test] + fn permit_releases_when_work_panics() { + let admission = Arc::new(ToolAdmission::with_limits(1, 1, 1)); + let admission_for_work = Arc::clone(&admission); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || { + smol::block_on(async move { + let _permit = admission_for_work + .acquire("agent", ToolAdmissionClass::Standard, &CancelToken::none()) + .await + .expect("permit"); + panic!("simulated tool panic"); + }); + })); + assert!(result.is_err()); + assert_eq!(admission.process_active(), 0); + } + + #[test] + fn orchestrator_is_a_noop_guard() { + smol::block_on(async { + let admission = ToolAdmission::with_limits(1, 1, 1); + let permit = admission + .acquire( + "agent", + ToolAdmissionClass::Orchestrator, + &CancelToken::none(), + ) + .await + .expect("orchestrator bypass"); + assert_eq!(admission.process_active(), 0); + drop(permit); + }); + } +} diff --git a/n00n-agent/src/tools/mod.rs b/n00n-agent/src/tools/mod.rs index 1f17a3909..038a7fc8f 100644 --- a/n00n-agent/src/tools/mod.rs +++ b/n00n-agent/src/tools/mod.rs @@ -5,6 +5,7 @@ //! quotes, camelCase keys, extra wrappers). Plan mode rejects writes to //! anything but the plan file before they reach the tool. +pub mod admission; mod file_tracker; pub mod grep; pub mod interpreter_bridge; @@ -12,6 +13,7 @@ pub mod registry; pub mod schema; pub mod tool_search; +pub use admission::{AdmissionError, ToolAdmission, ToolAdmissionClass, ToolAdmissionGuard}; pub use file_tracker::FileReadTracker; pub use registry::{ ActiveTools, BoxFuture, ExecFuture, HeaderFuture, HeaderResult, ParseError, PermissionScopes, @@ -322,6 +324,9 @@ pub struct ToolContext { pub opts: RequestOptions, pub subagent_cancels: Arc>, pub registry: Arc, + /// Stable identity used by registry-scoped per-agent admission. Child + /// sessions receive their own scope, while nested calls in one agent share it. + pub admission_scope: Arc, pub tool_filter: ToolFilter, pub workflow: bool, pub audience: ToolAudience, @@ -461,22 +466,34 @@ pub(crate) fn truncate_bytes(line: &str, max_bytes: usize) -> String { #[must_use] pub fn truncate_output(text: &str, max_lines: usize, max_bytes: usize) -> String { const TRUNCATED_MARKER: &str = "[truncated]"; + if max_bytes == 0 || max_lines == 0 { + return String::new(); + } let mut lines = text.lines(); let mut result = String::new(); let mut truncated = false; for _ in 0..max_lines { let Some(line) = lines.next() else { break }; - if !result.is_empty() { - result.push('\n'); - } - result.push_str(line); - if result.len() > max_bytes { - let boundary = result.floor_char_boundary(max_bytes); - result.truncate(boundary); + let separator = usize::from(!result.is_empty()); + if result + .len() + .saturating_add(separator) + .saturating_add(line.len()) + > max_bytes + { + let remaining = max_bytes.saturating_sub(result.len().saturating_add(separator)); + if separator != 0 && remaining > 0 { + result.push('\n'); + } + result.push_str(&line[..line.floor_char_boundary(remaining)]); truncated = true; break; } + if separator != 0 { + result.push('\n'); + } + result.push_str(line); } if !truncated && lines.next().is_some() { @@ -484,8 +501,19 @@ pub fn truncate_output(text: &str, max_lines: usize, max_bytes: usize) -> String } if truncated { - result.push('\n'); - result.push_str(TRUNCATED_MARKER); + let suffix = if result.is_empty() { + TRUNCATED_MARKER.to_owned() + } else { + format!("\n{TRUNCATED_MARKER}") + }; + if suffix.len() >= max_bytes { + result.clear(); + suffix[..suffix.floor_char_boundary(max_bytes)].clone_into(&mut result); + } else { + let content_limit = max_bytes - suffix.len(); + result.truncate(result.floor_char_boundary(content_limit)); + result.push_str(&suffix); + } } result } @@ -589,6 +617,7 @@ pub fn interpreter_ctx( opts: RequestOptions::default(), subagent_cancels: Arc::new(CancelMap::new()), registry, + admission_scope: crate::tools::ToolAdmission::new_scope(), tool_filter: ToolFilter::All, workflow: false, audience: ToolAudience::MAIN, @@ -794,6 +823,13 @@ mod tests { assert_eq!(result, expected); } + #[test_case("abcdefghij\nklmno", 2, 15, "abc\n[truncated]" ; "short_byte_limit_keeps_marker")] + #[test_case("one\ntwo", 0, 10, "" ; "zero_line_limit")] + #[test_case("one", 10, 0, "" ; "zero_byte_limit")] + fn truncate_output_edge_cases(input: &str, max_lines: usize, max_bytes: usize, expected: &str) { + assert_eq!(truncate_output(input, max_lines, max_bytes), expected); + } + #[test] fn truncate_output_respects_line_and_byte_limits() { const MAX_LINES: usize = 2000; @@ -805,10 +841,12 @@ mod tests { .join("\n"); let result = truncate_output(&many_lines, MAX_LINES, MAX_BYTES); assert!(result.ends_with("[truncated]")); + assert!(result.len() <= MAX_BYTES); let many_bytes = "x".repeat(MAX_BYTES + 1000); let result = truncate_output(&many_bytes, MAX_LINES, MAX_BYTES); assert!(result.ends_with("[truncated]")); + assert!(result.len() <= MAX_BYTES); } #[test] diff --git a/n00n-agent/src/tools/registry.rs b/n00n-agent/src/tools/registry.rs index 7ce8863ad..e5bf1e4e0 100644 --- a/n00n-agent/src/tools/registry.rs +++ b/n00n-agent/src/tools/registry.rs @@ -16,6 +16,7 @@ use serde_json::{Value, json}; use crate::template::Vars; use crate::{BufferSnapshot, ToolOutput}; +use super::admission::{ToolAdmission, ToolAdmissionClass}; use super::schema::sanitize_tool_input_schema; use super::{DescriptionContext, ToolContext}; @@ -241,6 +242,9 @@ pub trait Tool: Send + Sync + 'static { fn tool_kind(&self) -> Option<&str> { None } + fn admission_class(&self) -> ToolAdmissionClass { + ToolAdmissionClass::for_tool(self.name(), self.tool_kind()) + } fn defer_loading(&self) -> bool { false } @@ -336,6 +340,7 @@ impl<'a> IntoIterator for &'a ToolsSnapshot { /// Lock-free reads via `ArcSwap`, writes swap in a new snapshot atomically. pub struct ToolRegistry { tools: ArcSwap, + admission: Arc, } impl Default for ToolRegistry { @@ -353,11 +358,22 @@ pub enum RegistryError { impl ToolRegistry { #[must_use] pub fn new() -> Self { + Self::with_admission(Arc::new(ToolAdmission::new())) + } + + #[must_use] + pub fn with_admission(admission: Arc) -> Self { Self { tools: ArcSwap::from_pointee(ToolsSnapshot::empty()), + admission, } } + #[must_use] + pub fn admission(&self) -> &ToolAdmission { + &self.admission + } + /// The process-wide registry. Every tool in it comes from a Lua plugin /// or an MCP server; Rust itself registers nothing. #[must_use] diff --git a/n00n-agent/src/types.rs b/n00n-agent/src/types.rs index 8425c0b1a..cee5ef46d 100644 --- a/n00n-agent/src/types.rs +++ b/n00n-agent/src/types.rs @@ -470,6 +470,22 @@ fn lines_remaining_after(total: usize, start_line: usize, shown: usize) -> usize total.saturating_sub(end) } +fn bounded_text_output(mut output: TextOutput, max_lines: usize, max_bytes: usize) -> TextOutput { + output.text = crate::tools::truncate_output(&output.text, max_lines, max_bytes); + let state_limit = max_bytes.saturating_sub(output.text.len()); + let state_too_large = match output.state.as_ref() { + Some(state) => match serde_json::to_vec(state) { + Ok(encoded) => encoded.len() > state_limit, + Err(_) => true, + }, + None => false, + }; + if state_too_large { + output.state = None; + } + output +} + impl ToolOutput { fn grep_summary(entries: &[GrepFileEntry]) -> String { let matches: usize = entries.iter().map(GrepFileEntry::match_count).sum(); @@ -614,6 +630,98 @@ impl ToolOutput { } } + #[must_use] + pub fn bounded(self, max_lines: usize, max_bytes: usize) -> Self { + let bounded = match self { + Self::Plain(output) => Self::Plain(bounded_text_output(output, max_lines, max_bytes)), + Self::Markdown(output) => { + Self::Markdown(bounded_text_output(output, max_lines, max_bytes)) + } + Self::ReadDir(output) => { + Self::ReadDir(bounded_text_output(output, max_lines, max_bytes)) + } + Self::ReadCode { + path, + start_line, + mut lines, + total_lines, + instructions, + } => { + lines.truncate(max_lines); + let max_line_bytes = max_bytes.clamp(1, 4096); + for line in &mut lines { + *line = crate::tools::truncate_output(line, 1, max_line_bytes); + } + Self::ReadCode { + path, + start_line, + lines, + total_lines, + instructions, + } + } + Self::Diff { + path, + before, + after, + summary, + telemetry, + } => Self::Diff { + path, + before: crate::tools::truncate_output(&before, max_lines, max_bytes / 2), + after: crate::tools::truncate_output(&after, max_lines, max_bytes / 2), + summary: crate::tools::truncate_output(&summary, max_lines, max_bytes), + telemetry, + }, + Self::TodoList(mut items) => { + items.truncate(max_lines); + for item in &mut items { + item.content = crate::tools::truncate_output(&item.content, 1, max_bytes); + } + Self::TodoList(items) + } + Self::WriteCode { + path, + byte_count, + mut lines, + } => { + lines.truncate(max_lines); + for line in &mut lines { + *line = crate::tools::truncate_output(line, 1, max_bytes); + } + Self::WriteCode { + path, + byte_count, + lines, + } + } + Self::GrepResult { entries } => { + let text = Self::GrepResult { entries }.as_display_text(); + Self::Plain(crate::tools::truncate_output(&text, max_lines, max_bytes).into()) + } + Self::Batch { text } => Self::Batch { + text: crate::tools::truncate_output(&text, max_lines, max_bytes), + }, + Self::Instructions { blocks } => Self::Instructions { blocks }, + Self::Image { + source, + text, + telemetry, + } => Self::Image { + source, + text: crate::tools::truncate_output(&text, max_lines, max_bytes), + telemetry, + }, + }; + let text = bounded.as_text(); + let limited = crate::tools::truncate_output(&text, max_lines, max_bytes); + if limited == text { + bounded + } else { + Self::Plain(limited.into()) + } + } + #[must_use] pub fn with_telemetry(mut self, telemetry: Option) -> Self { let Some(telemetry) = telemetry else { @@ -1439,6 +1547,47 @@ mod tests { assert_eq!(summary.as_text(), expected_text); } + #[test] + fn bounded_output_limits_structured_display() { + let output = ToolOutput::ReadCode { + path: "a.rs".into(), + start_line: 1, + lines: vec!["x".repeat(200); 100], + total_lines: 100, + instructions: None, + }; + let bounded = output.bounded(10, 100); + assert!(bounded.as_text().len() <= 100); + assert!(bounded.as_text().lines().count() <= 11); + } + + #[test] + fn bounded_output_drops_oversized_state() { + let output = ToolOutput::Plain(TextOutput { + text: "body".into(), + instructions: None, + state: Some(serde_json::json!({ "payload": "x".repeat(100) })), + telemetry: None, + }); + let bounded = output.bounded(10, 16); + assert!(bounded.state().is_none()); + } + + #[test] + fn bounded_output_preserves_instructions_when_room_allows() { + let output = ToolOutput::Plain(TextOutput { + text: "body".into(), + instructions: Some(vec![InstructionBlock { + path: "AGENTS.md".into(), + content: "do stuff".into(), + }]), + state: None, + telemetry: None, + }); + let bounded = output.bounded(10, 100); + assert!(bounded.as_text().contains("Instructions from: AGENTS.md")); + } + #[test] fn summary_preserves_instructions() { let instructions = vec![InstructionBlock { diff --git a/n00n-interpreter/src/runner.rs b/n00n-interpreter/src/runner.rs index abb9e7cbb..cf60eaedf 100644 --- a/n00n-interpreter/src/runner.rs +++ b/n00n-interpreter/src/runner.rs @@ -20,6 +20,8 @@ use crate::convert::{json_to_monty, monty_to_json}; use crate::error::InterpreterError; const DEFAULT_MAX_RECURSION: usize = 100; +const MAX_STDOUT_BYTES: usize = 16 * 1024; +const MAX_RESULT_BYTES: usize = 16 * 1024; const SCRIPT_NAME: &str = "agent.py"; pub type ToolFn = Box, Vec<(String, Value)>) -> Result>; @@ -43,16 +45,34 @@ pub struct InterpreterResult { struct StreamingWriter<'a> { buffer: String, flushed_pos: usize, + truncated: bool, on_line: &'a mut dyn FnMut(&str), } impl PrintWriterCallback for StreamingWriter<'_> { fn stdout_write(&mut self, output: Cow<'_, str>) -> Result<(), MontyException> { - self.buffer.push_str(&output); + if self.truncated { + return Ok(()); + } + let remaining = MAX_STDOUT_BYTES.saturating_sub(self.buffer.len()); + let boundary = output.floor_char_boundary(remaining); + self.buffer.push_str(&output[..boundary]); + if boundary < output.len() { + self.truncated = true; + } Ok(()) } fn stdout_push(&mut self, ch: char) -> Result<(), MontyException> { + if self.truncated { + return Ok(()); + } + let mut encoded = [0; 4]; + let text = ch.encode_utf8(&mut encoded); + if self.buffer.len().saturating_add(text.len()) > MAX_STDOUT_BYTES { + self.truncated = true; + return Ok(()); + } self.buffer.push(ch); if ch == '\n' { (self.on_line)(&self.buffer[self.flushed_pos..]); @@ -62,6 +82,22 @@ impl PrintWriterCallback for StreamingWriter<'_> { } } +impl StreamingWriter<'_> { + fn finish(&mut self) { + if self.flushed_pos < self.buffer.len() { + (self.on_line)(&self.buffer[self.flushed_pos..]); + self.flushed_pos = self.buffer.len(); + } + if self.truncated { + const MARKER: &str = "[truncated]"; + let keep = MAX_STDOUT_BYTES.saturating_sub(MARKER.len()); + self.buffer.truncate(self.buffer.floor_char_boundary(keep)); + self.buffer.push_str(MARKER); + (self.on_line)(MARKER); + } + } +} + /// Runs Python code with the given tools and resource limits. /// /// # Errors @@ -76,10 +112,23 @@ pub fn run( resolver: Option<&AsyncResolver>, limits: ResourceLimits, ) -> Result { - let mut stdout = String::new(); - let mut print_writer = PrintWriter::CollectString(&mut stdout); - let output = run_inner(code, tools, resolver, limits, &mut print_writer)?; - Ok(InterpreterResult { output, stdout }) + let mut on_line = |_line: &str| {}; + let mut writer = StreamingWriter { + buffer: String::new(), + flushed_pos: 0, + truncated: false, + on_line: &mut on_line, + }; + let result = { + let mut print_writer = PrintWriter::Callback(&mut writer); + run_inner(code, tools, resolver, limits, &mut print_writer) + }; + writer.finish(); + let output = result?; + Ok(InterpreterResult { + output, + stdout: writer.buffer, + }) } /// Runs Python code with streaming stdout output via a callback. @@ -100,10 +149,15 @@ pub fn run_streaming( let mut writer = StreamingWriter { buffer: String::new(), flushed_pos: 0, + truncated: false, on_line: on_output, }; - let mut print_writer = PrintWriter::Callback(&mut writer); - let output = run_inner(code, tools, resolver, limits, &mut print_writer)?; + let result = { + let mut print_writer = PrintWriter::Callback(&mut writer); + run_inner(code, tools, resolver, limits, &mut print_writer) + }; + writer.finish(); + let output = result?; let stdout = writer.buffer; Ok(InterpreterResult { output, stdout }) } @@ -130,9 +184,19 @@ fn run_inner( loop { match progress { RunProgress::Complete(obj) => { - let output = match &obj { - MontyObject::None => None, - _ => Some(monty_to_json(&obj)), + let output = if matches!(obj, MontyObject::None) { + None + } else { + let value = monty_to_json(&obj); + let size = serde_json::to_vec(&value) + .map_err(|e| InterpreterError::Runtime(e.to_string()))? + .len(); + if size > MAX_RESULT_BYTES { + return Err(InterpreterError::Runtime(format!( + "result exceeds the {MAX_RESULT_BYTES} byte limit" + ))); + } + Some(value) }; return Ok(output); } @@ -390,6 +454,19 @@ mod tests { assert!(called); } + #[test] + fn stdout_is_bounded_and_marked() { + let result = run("print('x' * 20000)", &empty_tools(), None, default_limits()).unwrap(); + assert!(result.stdout.len() <= MAX_STDOUT_BYTES); + assert!(result.stdout.ends_with("[truncated]")); + } + + #[test] + fn oversized_result_is_rejected() { + let err = run("'x' * 20000", &empty_tools(), None, default_limits()).unwrap_err(); + assert!(err.to_string().contains("result exceeds")); + } + #[test] fn async_gather_resolves_concurrently() { let code = r" diff --git a/n00n-lua/src/api/agent.rs b/n00n-lua/src/api/agent.rs index 85c9fe7cd..478088f4c 100644 --- a/n00n-lua/src/api/agent.rs +++ b/n00n-lua/src/api/agent.rs @@ -59,6 +59,8 @@ const SAFE_ACTIVITY_DESCRIPTION_TOOLS: &[&str] = &[ "write", ]; const PROGRESS_TIMEOUT_MS: u64 = 500; +const LIVE_EVENT_QUEUE_CAPACITY: usize = 256; +const SUBAGENT_EVENT_QUEUE_CAPACITY: usize = 1024; const STEERING_QUEUE_CAPACITY: usize = 32; const TOOL_EXCLUSIONS_META_FIELD: &str = "__n00n_tool_exclusions"; @@ -116,6 +118,16 @@ fn parse_session_mode( type Pair = (Option, Option); #[allow(clippy::needless_pass_by_value)] +fn normalize_tool_definitions(tools: JsonValue) -> Result { + if tools.is_array() { + return Ok(tools); + } + if tools.as_object().is_some_and(serde_json::Map::is_empty) { + return Ok(JsonValue::Array(Vec::new())); + } + Err("tools must be an array".to_owned()) +} + fn explicit_tool_filter(tools: &JsonValue) -> Result { let definitions = tools .as_array() @@ -563,7 +575,7 @@ async fn call_tool( on_buf = o.get::>("on_live_buf")?; on_ann = o.get::>("on_annotation")?; if on_buf.is_some() || on_ann.is_some() { - let (tx, r) = flume::unbounded(); + let (tx, r) = flume::bounded(LIVE_EVENT_QUEUE_CAPACITY); tctx.live_sink = Some(tx); rx = Some(r); } @@ -694,15 +706,12 @@ async fn session( // A standalone task shows its model via SubagentInfo on the header; // a dispatching caller (batch) gets the same thing as a live annotation. if let Some(sink) = &agent_ctx.live_sink { - let _ = sink.send(ToolLive::Annotation(model.spec())); + let _ = sink.try_send(ToolLive::Annotation(model.spec())); } let explicit_tools = tools_val.is_some(); let (mut tools_json, mut tool_filter) = if let Some(val) = tools_val { - let tools = lua_to_json(&lua, &val)?; - if !tools.is_array() { - return Err(mlua::Error::runtime("tools must be an array")); - } + let tools = try_pair!(normalize_tool_definitions(lua_to_json(&lua, &val)?)); (tools, ToolFilter::All) } else { let vars = n00n_agent::template::Vars::new(); @@ -782,7 +791,7 @@ async fn session( let child_id = session_id.to_string(); let parent_tool_use_id = child_id.clone(); let start = Instant::now(); - let (sub_tx, sub_rx) = flume::unbounded::(); + let (sub_tx, sub_rx) = flume::bounded::(SUBAGENT_EVENT_QUEUE_CAPACITY); let sub_event_tx = EventSender::new(sub_tx, agent_ctx.event_tx.run_id()); let parent_tx = agent_ctx.event_tx.clone(); let (answer_tx, answer_rx) = flume::unbounded::(); @@ -1742,6 +1751,15 @@ mod tests { assert_eq!(definitions, json!([{"name": "read"}])); } + #[test] + fn empty_explicit_tools_are_normalized_to_an_array() { + assert_eq!( + normalize_tool_definitions(json!({})).unwrap(), + JsonValue::Array(Vec::new()) + ); + assert!(normalize_tool_definitions(json!({"name": "read"})).is_err()); + } + #[test] fn explicit_session_filter_uses_final_definition_names() { let tools = json!([ diff --git a/n00n-lua/src/api/tool.rs b/n00n-lua/src/api/tool.rs index 93a4e6924..2735b20c6 100644 --- a/n00n-lua/src/api/tool.rs +++ b/n00n-lua/src/api/tool.rs @@ -19,8 +19,8 @@ use n00n_agent::tools::registry::{RegisteredTool, ToolRegistry}; use n00n_agent::tools::schema::{ParamSchema, to_json_schema, try_from_json, validate}; use n00n_agent::tools::{ BoxFuture, Deadline, DescriptionContext, ExecFuture, HeaderFuture, HeaderResult, ParseError, - PermissionScopes, ToolAudience, ToolContext, ToolExecResult, ToolFilter, ToolInvocation, - is_tool_enabled, timeout_annotation, + PermissionScopes, ToolAdmissionClass, ToolAudience, ToolContext, ToolExecResult, ToolFilter, + ToolInvocation, is_tool_enabled, timeout_annotation, }; use n00n_agent::{ AgentEvent, BufferSnapshot, ImageMediaType, ImageSource, InstructionBlock, SharedBuf, @@ -126,6 +126,7 @@ pub(crate) struct PendingTool { pub(crate) schema: &'static ParamSchema, pub(crate) audience: ToolAudience, pub(crate) kind: Option>, + pub(crate) workload: Option, pub(crate) handler_key: RegistryKey, pub(crate) header_key: Option, pub(crate) restore_key: Option, @@ -148,6 +149,7 @@ pub(crate) struct LuaTool { pub(crate) schema: &'static ParamSchema, pub(crate) audience: ToolAudience, pub(crate) kind: Option>, + pub(crate) workload: Option, pub(crate) tx: Sender, pub(crate) plugin: Arc, pub(crate) has_header_fn: bool, @@ -214,6 +216,11 @@ impl Tool for LuaTool { self.kind.as_deref() } + fn admission_class(&self) -> ToolAdmissionClass { + self.workload + .unwrap_or_else(|| ToolAdmissionClass::for_tool(&self.name, self.kind.as_deref())) + } + fn examples(&self) -> Option { self.examples.clone() } @@ -1172,6 +1179,14 @@ fn register_tool_from_lua(lua: &Lua, spec: &Table, pending: PendingTools) -> Lua .get::("kind") .ok() .map(|s| Arc::from(s.as_str())); + let workload = match spec.get::>("workload")? { + Some(value) => Some(ToolAdmissionClass::from_workload(&value).ok_or_else(|| { + mlua::Error::runtime( + "register_tool: workload must be cheap, standard, expensive, or orchestrator", + ) + })?), + None => None, + }; let audience = parse_audience(audiences)?; let timeout = parse_timeout(spec)?; let start_annotation = parse_start_annotation(spec, &schema_val)?; @@ -1215,6 +1230,7 @@ fn register_tool_from_lua(lua: &Lua, spec: &Table, pending: PendingTools) -> Lua schema: param_schema, audience, kind, + workload, handler_key, header_key, restore_key, @@ -1721,6 +1737,7 @@ mod tests { schema, audience: ToolAudience::default(), kind: None, + workload: None, tx, plugin: Arc::from("test"), has_header_fn: false, diff --git a/n00n-lua/src/api/util/ctx.rs b/n00n-lua/src/api/util/ctx.rs index 7f79feb29..74a6396c7 100644 --- a/n00n-lua/src/api/util/ctx.rs +++ b/n00n-lua/src/api/util/ctx.rs @@ -33,7 +33,7 @@ fn send_live_buf(lua: &mlua::Lua, buf: &mlua::AnyUserData) -> mlua::Result<()> { }); } if let Some(sink) = sink { - let _ = sink.send(ToolLive::Buf(shared)); + let _ = sink.try_send(ToolLive::Buf(shared)); } Ok(()) } diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index 7be7cafa7..53edec22a 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -1469,6 +1469,7 @@ impl LuaRuntime { schema: t.schema, audience: t.audience, kind: t.kind.clone(), + workload: t.workload, tx: self.tx.clone(), plugin: Arc::clone(&name), has_header_fn: t.header_key.is_some(), @@ -2125,7 +2126,7 @@ async fn run_tool_call( }); } if let Some(sink) = sink { - let _ = sink.send(ToolLive::Buf(buf)); + let _ = sink.try_send(ToolLive::Buf(buf)); } } dispatch_async(&lua, Arc::clone(&handle), &plugin, &tool, finish_rx).await diff --git a/n00n-providers/src/admission.rs b/n00n-providers/src/admission.rs new file mode 100644 index 000000000..3f3778421 --- /dev/null +++ b/n00n-providers/src/admission.rs @@ -0,0 +1,139 @@ +//! Process and provider scoped admission for outbound model streams. + +use async_lock::{Semaphore, SemaphoreGuardArc}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, LazyLock, Mutex}; + +pub const DEFAULT_MAX_CONCURRENT_STREAMS: usize = 8; +pub const DEFAULT_MAX_CONCURRENT_STREAMS_PER_PROVIDER: usize = 4; +const MAX_PROVIDER_KEYS: usize = 64; + +pub struct ProviderAdmission { + process: Arc, + provider_limit: usize, + providers: Mutex>>, + overflow: Arc, + active: AtomicUsize, +} + +impl Default for ProviderAdmission { + fn default() -> Self { + Self::new() + } +} + +impl ProviderAdmission { + #[must_use] + pub fn new() -> Self { + Self::with_limits( + DEFAULT_MAX_CONCURRENT_STREAMS, + DEFAULT_MAX_CONCURRENT_STREAMS_PER_PROVIDER, + ) + } + + #[must_use] + pub fn with_limits(process_limit: usize, provider_limit: usize) -> Self { + Self { + process: Arc::new(Semaphore::new(process_limit.max(1))), + provider_limit: provider_limit.max(1), + providers: Mutex::new(HashMap::new()), + overflow: Arc::new(Semaphore::new(provider_limit.max(1))), + active: AtomicUsize::new(0), + } + } + + #[must_use] + pub fn global() -> &'static Self { + static GLOBAL: LazyLock = LazyLock::new(ProviderAdmission::new); + &GLOBAL + } + + pub async fn acquire(&self, provider: &str) -> ProviderAdmissionGuard<'_> { + let process = self.process.acquire_arc().await; + let provider_sem = { + let mut providers = self + .providers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if providers.len() < MAX_PROVIDER_KEYS || providers.contains_key(provider) { + Arc::clone( + providers + .entry(provider.to_owned()) + .or_insert_with(|| Arc::new(Semaphore::new(self.provider_limit))), + ) + } else { + Arc::clone(&self.overflow) + } + }; + let provider = provider_sem.acquire_arc().await; + self.active.fetch_add(1, Ordering::Relaxed); + ProviderAdmissionGuard { + process, + provider, + active: &self.active, + } + } + + #[must_use] + pub fn active(&self) -> usize { + self.active.load(Ordering::Relaxed) + } +} + +pub struct ProviderAdmissionGuard<'a> { + process: SemaphoreGuardArc, + provider: SemaphoreGuardArc, + active: &'a AtomicUsize, +} + +impl Drop for ProviderAdmissionGuard<'_> { + fn drop(&mut self) { + let _ = (&self.process, &self.provider); + self.active.fetch_sub(1, Ordering::Relaxed); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn process_limit_covers_distinct_providers() { + smol::block_on(async { + let admission = Arc::new(ProviderAdmission::with_limits(1, 1)); + let first = admission.acquire("one").await; + let (started_tx, started_rx) = flume::bounded(1); + let admission_for_second = Arc::clone(&admission); + let second = smol::spawn(async move { + started_tx + .send_async(()) + .await + .expect("test receiver remains available"); + let permit = admission_for_second.acquire("two").await; + drop(permit); + }); + started_rx + .recv_async() + .await + .expect("waiter reached admission"); + drop(first); + second.await; + assert_eq!(admission.active(), 0); + }); + } + + #[test] + fn permit_releases_when_work_panics() { + let admission = Arc::new(ProviderAdmission::with_limits(1, 1)); + let admission_for_work = Arc::clone(&admission); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || { + smol::block_on(async move { + let _permit = admission_for_work.acquire("one").await; + panic!("simulated stream panic"); + }); + })); + assert!(result.is_err()); + assert_eq!(admission.active(), 0); + } +} diff --git a/n00n-providers/src/lib.rs b/n00n-providers/src/lib.rs index a55c873c6..9d0000418 100644 --- a/n00n-providers/src/lib.rs +++ b/n00n-providers/src/lib.rs @@ -6,6 +6,7 @@ #![allow(clippy::unwrap_or_default)] #![allow(clippy::default_trait_access)] +pub mod admission; pub(crate) mod error; pub mod manifest; pub mod model; diff --git a/n00n-ui/src/app/btw.rs b/n00n-ui/src/app/btw.rs index ef61f2e8a..8be99b9ff 100644 --- a/n00n-ui/src/app/btw.rs +++ b/n00n-ui/src/app/btw.rs @@ -18,6 +18,7 @@ say so; do not offer to look it up.\n"; const BTW_FALLBACK_SYSTEM: &str = "You are a helpful coding assistant. Answer concisely \ from the conversation context."; +const PROVIDER_EVENT_QUEUE_CAPACITY: usize = 256; /// The reminder leads so the model treats the question as a quick aside, not a task to act on. pub(crate) fn btw_question(question: &str) -> Message { @@ -52,11 +53,14 @@ async fn run_btw( messages: Vec, btw_tx: Sender, ) { - let (event_tx, event_rx) = flume::unbounded(); + let (event_tx, event_rx) = flume::bounded(PROVIDER_EVENT_QUEUE_CAPACITY); let tools = Value::Array(vec![]); let messages = n00n_providers::adapt_images_for_model(&model, &messages); let messages = n00n_providers::adapt_files_for_model(&model, &messages); + let _permit = n00n_providers::admission::ProviderAdmission::global() + .acquire(model.provider.as_ref()) + .await; let stream_fut = provider.stream_message( &model, &messages, diff --git a/plugins/batch/init.lua b/plugins/batch/init.lua index 9c2bdc2ab..490950783 100644 --- a/plugins/batch/init.lua +++ b/plugins/batch/init.lua @@ -499,6 +499,7 @@ n00n.api.register_tool({ name = "batch", description = description, kind = "execute", + workload = "orchestrator", audiences = { "main", "research_sub", "general_sub" }, defer_loading = true, schema = schema, diff --git a/plugins/glob/init.lua b/plugins/glob/init.lua index 1dcbf8651..6b7a0c061 100644 --- a/plugins/glob/init.lua +++ b/plugins/glob/init.lua @@ -17,6 +17,7 @@ end n00n.api.register_tool({ name = "glob", kind = "search", + workload = "cheap", modes = { "default", "research", "build", "compact" }, description = "Find files by glob pattern. Respects .gitignore. Returns matching paths sorted by mtime.", diff --git a/plugins/grep/init.lua b/plugins/grep/init.lua index d2a3e7dc6..ee18e327f 100644 --- a/plugins/grep/init.lua +++ b/plugins/grep/init.lua @@ -203,6 +203,7 @@ n00n.api.register_prompt_hint({ n00n.api.register_tool({ name = "grep", kind = "search", + workload = "cheap", modes = { "default", "research", "build", "compact" }, description = [[Search file contents using regex. Respects .gitignore. Results grouped by file, sorted by modification time. Prefer speculative parallel searches over sequential glob+grep. Do NOT wrap pattern in quotes or double-escape (e.g. `\[` not `\\[`). Multi-line matching auto-enabled when pattern contains `\n`, `(?s)`, or `(?m)`.]], diff --git a/plugins/read/init.lua b/plugins/read/init.lua index 49de1e040..d8ca5f0dc 100644 --- a/plugins/read/init.lua +++ b/plugins/read/init.lua @@ -206,6 +206,7 @@ n00n.api.register_prompt_hint({ n00n.api.register_tool({ name = "read", kind = "read", + workload = "cheap", modes = { "default", "research", "build", "compact" }, description = DESCRIPTION, diff --git a/plugins/task/init.lua b/plugins/task/init.lua index 417fa55fb..6caefc803 100644 --- a/plugins/task/init.lua +++ b/plugins/task/init.lua @@ -333,6 +333,7 @@ n00n.api.register_tool({ name = "task", description = description, kind = "execute", + workload = "orchestrator", audiences = { "main", "workflow" }, schema = schema, handler = handler, diff --git a/plugins/team/init.lua b/plugins/team/init.lua index 2a996f9a8..8ee75af20 100644 --- a/plugins/team/init.lua +++ b/plugins/team/init.lua @@ -1091,6 +1091,7 @@ n00n.api.register_tool({ name = "team", description = description, kind = "execute", + workload = "orchestrator", audiences = { "main", "workflow" }, schema = schema, timeout = TEAM_TIMEOUT_SECS, diff --git a/plugins/workflow/init.lua b/plugins/workflow/init.lua index 4d88742a1..09188130a 100644 --- a/plugins/workflow/init.lua +++ b/plugins/workflow/init.lua @@ -44,6 +44,15 @@ local DEFAULT_CONCURRENT_WORKFLOWS = 2 local HARD_MAX_CONCURRENT_AGENTS = 8 local HARD_MAX_CONCURRENT_WORKFLOWS = 4 local HARD_MAX_AGGREGATE_AGENTS = 12 +local HARD_MAX_AGENTS_PER_RUN = 64 +local MAX_PARALLEL_BRANCHES = 32 +local MAX_PIPELINE_ITEMS = 32 +local MAX_PIPELINE_STAGES = 16 +local MAX_SCRIPT_BYTES = 64 * 1024 +local MAX_INPUT_BYTES = 64 * 1024 +local MAX_JOURNAL_BYTES = 1024 * 1024 +local MAX_RESULT_BYTES = 32 * 1024 +local RESULT_TRUNCATED_MARKER = "\n[truncated]" local INVALID_RUN_ID_ERROR = "resume must be a run_id (hex letters/digits only, no path separators)" local RUN_ID_PATTERN = "^[%x]+$" local DEFAULT_TIMEOUT_SECS = 600 @@ -53,7 +62,7 @@ local description = [[Run sandboxed Lua workflow for multi-stage agent orchestra Start with meta({ name, description, phases }). Globals: agent({ prompt, subagent_type?, model_tier?, label?, output_schema? }) returns agent result; parallel(fns, { concurrency? }) runs branches; pipeline(items, stages, { concurrency? }) runs stages per item; phase(name, fn), log(...), inputs. -No n00n, os, io, require, print, or load. Scripts must be deterministic for resume replay, must return the final string, and are capped by max_agents_per_run (default 24, no hard maximum) with a runaway guard for repeated prompts and consecutive errors. Use task for one agent.]] +No n00n, os, io, require, print, or load. Scripts must be deterministic for resume replay, must return the final string, and are capped by max_agents_per_run (default 24, hard maximum 64) with a runaway guard for repeated prompts and consecutive errors. Use task for one agent.]] local schema = { type = "object", @@ -83,7 +92,7 @@ local opts = n00n.api.register_options({ max_agents_per_run = { default = DEFAULT_AGENTS_PER_RUN, min = 1, - desc = "Agent-call budget per workflow (default 24, no hard maximum).", + desc = "Agent-call budget per workflow (default 24, hard maximum 64).", }, max_concurrent_agents = { default = DEFAULT_CONCURRENT_AGENTS, @@ -102,7 +111,7 @@ local opts = n00n.api.register_options({ }, }) -local max_agents_per_run = opts.max_agents_per_run or DEFAULT_AGENTS_PER_RUN +local max_agents_per_run = math.min(opts.max_agents_per_run or DEFAULT_AGENTS_PER_RUN, HARD_MAX_AGENTS_PER_RUN) local max_concurrent_agents = math.min(opts.max_concurrent_agents, HARD_MAX_CONCURRENT_AGENTS) local max_concurrent_workflows = math.min(opts.max_concurrent_workflows, HARD_MAX_CONCURRENT_WORKFLOWS) local workflow_semaphore = n00n.async.semaphore(max_concurrent_workflows) @@ -275,6 +284,13 @@ local function load_journal(run_id, required) end return cache, path, "" end + local metadata, metadata_err = n00n.fs.metadata(path) + if not metadata then + return nil, path, nil, "failed to inspect workflow journal: " .. tostring(metadata_err) + end + if metadata.size and metadata.size > MAX_JOURNAL_BYTES then + return nil, path, nil, "workflow journal exceeds the " .. MAX_JOURNAL_BYTES .. " byte limit" + end local text, read_err = n00n.fs.read(path) if type(text) ~= "string" then return nil, path, nil, "failed to read workflow journal: " .. tostring(read_err) @@ -386,10 +402,23 @@ local function new_run_id(script) return n00n.workflow.hash(script .. "\0" .. tostring(os.time()) .. "\0" .. tostring(run_seq)) end +local function bounded_text(text, limit) + if #text <= limit then + return text + end + if limit <= #RESULT_TRUNCATED_MARKER then + return RESULT_TRUNCATED_MARKER:sub(1, limit) + end + return text:sub(1, limit - #RESULT_TRUNCATED_MARKER) .. RESULT_TRUNCATED_MARKER +end + local function parallel(fns, popts) if type(fns) ~= "table" then error("parallel: fns must be an array of functions", 0) end + if #fns > MAX_PARALLEL_BRANCHES then + error("parallel: branch count exceeds " .. MAX_PARALLEL_BRANCHES, 0) + end popts = popts or {} local concurrency = max_concurrent_agents if type(popts.concurrency) == "number" then @@ -437,6 +466,12 @@ local function pipeline(items, stages, popts) if type(stages) ~= "table" then error("pipeline: stages must be an array of functions", 0) end + if #items > MAX_PIPELINE_ITEMS then + error("pipeline: item count exceeds " .. MAX_PIPELINE_ITEMS, 0) + end + if #stages > MAX_PIPELINE_STAGES then + error("pipeline: stage count exceeds " .. MAX_PIPELINE_STAGES, 0) + end for i, stage in ipairs(stages) do if type(stage) ~= "function" then error("pipeline: stages[" .. i .. "] must be a function", 0) @@ -561,6 +596,7 @@ local function make_agent(ctx, progress, journal, logger, run_guard) else out = "" end + out = bounded_text(out, MAX_RESULT_BYTES) local gate = journal.lock:acquire() local io_ok, io_err = pcall(function() @@ -579,6 +615,9 @@ local function make_agent(ctx, progress, journal, logger, run_guard) error("failed to encode workflow journal entry: " .. tostring(encode_err), 0) end local next_text = (journal.text or "") .. line .. "\n" + if #next_text > MAX_JOURNAL_BYTES then + error("workflow journal exceeds the " .. MAX_JOURNAL_BYTES .. " byte limit", 0) + end local write_ok, write_err = n00n.fs.write(journal.path, next_text) if not write_ok then error("failed to write workflow journal: " .. tostring(write_err), 0) @@ -800,6 +839,24 @@ local function handler(input, ctx) if type(input.script) ~= "string" or input.script == "" then return { llm_output = SCRIPT_REQUIRED_ERROR, is_error = true } end + if #input.script > MAX_SCRIPT_BYTES then + return { + llm_output = "workflow script exceeds the " .. MAX_SCRIPT_BYTES .. " byte limit", + is_error = true, + } + end + if input.inputs then + local encoded_inputs, inputs_err = n00n.json.encode(input.inputs) + if not encoded_inputs then + return { llm_output = "workflow inputs are not serializable: " .. tostring(inputs_err), is_error = true } + end + if #encoded_inputs > MAX_INPUT_BYTES then + return { + llm_output = "workflow inputs exceed the " .. MAX_INPUT_BYTES .. " byte limit", + is_error = true, + } + end + end if input.timeout_secs and input.timeout_secs < ASYNC_RUNTIME_MIN_TIMEOUT_SECS then return { llm_output = "timeout_secs must be at least " .. ASYNC_RUNTIME_MIN_TIMEOUT_SECS, @@ -922,7 +979,9 @@ local function handler(input, ctx) if type(output) ~= "string" then output = tostring(output) end - return output .. "\n\n_run_id: `" .. run_id .. "` (pass as `resume` to continue)_" + local run_suffix = "\n\n_run_id: `" .. run_id .. "` (pass as `resume` to continue)_" + output = bounded_text(output, MAX_RESULT_BYTES - #run_suffix) + return output .. run_suffix end) if permit then permit:release() @@ -975,6 +1034,7 @@ n00n.api.register_tool({ name = "workflow", description = description, kind = "execute", + workload = "orchestrator", audiences = { "main" }, schema = schema, handler = handler, From 6ff6d69a58963515e928cc448d7b20b9d7ec0d7f Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 3 Aug 2026 15:13:06 -0400 Subject: [PATCH 02/24] docs(config): regenerate workflow limits --- site/docs/content/configuration/_index.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/site/docs/content/configuration/_index.md b/site/docs/content/configuration/_index.md index 048fc13c6..c6ac9ea6e 100644 --- a/site/docs/content/configuration/_index.md +++ b/site/docs/content/configuration/_index.md @@ -264,7 +264,7 @@ n00n.setup({ | Field | Type | Default | Min | Description | |-------|------|---------|-----|-------------| -| `max_agents_per_run` | integer | `24` | 1 | Agent-call budget per workflow (default 24, no hard maximum). | +| `max_agents_per_run` | integer | `24` | 1 | Agent-call budget per workflow (default 24, hard maximum 64). | | `max_concurrent_agents` | integer | `4` | 1 | Concurrency per parallel()/pipeline() (default 4, hard max 8). | | `max_concurrent_workflows` | integer | `2` | 1 | Concurrent workflows (default 2, hard max 4). | | `timeout_secs` | integer | `600` | 60 | Maximum deadline for one workflow run; per-run timeout_secs may only shorten it. | From ae8a076de7576c16b309cf5f2aa56f4e372832d1 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 3 Aug 2026 16:06:17 -0400 Subject: [PATCH 03/24] refactor(session): propagate authoritative caller identity --- n00n-agent/src/agent/run.rs | 13 ++++++ n00n-agent/src/tools/mod.rs | 2 + n00n-lua/src/api/async.rs | 2 +- n00n-lua/src/api/session.rs | 68 +++++++++++++++++++++++++++ n00n-lua/src/api/util/command.rs | 3 ++ n00n-lua/src/api/util/ctx.rs | 20 +++++++- n00n-lua/src/runtime.rs | 79 ++++++++++++++++++++++++-------- n00n-ui/src/event_loop.rs | 2 + src/cmd/tui_bridge.rs | 21 +++++++-- 9 files changed, 185 insertions(+), 25 deletions(-) diff --git a/n00n-agent/src/agent/run.rs b/n00n-agent/src/agent/run.rs index 1acdae91c..60ab9efa9 100644 --- a/n00n-agent/src/agent/run.rs +++ b/n00n-agent/src/agent/run.rs @@ -884,6 +884,7 @@ impl<'h> Agent<'h> { event_tx: self.event_tx.clone(), mode: Arc::clone(&self.mode), tool_use_id: None, + session_id: self.session_id.clone(), user_response_rx: self.user_response_rx.clone(), loaded_instructions: self.loaded_instructions.clone(), cancel: self.cancel.clone(), @@ -2177,6 +2178,18 @@ mod tests { (agent, event_rx) } + #[test] + fn tool_context_preserves_agent_session_id() { + let mut history = History::new(Vec::new()); + let (mut agent, _event_rx) = make_agent(MockProvider::new(Vec::new()), &mut history); + let session_id = SessionRef::generate(); + agent.session_id = Some(session_id.clone()); + + let ctx = agent.tool_context(); + + assert_eq!(ctx.session_id, Some(session_id)); + } + fn make_agent_with_config( provider: MockProvider, history: &mut History, diff --git a/n00n-agent/src/tools/mod.rs b/n00n-agent/src/tools/mod.rs index 038a7fc8f..c7e57d1b6 100644 --- a/n00n-agent/src/tools/mod.rs +++ b/n00n-agent/src/tools/mod.rs @@ -309,6 +309,7 @@ pub struct ToolContext { pub event_tx: EventSender, pub mode: Arc, pub tool_use_id: Option, + pub session_id: Option, pub user_response_rx: Option>>>, pub loaded_instructions: LoadedInstructions, pub cancel: CancelToken, @@ -602,6 +603,7 @@ pub fn interpreter_ctx( event_tx: event_tx.clone(), mode: Arc::new(mode.clone()), tool_use_id: None, + session_id: None, user_response_rx, loaded_instructions: LoadedInstructions::new(), cancel, diff --git a/n00n-lua/src/api/async.rs b/n00n-lua/src/api/async.rs index 8ffb91ec9..ff7231027 100644 --- a/n00n-lua/src/api/async.rs +++ b/n00n-lua/src/api/async.rs @@ -605,7 +605,7 @@ mod tests { fn cancelled_task_handle() -> TaskHandle { let (trigger, token) = CancelToken::new(); trigger.cancel(); - Arc::new(Mutex::new(TaskCell::new(token, None, None))) + Arc::new(Mutex::new(TaskCell::new(token, None, None, None))) } #[test_case(0 ; "zero_clamps_to_capacity_one")] diff --git a/n00n-lua/src/api/session.rs b/n00n-lua/src/api/session.rs index 141c1c345..a9ca75b28 100644 --- a/n00n-lua/src/api/session.rs +++ b/n00n-lua/src/api/session.rs @@ -7,6 +7,7 @@ use n00n_lua_macro::{lua_fn, lua_table}; use crate::api::util::command::{SessionReply, SessionRequest, UiAction}; use crate::api::util::convert::json_to_lua; +use crate::runtime::active_session_id; const NO_UI_ERR: &str = "no interactive UI attached"; @@ -127,6 +128,7 @@ async fn new( #[ctx] tx: Option>, opts: Option, ) -> LuaResult { + let caller_id = active_session_id(&lua); let (prompt, focus, parent_id) = match opts { Some(opts) => ( opts.get("prompt")?, @@ -142,6 +144,7 @@ async fn new( prompt, focus, parent_id, + caller_id, }, ) .await @@ -167,6 +170,7 @@ async fn prompt( text: String, opts: Option
, ) -> LuaResult { + let caller_id = active_session_id(&lua); let (id, steer, control) = match opts { Some(opts) => ( opts.get("session")?, @@ -183,6 +187,7 @@ async fn prompt( text, steer, control, + caller_id, }, ) .await @@ -234,9 +239,13 @@ lua_table! { #[cfg(test)] mod tests { use super::*; + use n00n_agent::cancel::CancelToken; + use n00n_storage::id::SessionRef; use serde_json::json; use test_case::test_case; + use crate::runtime::{TaskCell, TaskScope}; + fn lua_with_session(tx: Option>) -> Lua { let lua = Lua::new(); let t = create_session_table(&lua, tx).unwrap(); @@ -244,6 +253,63 @@ mod tests { lua } + #[test] + fn session_requests_attach_runtime_caller_id_not_lua_option() { + let (tx, rx) = flume::unbounded::(); + let caller_id = SessionRef::generate(); + let lua = lua_with_session(Some(tx)); + let _scope = TaskScope::new( + &lua, + TaskCell::new(CancelToken::none(), None, None, Some(caller_id.clone())), + ); + let expected_caller_id = caller_id.clone(); + let checker = std::thread::spawn(move || { + let Ok(UiAction::Session { + req: + SessionRequest::New { + caller_id: actual_caller_id, + .. + }, + reply_tx, + }) = rx.recv() + else { + panic!("expected new request"); + }; + assert_eq!(actual_caller_id.as_ref(), Some(&expected_caller_id)); + reply_tx.send(Ok(json!("child"))).unwrap(); + let Ok(UiAction::Session { + req: + SessionRequest::Prompt { + caller_id: actual_caller_id, + .. + }, + reply_tx, + }) = rx.recv() + else { + panic!("expected prompt request"); + }; + assert_eq!(actual_caller_id.as_ref(), Some(&expected_caller_id)); + reply_tx.send(Ok(json!("queued"))).unwrap(); + }); + + let (child_id, prompt_status): (String, String) = smol::block_on( + lua.load( + r#" + local child, new_err = session.new({ caller_id = "spoof" }) + if new_err then error(new_err) end + local status, prompt_err = session.prompt("hello", { caller_id = "spoof" }) + if prompt_err then error(prompt_err) end + return child, status + "#, + ) + .eval_async(), + ) + .unwrap(); + checker.join().unwrap(); + assert_eq!(child_id, "child"); + assert_eq!(prompt_status, "queued"); + } + #[test] fn live_without_ui_returns_error_pair() { let lua = lua_with_session(None); @@ -340,6 +406,7 @@ mod tests { text, steer, control, + caller_id, }, reply_tx, }) = rx.recv() @@ -350,6 +417,7 @@ mod tests { assert_eq!(text, "hi"); assert_eq!(steer, expected_steer); assert_eq!(control, expected_control); + assert_eq!(caller_id, None); reply_tx.send(Ok(json!("queued"))).unwrap(); }); let (val, err): (String, Option) = diff --git a/n00n-lua/src/api/util/command.rs b/n00n-lua/src/api/util/command.rs index c3d55db15..cd0f08c99 100644 --- a/n00n-lua/src/api/util/command.rs +++ b/n00n-lua/src/api/util/command.rs @@ -8,6 +8,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use arc_swap::ArcSwap; use mlua::RegistryKey; use n00n_agent::SharedBuf; +use n00n_storage::id::SessionRef; #[derive(Clone)] pub struct LuaCommandInfo { @@ -415,12 +416,14 @@ pub enum SessionRequest { prompt: Option, focus: bool, parent_id: Option, + caller_id: Option, }, Prompt { id: Option, text: String, steer: bool, control: bool, + caller_id: Option, }, Cancel { id: String, diff --git a/n00n-lua/src/api/util/ctx.rs b/n00n-lua/src/api/util/ctx.rs index 74a6396c7..92b23ed9d 100644 --- a/n00n-lua/src/api/util/ctx.rs +++ b/n00n-lua/src/api/util/ctx.rs @@ -10,6 +10,7 @@ use n00n_agent::tools::{ Deadline, FileReadTracker, LocalTools, ToolAudience, ToolContext, ToolLive, }; use n00n_config::{AgentConfig, ToolOutputLines}; +use n00n_storage::id::SessionRef; use crate::api::tool::ToolCallReply; use crate::api::ui::buf::BufHandle; @@ -97,6 +98,7 @@ enum Caps { config: Arc, workflow: bool, audience: ToolAudience, + session_id: Option, }, Restore { state: Option, @@ -130,6 +132,7 @@ impl LuaCtx { config: Arc::clone(&ctx.config), workflow: ctx.workflow, audience: ctx.audience, + session_id: ctx.session_id.clone(), }, ) } @@ -154,6 +157,14 @@ impl LuaCtx { } } + pub(crate) fn session_id(&self) -> Option { + match &self.caps { + Caps::Handler { agent, .. } => agent.session_id.clone(), + Caps::Start { session_id, .. } => session_id.clone(), + Caps::Restore { .. } => None, + } + } + fn config(&self) -> Option<&AgentConfig> { match &self.caps { Caps::Handler { agent, .. } => Some(&agent.config), @@ -388,6 +399,7 @@ mod tests { fn populated_ctx() -> ToolContext { let mut ctx = stub_ctx_with(&AgentMode::Build, None, Some(TOOL_USE_ID)); + ctx.session_id = Some(n00n_storage::id::SessionRef::generate()); ctx.deadline = Deadline::after(Duration::from_mins(1)); ctx.tool_output_lines = ToolOutputLines { bash: 999, @@ -408,9 +420,12 @@ mod tests { } #[test] - fn agent_context_keeps_tool_use_id_and_resets_per_call_state() { - let agent = AgentContext::from(&populated_ctx()); + fn agent_context_keeps_tool_use_id_and_session_id_and_resets_per_call_state() { + let ctx = populated_ctx(); + let expected_session_id = ctx.session_id.clone(); + let agent = AgentContext::from(&ctx); assert_eq!(agent.tool_use_id.as_deref(), Some(TOOL_USE_ID)); + assert_eq!(agent.session_id, expected_session_id); assert!(matches!(agent.deadline, Deadline::None)); assert_eq!(agent.tool_output_lines, ToolOutputLines::default()); assert!(agent.local_tools.is_empty()); @@ -431,6 +446,7 @@ mod tests { ); let inner = agent.to_tool_context(); assert_eq!(inner.tool_use_id, None); + assert_eq!(inner.session_id, agent.session_id); assert!(inner.live_sink.is_none(), "sink must not be inherited"); assert_eq!(agent.tool_use_id.as_deref(), Some(TOOL_USE_ID)); } diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index 53edec22a..ac8368a05 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -22,6 +22,7 @@ use n00n_agent::{BufferSnapshot, SharedBuf, SnapshotLine, SnapshotSpan, SpanStyl use serde_json::Value; use n00n_config::RawConfig; +use n00n_storage::id::SessionRef; use crate::api::autocmd::AutocmdStore; use crate::api::create_n00n_global; @@ -298,6 +299,7 @@ pub(crate) struct TaskCell { /// Forwards live bufs and annotations to a parent /// `n00n.agent.call_tool(on_live_buf/on_annotation)`. pub(crate) live_sink: Option>, + pub(crate) session_id: Option, /// When `Some`, `n00n.async.run` tasks queue here instead of the global /// `SpawnQueue` so restore can run them inline before snapshotting. pub(crate) inline_spawn: Option>, @@ -315,6 +317,7 @@ impl TaskCell { cancel: CancelToken, deadline: Option, live: Option, + session_id: Option, ) -> Self { Self { cancel, @@ -325,6 +328,7 @@ impl TaskCell { live, root_buf: None, live_sink: None, + session_id, inline_spawn: None, bufs_claim: Weak::new(), async_tasks: Cell::new(0), @@ -485,7 +489,7 @@ impl TaskScope { /// (stale handle looks cancelled). Prefer [`run_detached`] over raw /// scopes. pub(crate) fn detached(lua: &Lua) -> Self { - Self::new(lua, TaskCell::new(CancelToken::none(), None, None)) + Self::new(lua, TaskCell::new(CancelToken::none(), None, None, None)) } pub(crate) fn handle(&self) -> &TaskHandle { @@ -592,6 +596,11 @@ pub(crate) fn active_task(lua: &Lua) -> TaskHandle { ) } +pub(crate) fn active_session_id(lua: &Lua) -> Option { + let handle = lua.app_data_ref::()?; + lock_cell(&handle).session_id.clone() +} + pub(crate) fn with_task_jobs(lua: &Lua, f: impl FnOnce(&mut JobStore) -> R) -> R { f(&mut lock_cell(&active_task(lua)).jobs) } @@ -608,12 +617,17 @@ pub(crate) fn with_live_ctx(lua: &Lua, f: impl FnOnce(&LiveCtx) -> R) -> Opti pub(crate) fn enqueue_async_task(lua: &Lua, work_fn: RegistryKey) -> Result<(), mlua::Error> { let handle = lua.app_data_ref::(); - let (cancel, live_ctx, parent_deadline) = match &handle { + let (cancel, live_ctx, parent_deadline, session_id) = match &handle { Some(h) => { let cell = lock_cell(h); - (cell.cancel.clone(), cell.live.clone(), cell.deadline.get()) - } - None => (CancelToken::none(), None, None), + ( + cell.cancel.clone(), + cell.live.clone(), + cell.deadline.get(), + cell.session_id.clone(), + ) + } + None => (CancelToken::none(), None, None, None), }; let deadline = @@ -624,6 +638,7 @@ pub(crate) fn enqueue_async_task(lua: &Lua, work_fn: RegistryKey) -> Result<(), cancel, deadline, live_ctx, + session_id, owner: None, parent: None, }; @@ -788,6 +803,7 @@ pub(crate) struct PendingAsyncTask { pub cancel: CancelToken, pub deadline: Option, pub live_ctx: Option, + pub session_id: Option, pub owner: Option>, /// Parent task that spawned this `noon.async.run` task, if any. /// Used to decrement the parent's `async_tasks` counter on completion. @@ -891,7 +907,12 @@ fn spawn_async_task( let scope = TaskScope::new( &lua, - TaskCell::new(task.cancel.clone(), task.deadline, task.live_ctx.clone()), + TaskCell::new( + task.cancel.clone(), + task.deadline, + task.live_ctx.clone(), + task.session_id.clone(), + ), ); let result = scope .scope_future(run_work_fn(&lua, &task.work_fn, task.deadline)) @@ -1731,6 +1752,7 @@ async fn restore_item( event_tx: n00n_agent::EventSender::new(dummy_tx, 0), tool_use_id: item.tool_use_id.clone(), }), + None, ); let ctx = lua @@ -2015,7 +2037,11 @@ async fn run_tool_start( live: LiveCtx, ctx: Box, ) { - let scope = TaskScope::new(lua, TaskCell::new(ctx.cancel.clone(), None, Some(live))); + let session_id = ctx.session_id(); + let scope = TaskScope::new( + lua, + TaskCell::new(ctx.cancel.clone(), None, Some(live), session_id), + ); let run = async { let input_lua = json_to_lua(lua, &input)?; let ctx_ud = lua.create_userdata(*ctx)?; @@ -2065,6 +2091,7 @@ async fn run_tool_call( let (finish_tx, finish_rx) = flume::bounded::(1); ctx.finish_tx = Some(finish_tx); let cancel = ctx.cancel.clone(); + let session_id = ctx.session_id(); let input_lua = match json_to_lua(&lua, &input) { Ok(v) => v, @@ -2081,7 +2108,7 @@ async fn run_tool_call( Err(e) => return ToolCallReply::err(strip_traceback(&e)), }; let live_id = live.as_ref().map(|l| l.tool_use_id.clone()); - let mut cell = TaskCell::new(cancel, deadline, live); + let mut cell = TaskCell::new(cancel, deadline, live, session_id.clone()); cell.live_sink = live_sink; let scope = TaskScope::new(&lua, cell); let handle = Arc::clone(scope.handle()); @@ -2164,7 +2191,7 @@ async fn run_tool_call( // A fresh cell, because the original's cancel token and // deadline are stale: the watchdog interrupt would use them to // kill warm clicks. - let mut cell = TaskCell::new(CancelToken::none(), None, None); + let mut cell = TaskCell::new(CancelToken::none(), None, None, session_id); cell.root_buf = Some(root); let mut warm = warm_tools.borrow_mut(); warm.push_back(WarmTool { @@ -2653,7 +2680,7 @@ mod tests { } fn task_cell(live: Option) -> TaskCell { - TaskCell::new(CancelToken::none(), None, live) + TaskCell::new(CancelToken::none(), None, live, None) } #[test] @@ -2812,7 +2839,7 @@ mod tests { #[test] fn enqueue_async_task_routes_to_inline_spawn_when_set() { let lua = enqueue_test_lua(); - let scope = set_active(&lua, TaskCell::new(CancelToken::none(), None, None)); + let scope = set_active(&lua, TaskCell::new(CancelToken::none(), None, None, None)); lock_cell(scope.handle()).inline_spawn = Some(Vec::new()); enqueue_async_task(&lua, enqueue_dummy(&lua)).unwrap(); @@ -2840,7 +2867,7 @@ mod tests { fn enqueue_async_task_inherits_cancel_token() { let lua = enqueue_test_lua(); let (trigger, token) = CancelToken::new(); - let _h = set_active(&lua, TaskCell::new(token, None, None)); + let _h = set_active(&lua, TaskCell::new(token, None, None, None)); enqueue_async_task(&lua, enqueue_dummy(&lua)).unwrap(); let queue = lua.app_data_ref::().unwrap(); @@ -2853,13 +2880,28 @@ mod tests { ); } + #[test] + fn enqueue_async_task_inherits_session_id() { + let lua = enqueue_test_lua(); + let session_id = n00n_storage::id::SessionRef::generate(); + let _h = set_active( + &lua, + TaskCell::new(CancelToken::none(), None, None, Some(session_id.clone())), + ); + enqueue_async_task(&lua, enqueue_dummy(&lua)).unwrap(); + + let queue = lua.app_data_ref::().unwrap(); + let queued = queue.rx.try_recv().unwrap(); + assert_eq!(queued.session_id, Some(session_id)); + } + #[test] fn enqueue_async_task_extends_expired_parent_to_minimum_deadline() { let lua = enqueue_test_lua(); let parent_deadline = Instant::now().checked_sub(Duration::from_secs(10)).unwrap(); let _h = set_active( &lua, - TaskCell::new(CancelToken::none(), Some(parent_deadline), None), + TaskCell::new(CancelToken::none(), Some(parent_deadline), None, None), ); let before = Instant::now(); @@ -2879,7 +2921,7 @@ mod tests { let parent_deadline = Instant::now() + Duration::from_mins(10); let _h = set_active( &lua, - TaskCell::new(CancelToken::none(), Some(parent_deadline), None), + TaskCell::new(CancelToken::none(), Some(parent_deadline), None, None), ); enqueue_async_task(&lua, enqueue_dummy(&lua)).unwrap(); @@ -2891,7 +2933,7 @@ mod tests { #[test] fn enqueue_async_task_without_parent_deadline_has_no_deadline() { let lua = enqueue_test_lua(); - let _h = set_active(&lua, TaskCell::new(CancelToken::none(), None, None)); + let _h = set_active(&lua, TaskCell::new(CancelToken::none(), None, None, None)); enqueue_async_task(&lua, enqueue_dummy(&lua)).unwrap(); @@ -2904,7 +2946,7 @@ mod tests { use crate::api::ui::buf::HandlerSlot; let lua = enqueue_test_lua(); - let scope = set_active(&lua, TaskCell::new(CancelToken::none(), None, None)); + let scope = set_active(&lua, TaskCell::new(CancelToken::none(), None, None, None)); let handle = Arc::clone(scope.handle()); let buf = Arc::new(SharedBuf::new()); @@ -2945,6 +2987,7 @@ mod tests { cancel, deadline, live_ctx: None, + session_id: None, owner: None, parent: None, } @@ -3000,7 +3043,7 @@ mod tests { fn cancelled_handle() -> TaskHandle { let (trigger, token) = CancelToken::new(); trigger.cancel(); - Arc::new(Mutex::new(TaskCell::new(token, None, None))) + Arc::new(Mutex::new(TaskCell::new(token, None, None, None))) } #[test] @@ -3040,7 +3083,7 @@ mod tests { apply_jit(&lua, true); let deadline = Instant::now() + Duration::from_millis(20); - let cell = TaskCell::new(CancelToken::none(), Some(deadline), None); + let cell = TaskCell::new(CancelToken::none(), Some(deadline), None, None); lua.set_app_data::(Arc::new(Mutex::new(cell))); let err = hot_loop_expecting_kill(&lua); diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index d23dd9959..abe65162f 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -945,6 +945,7 @@ impl<'t> EventLoop<'t> { prompt, focus, parent_id, + caller_id: _, } => { let mut session = { let slot = self.ctx.model_slot.load(); @@ -984,6 +985,7 @@ impl<'t> EventLoop<'t> { text, steer, control, + caller_id: _, } => { let idx = match id { None => Ok(self.focused), diff --git a/src/cmd/tui_bridge.rs b/src/cmd/tui_bridge.rs index 1aa3027d0..9a2a62e61 100644 --- a/src/cmd/tui_bridge.rs +++ b/src/cmd/tui_bridge.rs @@ -130,6 +130,7 @@ fn message_one( text: text.to_owned(), steer: opts.steer, control: opts.control, + caller_id: None, }, ) .map_err(|e| map_not_found(id, e))?; @@ -150,6 +151,7 @@ fn resume_one(tx: &flume::Sender, id: &str) -> ControlResult<()> { text: prompt, steer: true, control: true, + caller_id: None, }, ) .map_err(|e| map_not_found(id, e))?; @@ -361,10 +363,16 @@ mod tests { text, steer, control, + caller_id, } => { - if id.as_deref() != Some("sess-1") || text != "hi" || !steer || !control { + if id.as_deref() != Some("sess-1") + || text != "hi" + || !steer + || !control + || caller_id.is_some() + { let _ = reply_tx.send(Err(format!( - "unexpected prompt id={id:?} text={text:?} steer={steer} control={control}" + "unexpected prompt id={id:?} text={text:?} steer={steer} control={control} caller_id={caller_id:?}" ))); return; } @@ -412,10 +420,15 @@ mod tests { text, steer, control, + caller_id, } => { - if id.as_deref() != Some("sess-1") || !steer || !control { + if id.as_deref() != Some("sess-1") + || !steer + || !control + || caller_id.is_some() + { let _ = reply_tx.send(Err(format!( - "unexpected prompt id={id:?} steer={steer} control={control}" + "unexpected prompt id={id:?} steer={steer} control={control} caller_id={caller_id:?}" ))); return; } From f6061074324ac8fc616df76ef0c96617f47af1fa Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 3 Aug 2026 18:46:26 -0400 Subject: [PATCH 04/24] feat(session): add configurable lineage guard --- n00n-config/src/lib.rs | 52 +++ n00n-ui/src/event_loop.rs | 16 + n00n-ui/src/lib.rs | 1 + n00n-ui/src/session_lineage.rs | 566 +++++++++++++++++++++++++++++++++ 4 files changed, 635 insertions(+) create mode 100644 n00n-ui/src/session_lineage.rs diff --git a/n00n-config/src/lib.rs b/n00n-config/src/lib.rs index 900f1898d..1e401f738 100644 --- a/n00n-config/src/lib.rs +++ b/n00n-config/src/lib.rs @@ -29,6 +29,9 @@ pub const MIN_MAX_INPUT_LINES: u32 = 1; pub const DEFAULT_MCP_TOOL_DESC_MAX_CHARS: usize = 200; pub const DEFAULT_MAX_CONTINUATION_TURNS: u32 = 3; +pub const DEFAULT_MAX_DEPTH: usize = 4; +pub const DEFAULT_MAX_TOTAL_DESCENDANTS: usize = 16; +pub const DEFAULT_MAX_ACTIVE_DESCENDANTS: usize = 8; pub const DEFAULT_COMPACTION_BUFFER: CompactionBuffer = CompactionBuffer::Percent(20); pub const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 10; @@ -44,6 +47,9 @@ pub const DEFAULT_INPUT_HISTORY_SIZE: usize = 100; pub const MIN_OUTPUT_BYTES: usize = 1024; pub const MIN_OUTPUT_LINES: usize = 10; pub const MIN_MAX_CONTINUATION_TURNS: u32 = 1; +pub const MIN_MAX_DEPTH: usize = 1; +pub const MIN_MAX_TOTAL_DESCENDANTS: usize = 1; +pub const MIN_MAX_ACTIVE_DESCENDANTS: usize = 1; pub const MIN_COMPACTION_BUFFER: u32 = 1_000; const MAX_COMPACTION_PERCENT: u8 = 99; const COMPACTION_BUFFER_EXPECTED: &str = @@ -191,6 +197,16 @@ pub enum ConfigError { "invalid config: agent.fusion.sidekick_tier must be weak, medium, or strong, got {tier:?}" )] InvalidFusionSidekickTier { tier: Tier }, + #[error( + "invalid config: agent lineage limits require max_depth <= max_total_descendants and \ + max_active_descendants <= max_total_descendants (got max_depth={max_depth}, \ + max_total_descendants={max_total_descendants}, max_active_descendants={max_active_descendants})" + )] + InvalidLineageLimits { + max_depth: usize, + max_total_descendants: usize, + max_active_descendants: usize, + }, } fn check( @@ -492,6 +508,9 @@ pub struct AgentFileConfig { pub max_output_bytes: Option, pub max_output_lines: Option, pub max_continuation_turns: Option, + pub max_depth: Option, + pub max_total_descendants: Option, + pub max_active_descendants: Option, pub compaction_buffer: Option, pub mcp_tool_desc_max_chars: Option, pub dynamic_tools: Option, @@ -520,6 +539,9 @@ impl AgentFileConfig { max_output_bytes, max_output_lines, max_continuation_turns, + max_depth, + max_total_descendants, + max_active_descendants, compaction_buffer, mcp_tool_desc_max_chars ); @@ -1095,6 +1117,15 @@ pub struct AgentConfig { #[config(default = DEFAULT_MAX_CONTINUATION_TURNS, min = MIN_MAX_CONTINUATION_TURNS, desc = "Max automatic continuation turns")] pub max_continuation_turns: u32, + #[config(default = DEFAULT_MAX_DEPTH, min = MIN_MAX_DEPTH, desc = "Maximum session lineage depth")] + pub max_depth: usize, + + #[config(default = DEFAULT_MAX_TOTAL_DESCENDANTS, min = MIN_MAX_TOTAL_DESCENDANTS, desc = "Maximum total descendants per session lineage root")] + pub max_total_descendants: usize, + + #[config(default = DEFAULT_MAX_ACTIVE_DESCENDANTS, min = MIN_MAX_ACTIVE_DESCENDANTS, desc = "Maximum active descendants per session lineage root")] + pub max_active_descendants: usize, + #[config(default = DEFAULT_COMPACTION_BUFFER, ty = "u32 | string", default_doc = "20%", desc = "Context reserved for compaction: token count or percent of the context window (e.g. \"20%\")")] pub compaction_buffer: CompactionBuffer, @@ -1193,6 +1224,13 @@ impl AgentConfig { max_continuation_turns: file .max_continuation_turns .unwrap_or_else(|| DEFAULT_MAX_CONTINUATION_TURNS), + max_depth: file.max_depth.unwrap_or_else(|| DEFAULT_MAX_DEPTH), + max_total_descendants: file + .max_total_descendants + .unwrap_or_else(|| DEFAULT_MAX_TOTAL_DESCENDANTS), + max_active_descendants: file + .max_active_descendants + .unwrap_or_else(|| DEFAULT_MAX_ACTIVE_DESCENDANTS), compaction_buffer: file .compaction_buffer .unwrap_or_else(|| DEFAULT_COMPACTION_BUFFER), @@ -1206,6 +1244,19 @@ impl AgentConfig { fusion, } } + + fn validate_lineage_limits(&self) -> Result<(), ConfigError> { + if self.max_depth > self.max_total_descendants + || self.max_active_descendants > self.max_total_descendants + { + return Err(ConfigError::InvalidLineageLimits { + max_depth: self.max_depth, + max_total_descendants: self.max_total_descendants, + max_active_descendants: self.max_active_descendants, + }); + } + Ok(()) + } } #[derive(Debug, Clone, ConfigSection)] @@ -1380,6 +1431,7 @@ impl Config { pub fn validate(&self) -> Result<(), ConfigError> { self.ui.validate_all()?; self.agent.validate()?; + self.agent.validate_lineage_limits()?; if self.agent.fusion.sidekick_tier == Tier::Compaction { return Err(ConfigError::InvalidFusionSidekickTier { tier: self.agent.fusion.sidekick_tier, diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index abe65162f..80202aab6 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -48,6 +48,7 @@ use crate::components::{ Action, DisplayMessage, DisplayRole, ExitRequest, Status, SubmissionDispatch, }; use crate::input::InputReader; +use crate::session_lineage::{LineageLimits, LiveSession, SessionLineageGuard}; use crate::color_compat; use crate::storage_writer::StorageWriter; @@ -270,6 +271,7 @@ pub(crate) struct EventLoop<'t> { terminal: &'t mut ratatui::DefaultTerminal, sessions: Vec, focused: usize, + lineage: SessionLineageGuard, ctx: SpawnCtx, input: InputReader, warn_rx: flume::Receiver, @@ -487,6 +489,19 @@ impl<'t> EventLoop<'t> { let picker = Arc::new(terminal_image::picker()); + let lineage = SessionLineageGuard::from_live( + sessions.iter().map(|session| LiveSession { + id: session.id, + parent_id: session.meta.parent_id, + }), + LineageLimits { + max_depth: config.max_depth, + max_total_descendants: config.max_total_descendants, + max_active_descendants: config.max_active_descendants, + }, + ) + .map_err(|error| eyre!("invalid live session lineage: {error}"))?; + let ctx = SpawnCtx { storage, config, @@ -534,6 +549,7 @@ impl<'t> EventLoop<'t> { terminal, sessions: runtimes, focused, + lineage, ctx, input: InputReader::spawn()?, warn_rx: bg.warn_rx, diff --git a/n00n-ui/src/lib.rs b/n00n-ui/src/lib.rs index 32fe609e6..e425634f5 100644 --- a/n00n-ui/src/lib.rs +++ b/n00n-ui/src/lib.rs @@ -29,6 +29,7 @@ pub mod update; mod agent; mod event_loop; mod input; +mod session_lineage; mod terminal; mod terminal_image; diff --git a/n00n-ui/src/session_lineage.rs b/n00n-ui/src/session_lineage.rs new file mode 100644 index 000000000..6f41d7c14 --- /dev/null +++ b/n00n-ui/src/session_lineage.rs @@ -0,0 +1,566 @@ +use std::collections::{HashMap, HashSet}; + +use n00n_storage::id::n00nId; +use thiserror::Error; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct LineageLimits { + pub(crate) max_depth: usize, + pub(crate) max_total_descendants: usize, + pub(crate) max_active_descendants: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct LiveSession { + pub(crate) id: n00nId, + pub(crate) parent_id: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct SessionLineage { + pub(crate) caller: n00nId, + pub(crate) root: n00nId, + pub(crate) parent: Option, + pub(crate) depth: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct DescendantCounts { + pub(crate) total: usize, + pub(crate) active: usize, + pub(crate) reserved: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct NewReservation { + id: u64, +} + +#[derive(Debug, Error, PartialEq, Eq)] +pub(crate) enum LineageError { + #[error("caller session is not live: {0}")] + CallerNotLive(n00nId), + #[error("session is not live: {0}")] + TargetNotLive(n00nId), + #[error("session is not known: {0}")] + UnknownSession(n00nId), + #[error("session lineage parent must match caller")] + ParentMismatch, + #[error("session lineage contains a cycle at {0}")] + Cycle(n00nId), + #[error("session lineage depth limit exceeded: {limit}")] + DepthExceeded { limit: usize }, + #[error("session lineage total descendant limit exceeded: {limit}")] + TotalDescendantsExceeded { limit: usize }, + #[error("session lineage active descendant limit exceeded: {limit}")] + ActiveDescendantsExceeded { limit: usize }, + #[error("prompt target is outside the caller lineage")] + UnauthorizedTarget, + #[error("session already exists: {0}")] + DuplicateSession(n00nId), + #[error("session lineage parent changed for {id}")] + ParentChanged { id: n00nId }, + #[error("session lineage reservation is unknown")] + UnknownReservation, + #[error("session lineage reservation id space exhausted")] + ReservationIdExhausted, +} + +#[derive(Debug, Clone, Copy)] +struct SessionNode { + parent_id: Option, + active: bool, +} + +#[derive(Debug, Clone, Copy)] +struct PendingReservation { + caller: n00nId, + parent: n00nId, + root: n00nId, + depth: usize, +} + +pub(crate) struct SessionLineageGuard { + limits: LineageLimits, + sessions: HashMap, + reservations: HashMap, + next_reservation_id: u64, +} + +impl SessionLineageGuard { + pub(crate) fn from_live( + sessions: impl IntoIterator, + limits: LineageLimits, + ) -> Result { + let mut guard = Self { + limits, + sessions: HashMap::new(), + reservations: HashMap::new(), + next_reservation_id: 1, + }; + for session in sessions { + if guard + .sessions + .insert( + session.id, + SessionNode { + parent_id: session.parent_id, + active: true, + }, + ) + .is_some() + { + return Err(LineageError::DuplicateSession(session.id)); + } + } + guard.validate_graph()?; + Ok(guard) + } + + pub(crate) fn activate_runtime(&mut self, session: LiveSession) -> Result<(), LineageError> { + if let Some(existing) = self.sessions.get(&session.id) { + if existing.active { + return Err(LineageError::DuplicateSession(session.id)); + } + if existing.parent_id != session.parent_id { + return Err(LineageError::ParentChanged { id: session.id }); + } + self.sessions + .get_mut(&session.id) + .ok_or(LineageError::UnknownSession(session.id))? + .active = true; + if let Err(error) = self.validate_graph() { + if let Some(node) = self.sessions.get_mut(&session.id) { + node.active = false; + } + return Err(error); + } + return Ok(()); + } + + self.sessions.insert( + session.id, + SessionNode { + parent_id: session.parent_id, + active: true, + }, + ); + if let Err(error) = self.validate_graph() { + self.sessions.remove(&session.id); + return Err(error); + } + Ok(()) + } + + pub(crate) fn remove_runtime(&mut self, id: n00nId) -> Result<(), LineageError> { + let node = self + .sessions + .get_mut(&id) + .ok_or(LineageError::UnknownSession(id))?; + node.active = false; + Ok(()) + } + + pub(crate) fn lineage(&self, caller: n00nId) -> Result { + let node = self + .sessions + .get(&caller) + .ok_or(LineageError::CallerNotLive(caller))?; + if !node.active { + return Err(LineageError::CallerNotLive(caller)); + } + self.lineage_for(caller).map_err(|error| match error { + LineageError::UnknownSession(_) => LineageError::CallerNotLive(caller), + error => error, + }) + } + + pub(crate) fn reserve_new( + &mut self, + caller: n00nId, + explicit_parent: Option, + ) -> Result { + self.validate_graph()?; + let caller_lineage = self.lineage(caller)?; + let parent = explicit_parent.unwrap_or(caller); + if parent != caller { + return Err(LineageError::ParentMismatch); + } + let depth = caller_lineage + .depth + .checked_add(1) + .ok_or(LineageError::DepthExceeded { + limit: self.limits.max_depth, + })?; + if depth > self.limits.max_depth { + return Err(LineageError::DepthExceeded { + limit: self.limits.max_depth, + }); + } + + let counts = self.descendant_counts(caller_lineage.root)?; + if limit_reached( + counts.total, + counts.reserved, + self.limits.max_total_descendants, + ) { + return Err(LineageError::TotalDescendantsExceeded { + limit: self.limits.max_total_descendants, + }); + } + if limit_reached( + counts.active, + counts.reserved, + self.limits.max_active_descendants, + ) { + return Err(LineageError::ActiveDescendantsExceeded { + limit: self.limits.max_active_descendants, + }); + } + + let id = self.next_reservation_id; + self.next_reservation_id = self + .next_reservation_id + .checked_add(1) + .ok_or(LineageError::ReservationIdExhausted)?; + self.reservations.insert( + id, + PendingReservation { + caller, + parent, + root: caller_lineage.root, + depth, + }, + ); + Ok(NewReservation { id }) + } + + pub(crate) fn commit_new( + &mut self, + reservation: NewReservation, + child_id: n00nId, + ) -> Result<(), LineageError> { + let pending = self + .reservations + .remove(&reservation.id) + .ok_or(LineageError::UnknownReservation)?; + if self.sessions.contains_key(&child_id) { + return Err(LineageError::DuplicateSession(child_id)); + } + let caller_lineage = self.lineage(pending.caller)?; + if caller_lineage.root != pending.root || caller_lineage.depth + 1 != pending.depth { + return Err(LineageError::UnknownReservation); + } + self.sessions.insert( + child_id, + SessionNode { + parent_id: Some(pending.parent), + active: true, + }, + ); + if let Err(error) = self.validate_graph() { + self.sessions.remove(&child_id); + return Err(error); + } + Ok(()) + } + + pub(crate) fn release(&mut self, reservation: NewReservation) -> Result<(), LineageError> { + self.reservations + .remove(&reservation.id) + .map(|_| ()) + .ok_or(LineageError::UnknownReservation) + } + + pub(crate) fn authorize_prompt( + &self, + caller: n00nId, + explicit_target: Option, + ) -> Result { + self.validate_graph()?; + let caller_lineage = self.lineage(caller)?; + let target = explicit_target.unwrap_or(caller); + let target_node = self + .sessions + .get(&target) + .ok_or(LineageError::UnknownSession(target))?; + if !target_node.active { + return Err(LineageError::TargetNotLive(target)); + } + let caller_path = self.path_from(caller)?; + let target_path = self.path_from(target)?; + if caller_lineage.caller == target + || caller_path.contains(&target) + || target_path.contains(&caller) + { + return Ok(target); + } + Err(LineageError::UnauthorizedTarget) + } + + pub(crate) fn descendant_counts(&self, root: n00nId) -> Result { + let mut total = 0; + let mut active = 0; + for (&id, node) in &self.sessions { + let lineage = self.lineage_for(id)?; + if lineage.root != root || id == root { + continue; + } + total += 1; + if node.active { + active += 1; + } + } + let reserved = self + .reservations + .values() + .filter(|reservation| reservation.root == root) + .count(); + Ok(DescendantCounts { + total, + active, + reserved, + }) + } + + fn validate_graph(&self) -> Result<(), LineageError> { + for &id in self.sessions.keys() { + self.path_from(id)?; + } + Ok(()) + } + + fn lineage_for(&self, id: n00nId) -> Result { + let path = self.path_from(id)?; + let node = self + .sessions + .get(&id) + .ok_or(LineageError::UnknownSession(id))?; + let root = path + .last() + .copied() + .ok_or(LineageError::UnknownSession(id))?; + Ok(SessionLineage { + caller: id, + root, + parent: node.parent_id, + depth: path.len() - 1, + }) + } + + fn path_from(&self, start: n00nId) -> Result, LineageError> { + if !self.sessions.contains_key(&start) { + return Err(LineageError::UnknownSession(start)); + } + let mut path = Vec::new(); + let mut seen = HashSet::new(); + let mut current = start; + loop { + if !seen.insert(current) { + return Err(LineageError::Cycle(current)); + } + path.push(current); + let parent = self + .sessions + .get(¤t) + .ok_or(LineageError::UnknownSession(current))? + .parent_id; + let Some(parent) = parent else { + return Ok(path); + }; + if !self.sessions.contains_key(&parent) { + return Ok(path); + } + current = parent; + } + } +} + +fn limit_reached(committed: usize, reserved: usize, limit: usize) -> bool { + committed >= limit || reserved >= limit.saturating_sub(committed) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn id(value: u16) -> n00nId { + format!("00000000-0000-7000-8000-{value:012x}") + .parse() + .expect("valid test id") + } + + fn session(id: n00nId, parent_id: Option) -> LiveSession { + LiveSession { id, parent_id } + } + + fn limits( + max_depth: usize, + max_total_descendants: usize, + max_active_descendants: usize, + ) -> LineageLimits { + LineageLimits { + max_depth, + max_total_descendants, + max_active_descendants, + } + } + + #[test] + fn independent_roots_have_independent_limits() { + let root_a = id(1); + let root_b = id(2); + let child_a = id(3); + let child_b = id(4); + let mut guard = SessionLineageGuard::from_live( + [session(root_a, None), session(root_b, None)], + limits(4, 1, 1), + ) + .expect("valid roots"); + + let reservation = guard.reserve_new(root_a, None).expect("root A capacity"); + guard.commit_new(reservation, child_a).expect("child A"); + assert!(matches!( + guard.reserve_new(root_a, None), + Err(LineageError::TotalDescendantsExceeded { .. }) + )); + + let reservation = guard.reserve_new(root_b, None).expect("root B capacity"); + guard.commit_new(reservation, child_b).expect("child B"); + assert_eq!(guard.descendant_counts(root_a).expect("counts").total, 1); + assert_eq!(guard.descendant_counts(root_b).expect("counts").total, 1); + } + + #[test] + fn new_and_prompt_reject_spoofed_relationships() { + let root = id(1); + let sibling = id(2); + let foreign = id(3); + let mut guard = SessionLineageGuard::from_live( + [ + session(root, None), + session(sibling, Some(root)), + session(foreign, None), + ], + limits(4, 4, 4), + ) + .expect("valid graph"); + + assert!(matches!( + guard.reserve_new(root, Some(foreign)), + Err(LineageError::ParentMismatch) + )); + assert!(matches!( + guard.reserve_new(id(99), None), + Err(LineageError::CallerNotLive(_)) + )); + assert_eq!( + guard + .authorize_prompt(root, Some(sibling)) + .expect("descendant"), + sibling + ); + assert!(matches!( + guard.authorize_prompt(sibling, Some(foreign)), + Err(LineageError::UnauthorizedTarget) + )); + } + + #[test] + fn cycles_are_rejected() { + let first = id(1); + let second = id(2); + assert!(matches!( + SessionLineageGuard::from_live( + [session(first, Some(second)), session(second, Some(first))], + limits(4, 4, 4), + ), + Err(LineageError::Cycle(_)) + )); + } + + #[test] + fn depth_total_and_active_limits_are_distinct() { + let root = id(1); + let child = id(2); + let grandchild = id(3); + let mut guard = SessionLineageGuard::from_live( + [session(root, None), session(child, Some(root))], + limits(2, 2, 2), + ) + .expect("valid graph"); + let reservation = guard.reserve_new(child, None).expect("depth one"); + guard + .commit_new(reservation, grandchild) + .expect("grandchild"); + assert!(matches!( + guard.reserve_new(grandchild, None), + Err(LineageError::DepthExceeded { limit: 2 }) + )); + + let mut active_limited = SessionLineageGuard::from_live( + [session(root, None), session(child, Some(root))], + limits(4, 3, 1), + ) + .expect("valid graph"); + assert!(matches!( + active_limited.reserve_new(root, None), + Err(LineageError::ActiveDescendantsExceeded { limit: 1 }) + )); + } + + #[test] + fn reservation_release_is_exact_and_removal_releases_only_active_capacity() { + let root = id(1); + let child = id(2); + let mut guard = SessionLineageGuard::from_live([session(root, None)], limits(4, 1, 1)) + .expect("valid root"); + let reservation = guard.reserve_new(root, None).expect("reserve"); + assert_eq!( + guard.descendant_counts(root).expect("counts"), + DescendantCounts { + total: 0, + active: 0, + reserved: 1, + } + ); + guard.release(reservation).expect("release"); + assert_eq!( + guard.descendant_counts(root).expect("counts"), + DescendantCounts { + total: 0, + active: 0, + reserved: 0, + } + ); + + let reservation = guard.reserve_new(root, None).expect("reserve again"); + guard.commit_new(reservation, child).expect("commit"); + guard.remove_runtime(child).expect("remove"); + assert_eq!( + guard.descendant_counts(root).expect("counts"), + DescendantCounts { + total: 1, + active: 0, + reserved: 0, + } + ); + assert!(matches!( + guard.reserve_new(root, None), + Err(LineageError::TotalDescendantsExceeded { limit: 1 }) + )); + } + + #[test] + fn failed_commit_consumes_its_reservation() { + let root = id(1); + let mut guard = SessionLineageGuard::from_live([session(root, None)], limits(4, 1, 1)) + .expect("valid root"); + let reservation = guard.reserve_new(root, None).expect("reserve"); + assert!(matches!( + guard.commit_new(reservation, root), + Err(LineageError::DuplicateSession(_)) + )); + assert_eq!(guard.descendant_counts(root).expect("counts").reserved, 0); + } +} From ee576359be1259e55d83fcabc958fcfeb41963ff Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 3 Aug 2026 17:57:23 -0400 Subject: [PATCH 05/24] feat(lua): add scoped plugin session state runtime --- n00n-agent/src/agent/run.rs | 20 +- n00n-agent/src/headless.rs | 8 +- n00n-agent/src/tools/mod.rs | 73 +- n00n-lua/src/api/agent.rs | 23 +- n00n-lua/src/api/tool.rs | 33 +- n00n-lua/src/api/util/ctx.rs | 200 +++++- n00n-lua/src/api/util/mod.rs | 1 + n00n-lua/src/api/util/state_convert.rs | 495 ++++++++++++++ n00n-lua/src/error.rs | 2 + n00n-lua/src/lib.rs | 1 + n00n-lua/src/loader.rs | 75 ++- n00n-lua/src/runtime.rs | 525 ++++++++++++--- n00n-lua/src/state.rs | 888 +++++++++++++++++++++++++ n00n-lua/tests/plugin_host.rs | 426 +++++++++++- n00n-storage/src/sessions.rs | 566 ++++++++++++++-- n00n-ui/src/agent/agent_loop.rs | 9 +- 16 files changed, 3159 insertions(+), 186 deletions(-) create mode 100644 n00n-lua/src/api/util/state_convert.rs create mode 100644 n00n-lua/src/state.rs diff --git a/n00n-agent/src/agent/run.rs b/n00n-agent/src/agent/run.rs index 058e5870c..a98ea1d71 100644 --- a/n00n-agent/src/agent/run.rs +++ b/n00n-agent/src/agent/run.rs @@ -19,8 +19,8 @@ use crate::cancel::{CancelMap, CancelToken, PreDispatchGate}; use crate::mcp::McpSession; use crate::permissions::{PermissionAnswer, PermissionManager}; use crate::tools::{ - ActiveTools, Deadline, FileReadTracker, LocalTools, ToolAudience, ToolContext, ToolFilter, - ToolRegistry, + ActiveTools, Deadline, FileReadTracker, LocalTools, SessionIdentity, ToolAudience, ToolContext, + ToolFilter, ToolRegistry, }; use crate::{ AgentConfig, AgentError, AgentEvent, AgentInput, AgentMode, EventSender, ExtractedCommand, @@ -28,6 +28,7 @@ use crate::{ InterruptSource, ToolDoneEvent, TurnCompleteEvent, }; use n00n_config::{ToolKey, ToolOutputLines}; +#[cfg(test)] use n00n_storage::id::SessionRef; use crate::tokenize::{ @@ -167,7 +168,7 @@ pub struct AgentParams { pub config: Arc, pub tool_output_lines: ToolOutputLines, pub permissions: Arc, - pub session_id: Option, + pub identity: Option, pub timeouts: n00n_providers::Timeouts, pub openai_options: OpenAiOptions, pub file_tracker: Arc, @@ -214,7 +215,7 @@ pub struct Agent<'h> { thinking_empty_retried: bool, permissions: Arc, opts: RequestOptions, - session_id: Option, + identity: Option, timeouts: n00n_providers::Timeouts, openai_options: OpenAiOptions, file_tracker: Arc, @@ -280,7 +281,7 @@ impl<'h> Agent<'h> { post_tool_empty_retried: false, thinking_empty_retried: false, opts: RequestOptions::default(), - session_id: params.session_id, + identity: params.identity, file_tracker: params.file_tracker, prompt_slots: params.prompt_slots, subagent_cancels: params.subagent_cancels, @@ -510,7 +511,7 @@ impl<'h> Agent<'h> { event_tx: &self.event_tx, cancel: &self.cancel, opts, - session_id: self.session_id.as_ref(), + session_id: self.identity.as_ref().map(SessionIdentity::session_id), }) .await } @@ -895,6 +896,7 @@ impl<'h> Agent<'h> { prompt_slots: Arc::clone(&self.prompt_slots), opts: self.opts.clone(), subagent_cancels: Arc::clone(&self.subagent_cancels), + identity: self.identity.clone(), registry: Arc::clone(&self.registry), admission_scope: Arc::clone(&self.admission_scope), workflow: self.workflow, @@ -1154,7 +1156,7 @@ impl<'h> Agent<'h> { &self.event_tx, &self.cancel, CompactionTrigger::Auto, - self.session_id.as_ref(), + self.identity.as_ref().map(SessionIdentity::session_id), &cwd, None, ) @@ -2151,7 +2153,7 @@ mod tests { }, std::path::PathBuf::from("/tmp"), )), - session_id: None, + identity: None, timeouts: n00n_providers::Timeouts::default(), openai_options: OpenAiOptions::default(), file_tracker: FileReadTracker::fresh(), @@ -3376,7 +3378,7 @@ mod tests { }, std::path::PathBuf::from("/tmp"), )), - session_id: None, + identity: None, timeouts: n00n_providers::Timeouts::default(), openai_options: OpenAiOptions::default(), file_tracker: FileReadTracker::fresh(), diff --git a/n00n-agent/src/headless.rs b/n00n-agent/src/headless.rs index b8a5edbd7..862a8e591 100644 --- a/n00n-agent/src/headless.rs +++ b/n00n-agent/src/headless.rs @@ -21,7 +21,9 @@ use crate::cancel::{CancelMap, CancelToken}; use crate::permissions::PermissionManager; use crate::prompt::ResolvedSlots; use crate::template; -use crate::tools::{DescriptionContext, FileReadTracker, ToolAudience, ToolFilter, ToolRegistry}; +use crate::tools::{ + DescriptionContext, FileReadTracker, SessionIdentity, ToolAudience, ToolFilter, ToolRegistry, +}; use crate::{ Agent, AgentConfig, AgentEvent, AgentInput, AgentMode, AgentParams, AgentRunParams, Envelope, EventSender, ImageSource, McpHandle, McpSession, PermissionsConfig, ToolOutput, @@ -263,7 +265,7 @@ pub fn spawn(params: HeadlessParams) -> HeadlessHandle { params.permissions_config, working_dir_path, )), - session_id: Some(session_ref_clone.clone()), + identity: Some(SessionIdentity::root(session_ref_clone.clone())), timeouts: params.timeouts, openai_options: params.openai_options, file_tracker: FileReadTracker::fresh(), @@ -518,7 +520,7 @@ pub fn spawn_interactive(params: InteractiveParams) -> InteractiveHandle { config: Arc::clone(¶ms.config), tool_output_lines: ToolOutputLines::default(), permissions: Arc::clone(&permissions), - session_id: Some(session_ref_clone.clone()), + identity: Some(SessionIdentity::root(session_ref_clone.clone())), timeouts: params.timeouts, openai_options: params.openai_options, file_tracker: Arc::clone(&file_tracker), diff --git a/n00n-agent/src/tools/mod.rs b/n00n-agent/src/tools/mod.rs index c7e57d1b6..3c53d2c5a 100644 --- a/n00n-agent/src/tools/mod.rs +++ b/n00n-agent/src/tools/mod.rs @@ -302,6 +302,43 @@ pub fn timeout_annotation(secs: u64) -> String { pub type LocalToolFn = Arc Result + Send + Sync>; pub type LocalTools = Arc>; +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SessionIdentity { + session_id: n00n_storage::id::SessionRef, + root_session_id: n00n_storage::id::SessionRef, +} + +impl SessionIdentity { + #[must_use] + pub fn root(session_id: n00n_storage::id::SessionRef) -> Self { + Self { + root_session_id: session_id.clone(), + session_id, + } + } + + #[must_use] + pub fn child( + session_id: n00n_storage::id::SessionRef, + root_session_id: n00n_storage::id::SessionRef, + ) -> Self { + Self { + session_id, + root_session_id, + } + } + + #[must_use] + pub fn session_id(&self) -> &n00n_storage::id::SessionRef { + &self.session_id + } + + #[must_use] + pub fn root_session_id(&self) -> &n00n_storage::id::SessionRef { + &self.root_session_id + } +} + #[derive(Clone)] pub struct ToolContext { pub provider: Arc, @@ -324,6 +361,8 @@ pub struct ToolContext { pub prompt_slots: Arc, pub opts: RequestOptions, pub subagent_cancels: Arc>, + /// Immutable session and root identity of the agent executing this tool. + pub identity: Option, pub registry: Arc, /// Stable identity used by registry-scoped per-agent admission. Child /// sessions receive their own scope, while nested calls in one agent share it. @@ -618,6 +657,7 @@ pub fn interpreter_ctx( prompt_slots: Arc::new(crate::prompt::ResolvedSlots::default()), opts: RequestOptions::default(), subagent_cancels: Arc::new(CancelMap::new()), + identity: None, registry, admission_scope: crate::tools::ToolAdmission::new_scope(), tool_filter: ToolFilter::All, @@ -660,7 +700,8 @@ pub mod test_support { use super::{ AgentMode, Arc, CancelToken, DescriptionContext, FileReadTracker, LazyLock, - PermissionManager, ToolContext, ToolRegistry, Value, interpreter_ctx, registry, + PermissionManager, SessionIdentity, ToolContext, ToolRegistry, Value, interpreter_ctx, + registry, }; pub const GUARDED_TOOL_NAME: &str = "guarded_mock"; @@ -736,6 +777,9 @@ pub mod test_support { Arc::new(ToolRegistry::new()), ); ctx.tool_use_id = tool_use_id.map(String::from); + ctx.identity = Some(SessionIdentity::root( + n00n_storage::id::SessionRef::generate(), + )); ctx } @@ -776,6 +820,33 @@ mod tests { const LINE_LIMIT: usize = 500; + #[test] + fn root_session_identity_uses_one_id() { + let session_id = n00n_storage::id::SessionRef::generate(); + let identity = SessionIdentity::root(session_id.clone()); + + assert_eq!(identity.session_id(), &session_id); + assert_eq!(identity.root_session_id(), &session_id); + } + + #[test] + fn descendant_session_identities_inherit_root_and_remain_distinct() { + let root = SessionIdentity::root(n00n_storage::id::SessionRef::generate()); + let child = SessionIdentity::child( + n00n_storage::id::SessionRef::generate(), + root.root_session_id().clone(), + ); + let grandchild = SessionIdentity::child( + n00n_storage::id::SessionRef::generate(), + child.root_session_id().clone(), + ); + + assert_ne!(child.session_id(), root.session_id()); + assert_ne!(grandchild.session_id(), child.session_id()); + assert_eq!(child.root_session_id(), root.root_session_id()); + assert_eq!(grandchild.root_session_id(), root.root_session_id()); + } + #[test_case(true ; "vision_model_keeps_view_image")] #[test_case(false ; "text_only_model_loses_view_image")] fn from_config_gates_view_image_on_vision(vision: bool) { diff --git a/n00n-lua/src/api/agent.rs b/n00n-lua/src/api/agent.rs index 063a07707..bceadf932 100644 --- a/n00n-lua/src/api/agent.rs +++ b/n00n-lua/src/api/agent.rs @@ -17,8 +17,8 @@ use n00n_agent::cancel::CancelMap; use n00n_agent::tools::interpreter_bridge; use n00n_agent::tools::registry::ToolRegistry; use n00n_agent::tools::{ - Deadline, DescriptionContext, FileReadTracker, LocalToolFn, LocalTools, ToolAudience, - ToolContext, ToolFilter, ToolLive, + Deadline, DescriptionContext, FileReadTracker, LocalToolFn, LocalTools, SessionIdentity, + ToolAudience, ToolContext, ToolFilter, ToolLive, }; use n00n_agent::{ Agent, AgentEvent, AgentInput, AgentMode, AgentParams, AgentRunParams, Envelope, EventSender, @@ -36,6 +36,7 @@ use tracing::info; use crate::api::ui::buf::BufHandle; use crate::api::util::convert::{JSON_ARRAY_META_FIELD, json_to_lua, lua_to_json, lua_tool_result}; use crate::api::util::ctx::{AgentContext, LuaCtx}; +use crate::state::PluginStateStore; const SESSION_CLOSED_ERR: &str = "session closed"; const DEFAULT_SESSION_AUDIENCE: ToolAudience = ToolAudience::GENERAL_SUB; @@ -95,8 +96,7 @@ fn model_to_lua_table(lua: &Lua, model: &Model) -> LuaResult
{ } fn dispatch_ctx<'a>(ctx: &'a LuaCtx, method: &str) -> Result<&'a AgentContext, String> { - ctx.agent() - .ok_or_else(|| ctx.cap_err(&format!("n00n.agent.{method}"))) + ctx.dispatch_agent(method) } fn parse_session_mode( @@ -589,6 +589,7 @@ async fn call_tool( on_buf, on_ann, }; + let _nested_dispatch = crate::api::tool::enter_nested_dispatch(); let done = dispatch_racing_live(&tctx, &name, &input_json, rx, &cbs).await; // Same fallback the UI applies on tool completion, so a batch child's // header carries the annotation its standalone run would get. @@ -654,6 +655,10 @@ async fn session( opts: Table, ) -> LuaResult> { let agent_ctx = try_pair!(dispatch_ctx(&ctx, "session")).clone(); + let Some(parent_identity) = agent_ctx.identity.clone() else { + return Ok(err_pair("session identity is unavailable")); + }; + let plugin_state_store = try_pair!(ctx.plugin_state_store()); drop(ctx); let model_spec: Option = opts.get("model_spec")?; let system: Option = opts.get("system")?; @@ -853,7 +858,10 @@ async fn session( config: session_config(&agent_ctx.config, excluded_tools.clone()), tool_output_lines: n00n_config::ToolOutputLines::default(), permissions: Arc::clone(&agent_ctx.permissions), - session_id: Some(session_id.into()), + identity: Some(SessionIdentity::child( + session_id.into(), + parent_identity.root_session_id().clone(), + )), timeouts: agent_ctx.timeouts, openai_options: agent_ctx.openai_options, file_tracker: FileReadTracker::fresh(), @@ -885,6 +893,8 @@ async fn session( prompt_rx, prompt_tx: Some(prompt_tx), parent_cancels: Arc::clone(&agent_ctx.subagent_cancels), + plugin_state_store, + child_state_owner: session_id, child_id, parent_tool_use_id, parent_event_tx: parent_tx, @@ -1386,6 +1396,8 @@ struct SessionState { prompt_rx: flume::Receiver, prompt_tx: Option>, parent_cancels: Arc>, + plugin_state_store: Arc, + child_state_owner: n00nId, child_id: String, parent_tool_use_id: String, parent_event_tx: EventSender, @@ -1409,6 +1421,7 @@ impl SessionState { self.closed = true; self.progress.set_current_done(); self.parent_cancels.remove(&self.child_id); + self.plugin_state_store.drop_owner(self.child_state_owner); let messages = std::mem::replace(&mut self.history, History::new(Vec::new())).into_vec(); let _ = self.parent_event_tx.send(AgentEvent::SubagentHistory { tool_use_id: self.child_id.clone(), diff --git a/n00n-lua/src/api/tool.rs b/n00n-lua/src/api/tool.rs index 2735b20c6..fb5b752df 100644 --- a/n00n-lua/src/api/tool.rs +++ b/n00n-lua/src/api/tool.rs @@ -1,7 +1,7 @@ #![allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] use std::borrow::Cow; -use std::cell::RefCell; +use std::cell::{Cell, RefCell}; use std::path::Path; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -71,6 +71,24 @@ type ToolHandlesFn = Box Option>; thread_local! { static LOCAL_TOOL_HANDLES: RefCell> = const { RefCell::new(None) }; + static NESTED_DISPATCH_DEPTH: Cell = const { Cell::new(0) }; +} + +pub(crate) struct NestedDispatchGuard; + +pub(crate) fn enter_nested_dispatch() -> NestedDispatchGuard { + NESTED_DISPATCH_DEPTH.with(|depth| depth.set(depth.get().saturating_add(1))); + NestedDispatchGuard +} + +fn nested_dispatch_active() -> bool { + NESTED_DISPATCH_DEPTH.with(|depth| depth.get() > 0) +} + +impl Drop for NestedDispatchGuard { + fn drop(&mut self) { + NESTED_DISPATCH_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1))); + } } pub(crate) fn set_local_tool_handles(f: impl Fn(&str) -> Option + 'static) { @@ -254,6 +272,7 @@ impl Tool for LuaTool { input: validated, tx: self.tx.clone(), permission_state, + nested: nested_dispatch_active(), mutable_path_field: self.mutable_path_field.clone(), timeout: self.timeout, start_annotation: self.start_annotation.clone(), @@ -274,6 +293,7 @@ struct LuaToolInvocation { input: Value, tx: Sender, permission_state: PermissionState, + nested: bool, mutable_path_field: Option>, timeout: Option, start_annotation: Option, @@ -289,6 +309,7 @@ impl ToolInvocation for LuaToolInvocation { let plugin = Arc::clone(&self.plugin); let input = self.input.clone(); let tx = self.tx.clone(); + let nested = self.nested; let fallback = tool.to_string(); HeaderFuture::Pending { fallback: fallback.clone(), @@ -298,6 +319,7 @@ impl ToolInvocation for LuaToolInvocation { plugin: Arc::clone(&plugin), tool: Arc::clone(&tool), input, + nested, reply: reply_tx, }) .await; @@ -347,6 +369,7 @@ impl ToolInvocation for LuaToolInvocation { tool_use_id: id.clone(), }, ctx: Box::new(LuaCtx::start(ctx)), + nested: self.nested, reply: reply_tx, }; let tx = self.tx.clone(); @@ -366,6 +389,7 @@ impl ToolInvocation for LuaToolInvocation { let plugin = Arc::clone(&self.plugin); let tool = Arc::clone(&self.tool); let input = self.input.clone(); + let nested = self.nested; let fallback = input.to_string(); Box::pin(async move { if tx @@ -373,6 +397,7 @@ impl ToolInvocation for LuaToolInvocation { plugin, tool, input, + nested, reply: reply_tx, }) .await @@ -403,6 +428,7 @@ impl ToolInvocation for LuaToolInvocation { let input = self.input; let tx = self.tx; let tool_timeout = self.timeout; + let nested = self.nested; Box::pin(async move { let effective_secs: Option = match tool_timeout { @@ -436,6 +462,7 @@ impl ToolInvocation for LuaToolInvocation { Deadline::At(t) => Some(t), Deadline::None => None, }, + nested, reply: reply_tx, live, }) @@ -1688,6 +1715,7 @@ mod tests { input, tx, permission_state: PermissionState::Ready(None), + nested: false, mutable_path_field: None, timeout: Some(Duration::from_mins(1)), start_annotation: None, @@ -1848,6 +1876,7 @@ mod tests { input: serde_json::json!({"command": "ls"}), tx, permission_state: PermissionState::NeedsCompute, + nested: false, mutable_path_field: None, timeout: None, start_annotation: None, @@ -1866,6 +1895,7 @@ mod tests { input: serde_json::json!({"command": "echo hi"}), tx: tx2, permission_state: PermissionState::NeedsCompute, + nested: false, mutable_path_field: None, timeout: None, start_annotation: None, @@ -1890,6 +1920,7 @@ mod tests { input: serde_json::json!({"command": "cargo test"}), tx, permission_state: PermissionState::NeedsCompute, + nested: false, mutable_path_field: None, timeout: None, start_annotation: None, diff --git a/n00n-lua/src/api/util/ctx.rs b/n00n-lua/src/api/util/ctx.rs index 92b23ed9d..f6d168545 100644 --- a/n00n-lua/src/api/util/ctx.rs +++ b/n00n-lua/src/api/util/ctx.rs @@ -1,6 +1,9 @@ use std::ops::Deref; use std::path::{Path, PathBuf}; -use std::sync::Arc; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; use std::time::{Duration, Instant}; use mlua::{LuaSerdeExt, MultiValue, UserData, UserDataMethods, Value as LuaValue}; @@ -15,8 +18,13 @@ use n00n_storage::id::SessionRef; use crate::api::tool::ToolCallReply; use crate::api::ui::buf::BufHandle; use crate::api::util::convert::json_to_lua; +use crate::api::util::state_convert::{ + json_to_lua as state_json_to_lua, lua_to_json as state_lua_to_json, +}; use crate::runtime::{active_task, lock_cell}; +use crate::state::{PluginStateIdentity, PluginStateScope, PluginStateStore}; +const CONTEXT_INACTIVE_MSG: &str = "state context is no longer active"; const DEADLINE_ALREADY_SET_MSG: &str = "ctx:set_deadline() already called"; fn send_live_buf(lua: &mlua::Lua, buf: &mlua::AnyUserData) -> mlua::Result<()> { @@ -83,6 +91,14 @@ pub(crate) struct LuaCtx { pub(crate) cancel: CancelToken, tool_output_lines: ToolOutputLines, pub(crate) finish_tx: Option>, + active: Arc, + plugin_state: Option, +} + +struct PluginStateAccess { + plugin: Arc, + identity: PluginStateIdentity, + store: Arc, } enum Caps { @@ -112,6 +128,8 @@ impl LuaCtx { cancel: ctx.cancel.clone(), tool_output_lines: ctx.tool_output_lines, finish_tx: None, + active: Arc::new(AtomicBool::new(true)), + plugin_state: None, } } @@ -146,6 +164,8 @@ impl LuaCtx { cancel: CancelToken::none(), tool_output_lines, finish_tx: None, + active: Arc::new(AtomicBool::new(true)), + plugin_state: None, } } @@ -218,6 +238,62 @@ impl LuaCtx { } } + pub(crate) fn attach_plugin_state(&mut self, plugin: Arc, store: Arc) { + self.plugin_state = self + .agent() + .and_then(|agent| agent.identity.as_ref()) + .map(PluginStateIdentity::from) + .map(|identity| PluginStateAccess { + plugin, + identity, + store, + }); + } + + pub(crate) fn context_liveness(&self) -> Arc { + Arc::clone(&self.active) + } + + fn ensure_active(&self) -> Result<(), String> { + if self.active.load(Ordering::Acquire) { + Ok(()) + } else { + Err(CONTEXT_INACTIVE_MSG.to_owned()) + } + } + + fn inactive_pair(&self) -> Option<(LuaValue, Option)> { + self.ensure_active() + .err() + .map(|error| (LuaValue::Nil, Some(error))) + } + + pub(crate) fn plugin_state_store(&self) -> Result, String> { + self.plugin_state("n00n.agent.session") + .map(|access| Arc::clone(&access.store)) + } + + pub(crate) fn dispatch_agent(&self, method: &str) -> Result<&AgentContext, String> { + let agent = self + .agent() + .ok_or_else(|| self.cap_err(&format!("n00n.agent.{method}")))?; + self.ensure_active()?; + Ok(agent) + } + + fn plugin_state(&self, method: &str) -> Result<&PluginStateAccess, String> { + if !matches!(self.caps, Caps::Handler { .. }) { + return Err(self.cap_err(method)); + } + + let access = self + .plugin_state + .as_ref() + .ok_or_else(|| "session identity is unavailable".to_owned())?; + self.ensure_active()?; + Ok(access) + } + pub(crate) fn cap_err(&self, method: &str) -> String { format!("{method} not available in {} ctx", self.kind()) } @@ -230,9 +306,15 @@ impl LuaCtx { #[allow(clippy::too_many_lines)] impl UserData for LuaCtx { fn add_methods>(methods: &mut M) { - methods.add_method("cancelled", |_, this, ()| Ok(this.cancel.is_cancelled())); + methods.add_method("cancelled", |_, this, ()| { + this.ensure_active().map_err(mlua::Error::runtime)?; + Ok(this.cancel.is_cancelled()) + }); methods.add_method("workflow", |_, this, ()| { + if let Some(error) = this.inactive_pair() { + return Ok(error); + } let Some(workflow) = this.workflow() else { return Ok(this.cap_err_pair("workflow")); }; @@ -241,6 +323,9 @@ impl UserData for LuaCtx { methods.add_method("audience", |lua, this, ()| { const DEFAULT_AUDIENCE: &str = "main"; + if let Some(error) = this.inactive_pair() { + return Ok(error); + } let Some(audience) = this.audience() else { return Ok(this.cap_err_pair("audience")); }; @@ -249,6 +334,9 @@ impl UserData for LuaCtx { }); methods.add_method("live_buf", |lua, this, buf: mlua::AnyUserData| { + if let Some(error) = this.inactive_pair() { + return Ok(error); + } if matches!(this.caps, Caps::Restore { .. }) { return Ok(this.cap_err_pair("live_buf")); } @@ -257,6 +345,9 @@ impl UserData for LuaCtx { }); methods.add_method("config", |lua, this, args: MultiValue| { + if let Some(error) = this.inactive_pair() { + return Ok(error); + } let Some(config) = this.config() else { return Ok(this.cap_err_pair("config")); }; @@ -281,15 +372,102 @@ impl UserData for LuaCtx { }); methods.add_method("tool_output_lines", |lua, this, ()| { + this.ensure_active().map_err(mlua::Error::runtime)?; lua.to_value(&this.tool_output_lines) }); - methods.add_method("state", |lua, this, ()| match this.state() { - Some(v) => json_to_lua(lua, v), - None => Ok(LuaValue::Nil), + methods.add_method("state", |lua, this, ()| { + this.ensure_active().map_err(mlua::Error::runtime)?; + match this.state() { + Some(v) => json_to_lua(lua, v), + None => Ok(LuaValue::Nil), + } + }); + methods.add_method("state_get", |lua, this, scope: String| { + let Some(scope) = PluginStateScope::parse(&scope) else { + return Ok(( + LuaValue::Nil, + Some("state scope must be 'session' or 'root'".to_owned()), + )); + }; + let access = match this.plugin_state("state_get") { + Ok(access) => access, + Err(error) => return Ok((LuaValue::Nil, Some(error))), + }; + let Some(value) = access.store.get(&access.plugin, scope, &access.identity) else { + return Ok((LuaValue::Nil, None)); + }; + match state_json_to_lua(lua, &value) { + Ok(value) => Ok((value, None)), + Err(error) => Ok((LuaValue::Nil, Some(error.to_string()))), + } + }); + + methods.add_method( + "state_replace", + |lua, this, (scope, value): (String, LuaValue)| { + let Some(scope) = PluginStateScope::parse(&scope) else { + return Ok(( + LuaValue::Nil, + Some("state scope must be 'session' or 'root'".to_owned()), + )); + }; + let access = match this.plugin_state("state_replace") { + Ok(access) => access, + Err(error) => return Ok((LuaValue::Nil, Some(error))), + }; + let value = match state_lua_to_json(lua, &value) { + Ok(value) => value, + Err(error) => return Ok((LuaValue::Nil, Some(error.to_string()))), + }; + let previous = access.store.get(&access.plugin, scope, &access.identity); + let previous = match previous.as_ref() { + Some(previous) => match state_json_to_lua(lua, previous) { + Ok(previous) => previous, + Err(error) => return Ok((LuaValue::Nil, Some(error.to_string()))), + }, + None => LuaValue::Nil, + }; + if let Err(error) = + access + .store + .replace(&access.plugin, scope, &access.identity, value) + { + return Ok((LuaValue::Nil, Some(error.to_string()))); + } + Ok((previous, None)) + }, + ); + + methods.add_method("state_remove", |lua, this, scope: String| { + let Some(scope) = PluginStateScope::parse(&scope) else { + return Ok(( + LuaValue::Nil, + Some("state scope must be 'session' or 'root'".to_owned()), + )); + }; + let access = match this.plugin_state("state_remove") { + Ok(access) => access, + Err(error) => return Ok((LuaValue::Nil, Some(error))), + }; + let previous = access.store.get(&access.plugin, scope, &access.identity); + let previous = match previous.as_ref() { + Some(previous) => match state_json_to_lua(lua, previous) { + Ok(previous) => previous, + Err(error) => return Ok((LuaValue::Nil, Some(error.to_string()))), + }, + None => LuaValue::Nil, + }; + if let Err(error) = access.store.remove(&access.plugin, scope, &access.identity) { + return Ok((LuaValue::Nil, Some(error.to_string()))); + } + Ok((previous, None)) }); methods.add_method("set_deadline", |lua, this, secs: u64| { + if let Some(error) = this.inactive_pair() { + return Ok(error); + } if !matches!(this.caps, Caps::Handler { .. }) { return Ok(this.cap_err_pair("set_deadline")); } @@ -309,6 +487,9 @@ impl UserData for LuaCtx { }); methods.add_method("record_read", |_, this, path: String| { + if let Some(error) = this.inactive_pair() { + return Ok(error); + } let Some(tracker) = this.file_tracker() else { return Ok(this.cap_err_pair("record_read")); }; @@ -317,6 +498,9 @@ impl UserData for LuaCtx { }); methods.add_method("check_before_edit", |_, this, path: String| { + if let Some(error) = this.inactive_pair() { + return Ok(error); + } let Some(tracker) = this.file_tracker() else { return Ok(this.cap_err_pair("check_before_edit")); }; @@ -329,6 +513,9 @@ impl UserData for LuaCtx { methods.add_async_method( "find_instructions", |lua, this, dir_path: String| async move { + if let Some(error) = this.inactive_pair() { + return Ok(error); + } let Some(loaded) = this.loaded_instructions().cloned() else { return Ok(this.cap_err_pair("find_instructions")); }; @@ -358,6 +545,9 @@ impl UserData for LuaCtx { }); methods.add_method_mut("finish", |lua, this, val: LuaValue| { + if let Some(error) = this.inactive_pair() { + return Ok(error); + } if !matches!(this.caps, Caps::Handler { .. }) { return Ok(this.cap_err_pair("finish")); } diff --git a/n00n-lua/src/api/util/mod.rs b/n00n-lua/src/api/util/mod.rs index 7e911445b..d527a7395 100644 --- a/n00n-lua/src/api/util/mod.rs +++ b/n00n-lua/src/api/util/mod.rs @@ -3,3 +3,4 @@ pub(crate) mod convert; pub(crate) mod ctx; pub(crate) mod dispatch; pub(crate) mod setup; +pub(crate) mod state_convert; diff --git a/n00n-lua/src/api/util/state_convert.rs b/n00n-lua/src/api/util/state_convert.rs new file mode 100644 index 000000000..f124e39d5 --- /dev/null +++ b/n00n-lua/src/api/util/state_convert.rs @@ -0,0 +1,495 @@ +use std::{collections::HashSet, ffi::c_void}; + +use mlua::{Lua, LuaSerdeExt, Table, Value}; +use n00n_storage::sessions::MAX_PLUGIN_STATE_BYTES; +use serde_json::Value as JsonValue; + +pub(crate) const MAX_STATE_DEPTH: usize = 64; + +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub(crate) enum StateConvertError { + #[error("state contains unsupported Lua value type '{0}'")] + UnsupportedValue(&'static str), + #[error("state contains a non-finite number")] + NonFiniteNumber, + #[error("state contains a string that is not valid UTF-8")] + InvalidUtf8String, + #[error("state object keys must be UTF-8 strings")] + NonStringObjectKey, + #[error("state array keys must be positive integers")] + InvalidArrayKey, + #[error("state arrays must have contiguous keys starting at 1")] + SparseArray, + #[error("state contains a cycle")] + Cycle, + #[error("state exceeds the maximum size of {maximum} bytes")] + MaximumBytesExceeded { maximum: usize }, + #[error("state exceeds the maximum nesting depth of {maximum}")] + MaximumDepthExceeded { maximum: usize }, + #[error("JSON number cannot be represented as a Lua number")] + UnrepresentableNumber, + #[error("Lua operation failed during state conversion: {0}")] + Lua(String), +} + +struct StateBudget { + used: usize, +} + +impl StateBudget { + fn consume(&mut self, bytes: usize) -> Result<(), StateConvertError> { + self.used = self.used.saturating_add(bytes); + if self.used > MAX_PLUGIN_STATE_BYTES { + Err(StateConvertError::MaximumBytesExceeded { + maximum: MAX_PLUGIN_STATE_BYTES, + }) + } else { + Ok(()) + } + } +} + +impl From for StateConvertError { + fn from(error: mlua::Error) -> Self { + Self::Lua(error.to_string()) + } +} + +pub(crate) fn json_to_lua(lua: &Lua, value: &JsonValue) -> Result { + json_to_lua_at_depth(lua, value, 0) +} + +fn json_to_lua_at_depth( + lua: &Lua, + value: &JsonValue, + depth: usize, +) -> Result { + check_depth(depth)?; + match value { + JsonValue::Null => Ok(lua.null()), + JsonValue::Bool(value) => Ok(Value::Boolean(*value)), + JsonValue::Number(value) => { + if let Some(integer) = value.as_i64() { + return Ok(Value::Integer(integer)); + } + if value.as_u64().is_some() { + return Err(StateConvertError::UnrepresentableNumber); + } + let number = value + .as_f64() + .filter(|number| number.is_finite()) + .ok_or(StateConvertError::UnrepresentableNumber)?; + let round_trip = serde_json::Number::from_f64(number) + .ok_or(StateConvertError::UnrepresentableNumber)?; + if &round_trip != value { + return Err(StateConvertError::UnrepresentableNumber); + } + Ok(Value::Number(number)) + } + JsonValue::String(value) => Ok(Value::String(lua.create_string(value)?)), + JsonValue::Array(values) => { + let table = lua.create_table_with_capacity(values.len(), 0)?; + table.set_metatable(Some(lua.array_metatable()))?; + for (offset, value) in values.iter().enumerate() { + table.raw_set(offset + 1, json_to_lua_at_depth(lua, value, depth + 1)?)?; + } + Ok(Value::Table(table)) + } + JsonValue::Object(values) => { + let table = lua.create_table_with_capacity(0, values.len())?; + for (key, value) in values { + table.raw_set(key.as_str(), json_to_lua_at_depth(lua, value, depth + 1)?)?; + } + Ok(Value::Table(table)) + } + } +} + +pub(crate) fn lua_to_json(lua: &Lua, value: &Value) -> Result { + lua_to_json_at_depth( + lua, + value, + 0, + &mut HashSet::new(), + &mut StateBudget { used: 0 }, + ) +} + +fn lua_to_json_at_depth( + lua: &Lua, + value: &Value, + depth: usize, + active_tables: &mut HashSet<*const c_void>, + budget: &mut StateBudget, +) -> Result { + check_depth(depth)?; + match value { + value if value.is_null() => { + budget.consume(4)?; + Ok(JsonValue::Null) + } + Value::Boolean(value) => { + budget.consume(if *value { 4 } else { 5 })?; + Ok(JsonValue::Bool(*value)) + } + Value::Integer(value) => { + budget.consume(value.to_string().len())?; + Ok(JsonValue::Number((*value).into())) + } + Value::Number(value) => { + let number = + serde_json::Number::from_f64(*value).ok_or(StateConvertError::NonFiniteNumber)?; + budget.consume(number.to_string().len())?; + Ok(JsonValue::Number(number)) + } + Value::String(value) => { + let value = value + .to_str() + .map_err(|_| StateConvertError::InvalidUtf8String)?; + budget.consume(serialized_string_len(&value))?; + Ok(JsonValue::String(value.to_owned())) + } + Value::Table(table) => table_to_json(lua, table, depth, active_tables, budget), + value => Err(StateConvertError::UnsupportedValue(value.type_name())), + } +} + +fn table_to_json( + lua: &Lua, + table: &Table, + depth: usize, + active_tables: &mut HashSet<*const c_void>, + budget: &mut StateBudget, +) -> Result { + budget.consume(2)?; + let pointer = table.to_pointer(); + if !active_tables.insert(pointer) { + return Err(StateConvertError::Cycle); + } + + let result = if table + .metatable() + .is_some_and(|metatable| metatable == lua.array_metatable()) + { + array_to_json(lua, table, depth, active_tables, budget) + } else { + object_to_json(lua, table, depth, active_tables, budget) + }; + active_tables.remove(&pointer); + result +} + +fn array_to_json( + lua: &Lua, + table: &Table, + depth: usize, + active_tables: &mut HashSet<*const c_void>, + budget: &mut StateBudget, +) -> Result { + let mut entries = Vec::new(); + for pair in table.pairs::() { + let (key, value) = pair?; + let Value::Integer(key) = key else { + return Err(StateConvertError::InvalidArrayKey); + }; + let index = usize::try_from(key).map_err(|_| StateConvertError::InvalidArrayKey)?; + if index == 0 { + return Err(StateConvertError::InvalidArrayKey); + } + if !entries.is_empty() { + budget.consume(1)?; + } + entries.push((index, value)); + } + entries.sort_unstable_by_key(|(index, _)| *index); + + let mut values = Vec::with_capacity(entries.len()); + for (offset, (index, value)) in entries.into_iter().enumerate() { + if index != offset + 1 { + return Err(StateConvertError::SparseArray); + } + values.push(lua_to_json_at_depth( + lua, + &value, + depth + 1, + active_tables, + budget, + )?); + } + Ok(JsonValue::Array(values)) +} + +fn object_to_json( + lua: &Lua, + table: &Table, + depth: usize, + active_tables: &mut HashSet<*const c_void>, + budget: &mut StateBudget, +) -> Result { + let mut values = serde_json::Map::new(); + for pair in table.pairs::() { + let (key, value) = pair?; + let Value::String(key) = key else { + return Err(StateConvertError::NonStringObjectKey); + }; + let key = key + .to_str() + .map_err(|_| StateConvertError::InvalidUtf8String)?; + if !values.is_empty() { + budget.consume(1)?; + } + budget.consume(serialized_string_len(&key).saturating_add(1))?; + values.insert( + key.to_owned(), + lua_to_json_at_depth(lua, &value, depth + 1, active_tables, budget)?, + ); + } + Ok(JsonValue::Object(values)) +} +fn serialized_string_len(value: &str) -> usize { + value.chars().fold(2, |bytes, character| { + bytes.saturating_add(match character { + '"' | '\\' | '\u{0008}' | '\u{0009}' | '\u{000a}' | '\u{000c}' | '\u{000d}' => 2, + '\u{0000}'..='\u{001f}' => 6, + character => character.len_utf8(), + }) + }) +} + +const fn check_depth(depth: usize) -> Result<(), StateConvertError> { + if depth > MAX_STATE_DEPTH { + Err(StateConvertError::MaximumDepthExceeded { + maximum: MAX_STATE_DEPTH, + }) + } else { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use mlua::{Lua, LuaSerdeExt, Value}; + use n00n_storage::sessions::MAX_PLUGIN_STATE_BYTES; + use serde_json::json; + use test_case::test_case; + + use super::{ + MAX_STATE_DEPTH, StateConvertError, json_to_lua, lua_to_json, serialized_string_len, + }; + + #[test] + fn round_trip_preserves_nested_null_and_array_shape() { + let lua = Lua::new(); + let input = json!({"items": [null, {"enabled": true}], "name": "state"}); + + let lua_value = json_to_lua(&lua, &input).unwrap(); + let items: Value = lua_value.as_table().unwrap().raw_get("items").unwrap(); + let first: Value = items.as_table().unwrap().raw_get(1).unwrap(); + + assert!(first.is_null()); + assert_eq!(lua_to_json(&lua, &lua_value).unwrap(), input); + } + + #[test_case(f64::NAN ; "nan")] + #[test_case(f64::INFINITY ; "positive_infinity")] + #[test_case(f64::NEG_INFINITY ; "negative_infinity")] + fn rejects_non_finite_numbers(number: f64) { + let lua = Lua::new(); + + assert_eq!( + lua_to_json(&lua, &Value::Number(number)).unwrap_err(), + StateConvertError::NonFiniteNumber + ); + } + + #[test_case("18446744073709551617" ; "above_u64")] + #[test_case("-9223372036854775809" ; "below_i64")] + #[test_case("0.12345678901234567890123456789" ; "precision_loss")] + fn rejects_json_numbers_that_cannot_round_trip_exactly(raw: &str) { + let lua = Lua::new(); + let value: serde_json::Value = serde_json::from_str(raw).unwrap(); + + assert_eq!( + json_to_lua(&lua, &value).unwrap_err(), + StateConvertError::UnrepresentableNumber + ); + } + + #[test] + fn rejects_non_utf8_strings() { + let lua = Lua::new(); + let value = Value::String(lua.create_string([0xff]).unwrap()); + + assert_eq!( + lua_to_json(&lua, &value).unwrap_err(), + StateConvertError::InvalidUtf8String + ); + } + + #[test] + fn rejects_unsupported_values() { + let lua = Lua::new(); + let value = Value::Function(lua.create_function(|_, ()| Ok(())).unwrap()); + + assert_eq!( + lua_to_json(&lua, &value).unwrap_err(), + StateConvertError::UnsupportedValue("function") + ); + } + + #[test] + fn rejects_nil_but_accepts_mlua_null() { + let lua = Lua::new(); + + assert_eq!( + lua_to_json(&lua, &Value::Nil).unwrap_err(), + StateConvertError::UnsupportedValue("nil") + ); + assert_eq!(lua_to_json(&lua, &lua.null()).unwrap(), json!(null)); + } + + #[test] + fn unmarked_tables_are_objects_with_string_keys_only() { + let lua = Lua::new(); + let object = lua.create_table().unwrap(); + object.raw_set("name", "plugin").unwrap(); + assert_eq!( + lua_to_json(&lua, &Value::Table(object)).unwrap(), + json!({"name": "plugin"}) + ); + + let invalid = lua.create_table().unwrap(); + invalid.raw_set(1, "item").unwrap(); + assert_eq!( + lua_to_json(&lua, &Value::Table(invalid)).unwrap_err(), + StateConvertError::NonStringObjectKey + ); + } + + #[test] + fn marked_arrays_require_only_contiguous_integer_keys() { + let lua = Lua::new(); + let array = lua.create_table().unwrap(); + array.set_metatable(Some(lua.array_metatable())).unwrap(); + array.raw_set(1, "a").unwrap(); + array.raw_set(2, "b").unwrap(); + assert_eq!( + lua_to_json(&lua, &Value::Table(array)).unwrap(), + json!(["a", "b"]) + ); + + let mixed = lua.create_table().unwrap(); + mixed.set_metatable(Some(lua.array_metatable())).unwrap(); + mixed.raw_set(1, "a").unwrap(); + mixed.raw_set("name", "state").unwrap(); + assert_eq!( + lua_to_json(&lua, &Value::Table(mixed)).unwrap_err(), + StateConvertError::InvalidArrayKey + ); + + let sparse = lua.create_table().unwrap(); + sparse.set_metatable(Some(lua.array_metatable())).unwrap(); + sparse.raw_set(1, "a").unwrap(); + sparse.raw_set(3, "c").unwrap(); + assert_eq!( + lua_to_json(&lua, &Value::Table(sparse)).unwrap_err(), + StateConvertError::SparseArray + ); + } + + #[test] + fn rejects_direct_and_indirect_cycles() { + let lua = Lua::new(); + let direct = lua.create_table().unwrap(); + direct.raw_set("self", direct.clone()).unwrap(); + assert_eq!( + lua_to_json(&lua, &Value::Table(direct)).unwrap_err(), + StateConvertError::Cycle + ); + + let first = lua.create_table().unwrap(); + let second = lua.create_table().unwrap(); + first.raw_set("second", second.clone()).unwrap(); + second.raw_set("first", first.clone()).unwrap(); + assert_eq!( + lua_to_json(&lua, &Value::Table(first)).unwrap_err(), + StateConvertError::Cycle + ); + } + + #[test] + fn allows_repeated_acyclic_table_references() { + let lua = Lua::new(); + let shared = lua.create_table().unwrap(); + shared.raw_set("value", 7).unwrap(); + let root = lua.create_table().unwrap(); + root.raw_set("left", shared.clone()).unwrap(); + root.raw_set("right", shared).unwrap(); + + assert_eq!( + lua_to_json(&lua, &Value::Table(root)).unwrap(), + json!({"left": {"value": 7}, "right": {"value": 7}}) + ); + } + + #[test] + fn enforces_serialized_size_during_lua_conversion() { + let lua = Lua::new(); + let exact = Value::String( + lua.create_string("x".repeat(MAX_PLUGIN_STATE_BYTES - 2)) + .unwrap(), + ); + assert!(lua_to_json(&lua, &exact).is_ok()); + + let oversized = Value::String( + lua.create_string("x".repeat(MAX_PLUGIN_STATE_BYTES - 1)) + .unwrap(), + ); + assert_eq!( + lua_to_json(&lua, &oversized).unwrap_err(), + StateConvertError::MaximumBytesExceeded { + maximum: MAX_PLUGIN_STATE_BYTES + } + ); + } + + #[test_case("plain" ; "plain")] + #[test_case("quote\"slash\\" ; "escapes")] + #[test_case("line\ncontrol\u{0001}" ; "controls")] + #[test_case("héllo" ; "utf8")] + fn serialized_string_size_matches_serde_json(value: &str) { + assert_eq!( + serialized_string_len(value), + serde_json::to_vec(value).unwrap().len() + ); + } + + #[test] + fn enforces_maximum_depth_in_both_directions() { + let lua = Lua::new(); + let mut json_value = json!(true); + for _ in 0..=MAX_STATE_DEPTH { + json_value = json!({"child": json_value}); + } + assert_eq!( + json_to_lua(&lua, &json_value).unwrap_err(), + StateConvertError::MaximumDepthExceeded { + maximum: MAX_STATE_DEPTH + } + ); + + let root = lua.create_table().unwrap(); + let mut current = root.clone(); + for _ in 0..=MAX_STATE_DEPTH { + let child = lua.create_table().unwrap(); + current.raw_set("child", child.clone()).unwrap(); + current = child; + } + assert_eq!( + lua_to_json(&lua, &Value::Table(root)).unwrap_err(), + StateConvertError::MaximumDepthExceeded { + maximum: MAX_STATE_DEPTH + } + ); + } +} diff --git a/n00n-lua/src/error.rs b/n00n-lua/src/error.rs index f8a6f4abd..d04d94852 100644 --- a/n00n-lua/src/error.rs +++ b/n00n-lua/src/error.rs @@ -25,4 +25,6 @@ pub enum PluginError { UnknownPlugin { plugin: String }, #[error("plugin host is not running")] HostDead, + #[error("plugin state operation failed: {message}")] + State { message: String }, } diff --git a/n00n-lua/src/lib.rs b/n00n-lua/src/lib.rs index 23ef18d8a..b9e424fd5 100644 --- a/n00n-lua/src/lib.rs +++ b/n00n-lua/src/lib.rs @@ -10,6 +10,7 @@ pub mod language; mod loader; pub(crate) mod plugin_permissions; mod runtime; +mod state; pub use api::keymap::{KeymapEntry, KeymapReader, KeymapSnapshot}; pub use api::options::{OptionSpec, OptionType, PluginOptionSpecs}; diff --git a/n00n-lua/src/loader.rs b/n00n-lua/src/loader.rs index ecabdc2ed..211f96ca7 100644 --- a/n00n-lua/src/loader.rs +++ b/n00n-lua/src/loader.rs @@ -6,8 +6,10 @@ use std::sync::{Arc, LazyLock}; use std::time::Duration; use include_dir::{Dir, include_dir}; -use n00n_agent::tools::ToolRegistry; +use n00n_agent::tools::{SessionIdentity, ToolRegistry}; use n00n_config::{PluginsConfig, RawConfig}; +use n00n_storage::id::n00nId; +use n00n_storage::sessions::StoredSessionStateSnapshot; use crate::api::keymap::KeymapReader; use crate::api::options::{PluginOptionSpecs, PluginOpts}; @@ -15,6 +17,7 @@ use crate::api::util::command::{HintReader, LuaCommandReader, UiAction}; use crate::error::PluginError; use crate::plugin_permissions::{PluginPermissions, load_plugin_permissions}; use crate::runtime::{self, ClickFallback, LuaThread, Request, RestoreItem}; +use crate::state::PluginStateIdentity; use n00n_agent::prompt::ResolvedSlots; const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2); @@ -605,6 +608,76 @@ impl EventHandle { .await .unwrap_or_else(|_| ResolvedSlots::default()) } + /// Hydrates host-owned plugin state after all in-flight Lua work drains. + /// + /// # Errors + /// Returns an error when the host is unavailable or the snapshot cannot be applied. + pub fn hydrate_state( + &self, + identity: &SessionIdentity, + snapshot: Option, + ) -> Result<(), PluginError> { + let (reply, recv) = flume::bounded(1); + self.tx + .send(Request::HydrateState { + identity: PluginStateIdentity::from(identity), + snapshot, + reply, + }) + .map_err(|_| PluginError::HostDead)?; + recv.recv() + .map_err(|_| PluginError::HostDead)? + .map_err(|message| PluginError::State { message }) + } + + /// Captures host-owned plugin state after all in-flight Lua work drains. + /// + /// # Errors + /// Returns an error when the host is unavailable or capture validation fails. + pub fn capture_state( + &self, + identity: &SessionIdentity, + revision: u64, + ) -> Result { + let (reply, recv) = flume::bounded(1); + self.tx + .send(Request::CaptureState { + identity: PluginStateIdentity::from(identity), + revision, + reply, + }) + .map_err(|_| PluginError::HostDead)?; + recv.recv() + .map_err(|_| PluginError::HostDead)? + .map_err(|message| PluginError::State { message }) + } + + /// Clears both scopes for an identity and records removals for the next capture. + /// + /// # Errors + /// Returns an error when the host is unavailable. + pub fn reset_state(&self, identity: &SessionIdentity) -> Result<(), PluginError> { + let (reply, recv) = flume::bounded(1); + self.tx + .send(Request::ResetState { + identity: PluginStateIdentity::from(identity), + reply, + }) + .map_err(|_| PluginError::HostDead)?; + recv.recv().map_err(|_| PluginError::HostDead) + } + + /// Drops in-memory state for one canonical owner without persisting removals. + /// + /// # Errors + /// Returns an error when the host is unavailable. + pub fn drop_state_owner(&self, owner: n00nId) -> Result<(), PluginError> { + let (reply, recv) = flume::bounded(1); + self.tx + .send(Request::DropStateOwner { owner, reply }) + .map_err(|_| PluginError::HostDead)?; + recv.recv().map_err(|_| PluginError::HostDead) + } pub fn request_restore(&self, item: RestoreItem, event_tx: n00n_agent::EventSender) { let _ = self.tx.send(Request::RestoreToolAsync { item, event_tx }); diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index ac8368a05..4668c47d5 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -39,10 +39,15 @@ use crate::api::util::command::{LuaCommandReader, LuaCommandWriter, UiAction}; use crate::api::util::convert::json_to_lua; use crate::api::util::ctx::LuaCtx; use crate::api::util::setup::ConfigStore; +use crate::api::util::state_convert::json_to_lua as state_json_to_lua; use crate::docs_render; use crate::error::PluginError; use crate::plugin_permissions::{PluginPermissions, load_plugin_permissions}; +use n00n_storage::id::n00nId; +use n00n_storage::sessions::{StoredSessionStateSnapshot, StoredStateScope}; + +use crate::state::{PLUGIN_STATE_SCHEMA_VERSION, PluginStateIdentity, PluginStateStore}; fn register_builtin_tools(registry: &Arc) -> Result<(), PluginError> { let tools: [(Arc, ToolSource); 2] = [ ( @@ -140,6 +145,7 @@ pub enum Request { input: Value, ctx: Box, deadline: Option, + nested: bool, reply: flume::Sender, live: Option, }, @@ -147,12 +153,14 @@ pub enum Request { plugin: Arc, tool: Arc, input: Value, + nested: bool, reply: flume::Sender, }, ComputePermissionScopes { plugin: Arc, tool: Arc, input: Value, + nested: bool, reply: flume::Sender>, }, ClearPlugin { @@ -176,6 +184,24 @@ pub enum Request { CollectPluginOptions { reply: flume::Sender, }, + HydrateState { + identity: PluginStateIdentity, + snapshot: Option, + reply: flume::Sender>, + }, + CaptureState { + identity: PluginStateIdentity, + revision: u64, + reply: flume::Sender>, + }, + ResetState { + identity: PluginStateIdentity, + reply: flume::Sender<()>, + }, + DropStateOwner { + owner: n00nId, + reply: flume::Sender<()>, + }, Shutdown, RestoreToolAsync { item: RestoreItem, @@ -221,6 +247,7 @@ pub enum Request { input: Value, live: LiveCtx, ctx: Box, + nested: bool, reply: flume::Sender<()>, }, } @@ -282,6 +309,13 @@ pub struct LiveCtx { pub event_tx: n00n_agent::EventSender, pub tool_use_id: String, } +struct ContextLivenessGuard(Arc); + +impl Drop for ContextLivenessGuard { + fn drop(&mut self) { + self.0.store(false, Ordering::Release); + } +} /// Lua is single-threaded so this Mutex never contends, but /// `Lua::app_data` requires `Send + Sync` with the `send` feature. @@ -672,6 +706,40 @@ pub(crate) fn enqueue_async_task(lua: &Lua, work_fn: RegistryKey) -> Result<(), /// Caps concurrent coroutines to avoid blowing the Lua stack. /// Also serves as a drain barrier for load/clear ops. +#[derive(Default)] +struct LifecycleGate { + count: Cell, + event: Event, +} + +impl LifecycleGate { + fn start(self: &Rc) -> LifecycleGuard { + self.count.set(self.count.get() + 1); + LifecycleGuard(Rc::clone(self)) + } + + fn is_idle(&self) -> bool { + self.count.get() == 0 + } + + async fn changed(&self) { + let listener = self.event.listen(); + if self.is_idle() { + return; + } + listener.await; + } +} + +struct LifecycleGuard(Rc); + +impl Drop for LifecycleGuard { + fn drop(&mut self) { + self.0.count.set(self.0.count.get().saturating_sub(1)); + self.0.event.notify(usize::MAX); + } +} + struct InflightGate { lua: Lua, count: Cell, @@ -962,10 +1030,216 @@ async fn drain_barrier( } } +fn spawn_runtime_request( + rt: &LuaRuntime, + ex: &Rc>, + gate: &Rc, + lifecycle: &Rc, + request: Request, + nested_only: bool, +) -> Option { + if nested_only + && !matches!( + &request, + Request::CallTool { nested: true, .. } + | Request::ComputeHeader { nested: true, .. } + | Request::ComputePermissionScopes { nested: true, .. } + | Request::StartTool { nested: true, .. } + ) + { + return Some(request); + } + match request { + Request::CallTool { + plugin, + tool, + input, + mut ctx, + deadline, + nested: _, + reply, + live, + } => { + ctx.attach_plugin_state(Arc::clone(&plugin), Arc::clone(&rt.state)); + let lua = rt.lua.clone(); + let plugins = Rc::clone(&rt.plugins); + let live_tasks = Rc::clone(&rt.live_tasks); + let warm_tools = Rc::clone(&rt.warm_tools); + let shutdown = Arc::clone(&rt.shutdown); + let gate = Rc::clone(gate); + let lifecycle = lifecycle.start(); + ex.spawn(async move { + let result = run_tool_call( + lua, plugin, tool, input, ctx, deadline, live, live_tasks, warm_tools, plugins, + shutdown, gate, lifecycle, + ) + .await; + let _ = reply.send(result); + }) + .detach(); + None + } + Request::ComputeHeader { + plugin, + tool, + input, + nested: _, + reply, + } => { + let lua = rt.lua.clone(); + let plugins = Rc::clone(&rt.plugins); + let lifecycle = lifecycle.start(); + ex.spawn(async move { + let result = compute_header(&lua, &plugins, &plugin, &tool, input).await; + let _ = reply.send(result); + drop(lifecycle); + }) + .detach(); + None + } + Request::ComputePermissionScopes { + plugin, + tool, + input, + nested: _, + reply, + } => { + let lua = rt.lua.clone(); + let plugins = Rc::clone(&rt.plugins); + let lifecycle = lifecycle.start(); + ex.spawn(async move { + let result = + LuaRuntime::compute_permission_scopes(&lua, &plugins, &plugin, &tool, input) + .await; + let _ = reply.send(result); + drop(lifecycle); + }) + .detach(); + None + } + Request::StartTool { + plugin, + tool, + input, + live, + ctx, + nested: _, + reply, + } => { + let func = { + let plugins = rt.plugins.borrow(); + plugins + .get(&*plugin) + .and_then(|tools| tools.get(&*tool)) + .and_then(|keys| keys.start.as_ref()) + .and_then(|key| rt.lua.registry_value::(key).ok()) + }; + let Some(func) = func else { + let _ = reply.send(()); + return None; + }; + let lua = rt.lua.clone(); + let gate = Rc::clone(gate); + ex.spawn(async move { + let _gate_guard = gate.acquire().await; + run_tool_start(&lua, func, &tool, input, live, ctx).await; + let _ = reply.send(()); + }) + .detach(); + None + } + request => Some(request), + } +} + +enum RuntimeWake { + Lifecycle, + Spawn(PendingAsyncTask), + Request(Request), + Closed, +} + +async fn drain_runtime( + rt: &LuaRuntime, + ex: &Rc>, + gate: &Rc, + lifecycle: &Rc, + spawn_rx: &flume::Receiver, + request_rx: &flume::Receiver, + deferred: &mut VecDeque, +) { + while !lifecycle.is_idle() { + while let Ok(task) = spawn_rx.try_recv() { + spawn_async_task(&rt.lua, ex, gate, task); + } + while let Ok(request) = request_rx.try_recv() { + if let Some(request) = spawn_runtime_request(rt, ex, gate, lifecycle, request, true) { + deferred.push_back(request); + } + } + if lifecycle.is_idle() { + break; + } + let wake = smol::future::or( + async { + lifecycle.changed().await; + RuntimeWake::Lifecycle + }, + smol::future::or( + async { + spawn_rx + .recv_async() + .await + .map_or(RuntimeWake::Closed, RuntimeWake::Spawn) + }, + async { + request_rx + .recv_async() + .await + .map_or(RuntimeWake::Closed, RuntimeWake::Request) + }, + ), + ) + .await; + match wake { + RuntimeWake::Spawn(task) => spawn_async_task(&rt.lua, ex, gate, task), + RuntimeWake::Request(request) => { + if let Some(request) = spawn_runtime_request(rt, ex, gate, lifecycle, request, true) + { + deferred.push_back(request); + } + } + RuntimeWake::Lifecycle | RuntimeWake::Closed => {} + } + } + drain_barrier(&rt.lua, ex, gate, spawn_rx).await; +} + +fn validate_snapshot_lua_values( + lua: &Lua, + identity: &PluginStateIdentity, + snapshot: Option<&StoredSessionStateSnapshot>, +) -> Result<(), String> { + let Some(snapshot) = snapshot else { + return Ok(()); + }; + let entries = snapshot + .plugin_entries_for_apply(PLUGIN_STATE_SCHEMA_VERSION) + .map_err(|error| error.to_string())?; + for entry in entries { + if !identity.is_root() && matches!(entry.scope, StoredStateScope::Root) { + continue; + } + state_json_to_lua(lua, entry.payload).map_err(|error| error.to_string())?; + } + Ok(()) +} + struct ToolKeys { handler: RegistryKey, header: Option, restore: Option, + start: Option, permission_scopes: Option, describe: Option, @@ -983,6 +1257,7 @@ struct LuaRuntime { live_tasks: LiveTasks, warm_tools: WarmTools, registry: Arc, + state: Arc, tx: flume::Sender, shutdown: Arc, bundled_dirs: &'static [&'static Dir<'static>], @@ -1075,6 +1350,7 @@ impl LuaRuntime { live_tasks: Rc::new(RefCell::new(HashMap::new())), warm_tools: Rc::new(RefCell::new(VecDeque::new())), registry, + state: Arc::new(PluginStateStore::default()), tx, shutdown, bundled_dirs, @@ -1579,21 +1855,22 @@ impl LuaRuntime { } async fn compute_permission_scopes( - &self, + lua: &Lua, + plugins: &PluginMap, plugin: &str, tool: &str, input: Value, ) -> Option { let (func, lua_input) = plugin_fn( - &self.lua, - &self.plugins, + lua, + plugins, plugin, tool, "permission_scopes", |tk| tk.permission_scopes.as_ref(), &input, )?; - let result: LuaValue = match run_detached(&self.lua, func.call_async(lua_input)).await { + let result: LuaValue = match run_detached(lua, func.call_async(lua_input)).await { Ok(v) => v, Err(e) => { tracing::warn!(plugin, tool, error = %e, "permission_scopes callback failed"); @@ -1755,8 +2032,10 @@ async fn restore_item( None, ); + let ctx = LuaCtx::restore(item.tool_output_lines, item.state); + let _context_liveness = ContextLivenessGuard(ctx.context_liveness()); let ctx = lua - .create_userdata(LuaCtx::restore(item.tool_output_lines, item.state)) + .create_userdata(ctx) .map_err(|e| format!("restore context creation failed: {e}"))?; let inner = thread .into_async::((input_lua, &*item.output, item.is_error, ctx)) @@ -2037,6 +2316,7 @@ async fn run_tool_start( live: LiveCtx, ctx: Box, ) { + let _context_liveness = ContextLivenessGuard(ctx.context_liveness()); let session_id = ctx.session_id(); let scope = TaskScope::new( lua, @@ -2070,7 +2350,9 @@ async fn run_tool_call( plugins: PluginMap, shutdown: Arc, gate: Rc, + _lifecycle: LifecycleGuard, ) -> ToolCallReply { + let _context_liveness = ContextLivenessGuard(ctx.context_liveness()); let handler: Function = { let plugins_ref = plugins.borrow(); let Some(keys) = plugins_ref.get(&*plugin) else { @@ -2267,6 +2549,7 @@ pub fn spawn( let ex = Rc::new(smol::LocalExecutor::new()); let gate = Rc::new(InflightGate::new(rt.lua.clone())); + let lifecycle = Rc::new(LifecycleGate::default()); let restores = Rc::new(RestoreTracker::default()); let spawn_rx = rt .lua @@ -2276,6 +2559,7 @@ pub fn spawn( .clone(); smol::block_on(ex.run(async { + let mut deferred = VecDeque::new(); loop { while let Ok(task) = spawn_rx.try_recv() { spawn_async_task(&rt.lua, &ex, &gate, task); @@ -2284,25 +2568,29 @@ pub fn spawn( // ahead of bulk work like session restores so the UI stays // snappy, and queued `n00n.async.run` tasks jump ahead of // plain requests. - let next = smol::future::or( - async { prio_rx.recv_async().await.map(Some) }, - smol::future::or( - async { - let task = spawn_rx.recv_async().await?; - spawn_async_task(&rt.lua, &ex, &gate, task); - Ok(None) - }, - async { rx.recv_async().await.map(Some) }, - ), - ) - .await; - let msg = match next { - Ok(Some(m)) => m, - Ok(None) => { - smol::future::yield_now().await; - continue; + let msg = if let Some(request) = deferred.pop_front() { + request + } else { + let next = smol::future::or( + async { prio_rx.recv_async().await.map(Some) }, + smol::future::or( + async { + let task = spawn_rx.recv_async().await?; + spawn_async_task(&rt.lua, &ex, &gate, task); + Ok(None) + }, + async { rx.recv_async().await.map(Some) }, + ), + ) + .await; + match next { + Ok(Some(request)) => request, + Ok(None) => { + smol::future::yield_now().await; + continue; + } + Err(_) => break, } - Err(_) => break, }; match msg { Request::Shutdown => break, @@ -2314,48 +2602,37 @@ pub fn spawn( opts, reply, } => { - drain_barrier(&rt.lua, &ex, &gate, &spawn_rx).await; + drain_runtime( + &rt, + &ex, + &gate, + &lifecycle, + &spawn_rx, + &rx, + &mut deferred, + ) + .await; let res = rt.load_source(Arc::clone(&name), &source, plugin_dir, &permissions, opts, None).await; let _ = reply.send(res); } - Request::CallTool { - plugin, - tool, - input, - ctx, - deadline, - reply, - live, - } => { - let lua = rt.lua.clone(); - let plugins = Rc::clone(&rt.plugins); - let live_tasks = Rc::clone(&rt.live_tasks); - let warm_tools = Rc::clone(&rt.warm_tools); - let shutdown_ref = Arc::clone(&rt.shutdown); - let g = Rc::clone(&gate); - ex.spawn(async move { - let res = run_tool_call( - lua.clone(), - plugin, - tool, - input, - ctx, - deadline, - live, - live_tasks, - warm_tools, - plugins, - shutdown_ref, - g, - ) - .await; - let _ = reply.send(res); - }) - .detach(); + request @ Request::CallTool { .. } => { + let deferred_request = + spawn_runtime_request(&rt, &ex, &gate, &lifecycle, request, false); + debug_assert!(deferred_request.is_none()); } Request::ClearPlugin { plugin, reply } => { - drain_barrier(&rt.lua, &ex, &gate, &spawn_rx).await; + drain_runtime( + &rt, + &ex, + &gate, + &lifecycle, + &spawn_rx, + &rx, + &mut deferred, + ) + .await; rt.clear_plugin(&plugin); + rt.state.drop_plugin(&plugin); let _ = reply.send(()); } Request::RunCommand { @@ -2386,6 +2663,7 @@ pub fn spawn( plugin, tool, input, + nested: _, reply, } => { let res = @@ -2396,9 +2674,17 @@ pub fn spawn( plugin, tool, input, + nested: _, reply, } => { - let res = rt.compute_permission_scopes(&plugin, &tool, input).await; + let res = LuaRuntime::compute_permission_scopes( + &rt.lua, + &rt.plugins, + &plugin, + &tool, + input, + ) + .await; let _ = reply.send(res); } Request::RunInitLua { @@ -2407,7 +2693,16 @@ pub fn spawn( plugin_dir, reply, } => { - drain_barrier(&rt.lua, &ex, &gate, &spawn_rx).await; + drain_runtime( + &rt, + &ex, + &gate, + &lifecycle, + &spawn_rx, + &rx, + &mut deferred, + ) + .await; let res = rt.run_init_lua(&source, &source_name, plugin_dir).await; let _ = reply.send(res); } @@ -2418,6 +2713,82 @@ pub fn spawn( Request::CollectPluginOptions { reply } => { let _ = reply.send(collect_plugin_options(&rt.lua)); } + Request::HydrateState { + identity, + snapshot, + reply, + } => { + drain_runtime( + &rt, + &ex, + &gate, + &lifecycle, + &spawn_rx, + &rx, + &mut deferred, + ) + .await; + let result = validate_snapshot_lua_values( + &rt.lua, + &identity, + snapshot.as_ref(), + ) + .and_then(|()| { + rt.state + .hydrate(identity, snapshot) + .map_err(|error| error.to_string()) + }); + let _ = reply.send(result); + } + Request::CaptureState { + identity, + revision, + reply, + } => { + drain_runtime( + &rt, + &ex, + &gate, + &lifecycle, + &spawn_rx, + &rx, + &mut deferred, + ) + .await; + let result = rt + .state + .capture(&identity, revision) + .map_err(|error| error.to_string()); + let _ = reply.send(result); + } + Request::ResetState { identity, reply } => { + drain_runtime( + &rt, + &ex, + &gate, + &lifecycle, + &spawn_rx, + &rx, + &mut deferred, + ) + .await; + rt.state.reset(&identity); + let _ = reply.send(()); + } + Request::DropStateOwner { owner, reply } => { + drain_runtime( + &rt, + &ex, + &gate, + &lifecycle, + &spawn_rx, + &rx, + &mut deferred, + ) + .await; + rt.state.drop_owner(owner); + let _ = reply.send(()); + } Request::RestoreToolAsync { item, event_tx } => { spawn_restore(&ex, &gate, &restores, &rt, item, event_tx); } @@ -2512,34 +2883,10 @@ pub fn spawn( let _ = reply .send(run_describe(&rt.lua, &rt.plugins, &plugin, &tool, &dctx)); } - Request::StartTool { - plugin, - tool, - input, - live, - ctx, - reply, - } => { - let func = { - let plugins = rt.plugins.borrow(); - plugins - .get(&*plugin) - .and_then(|p| p.get(&*tool)) - .and_then(|tk| tk.start.as_ref()) - .and_then(|key| rt.lua.registry_value::(key).ok()) - }; - let Some(func) = func else { - let _ = reply.send(()); - continue; - }; - let lua = rt.lua.clone(); - let g = Rc::clone(&gate); - ex.spawn(async move { - let _gate_guard = g.acquire().await; - run_tool_start(&lua, func, &tool, input, live, ctx).await; - let _ = reply.send(()); - }) - .detach(); + request @ Request::StartTool { .. } => { + let deferred_request = + spawn_runtime_request(&rt, &ex, &gate, &lifecycle, request, false); + debug_assert!(deferred_request.is_none()); } Request::RunKeybindCallback { id } => { let func = rt.lua.app_data_ref::().and_then(|store| { diff --git a/n00n-lua/src/state.rs b/n00n-lua/src/state.rs new file mode 100644 index 000000000..ad39f202e --- /dev/null +++ b/n00n-lua/src/state.rs @@ -0,0 +1,888 @@ +use n00n_agent::tools::SessionIdentity; +use n00n_storage::{ + id::n00nId, + sessions::{ + MAX_PLUGIN_STATE_BYTES, SessionStateError, StoredSessionStateSnapshot, StoredStateScope, + }, +}; +use serde_json::Value; +use std::{ + collections::{HashMap, HashSet}, + sync::{Mutex, MutexGuard}, +}; + +pub(crate) const PLUGIN_STATE_SCHEMA_VERSION: u32 = 1; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum PluginStateScope { + Session, + Root, +} + +impl PluginStateScope { + pub(crate) fn parse(value: &str) -> Option { + match value { + "session" => Some(Self::Session), + "root" => Some(Self::Root), + _ => None, + } + } + + const fn stored(self) -> StoredStateScope { + match self { + Self::Session => StoredStateScope::Session, + Self::Root => StoredStateScope::Root, + } + } +} + +impl From for PluginStateScope { + fn from(scope: StoredStateScope) -> Self { + match scope { + StoredStateScope::Session => Self::Session, + StoredStateScope::Root => Self::Root, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct PluginStateIdentity { + session_id: n00nId, + root_session_id: n00nId, +} + +impl PluginStateIdentity { + fn owner(&self, scope: PluginStateScope) -> n00nId { + match scope { + PluginStateScope::Session => self.session_id, + PluginStateScope::Root => self.root_session_id, + } + } + + fn owns_scope(&self, scope: PluginStateScope) -> bool { + matches!(scope, PluginStateScope::Session) || self.is_root() + } + + pub(crate) fn is_root(&self) -> bool { + self.session_id == self.root_session_id + } + + fn scope_identity(&self, scope: PluginStateScope) -> Self { + match scope { + PluginStateScope::Session => self.clone(), + PluginStateScope::Root => Self { + session_id: self.root_session_id, + root_session_id: self.root_session_id, + }, + } + } +} + +impl From<&SessionIdentity> for PluginStateIdentity { + fn from(identity: &SessionIdentity) -> Self { + Self { + session_id: identity.session_id().id(), + root_session_id: identity.root_session_id().id(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct StateKey { + plugin: String, + scope: PluginStateScope, + owner: n00nId, +} + +impl StateKey { + fn new(plugin: &str, scope: PluginStateScope, identity: &PluginStateIdentity) -> Self { + Self { + plugin: plugin.to_owned(), + scope, + owner: identity.owner(scope), + } + } +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum PluginStateError { + #[error("plugin state is {bytes} bytes (maximum {maximum})")] + ValueTooLarge { bytes: usize, maximum: usize }, + #[error(transparent)] + Serialize(#[from] serde_json::Error), + #[error(transparent)] + Snapshot(#[from] SessionStateError), +} + +#[derive(Default)] +struct StateInner { + values: HashMap, + managed: HashSet, + bases: HashMap, +} + +#[derive(Default)] +pub(crate) struct PluginStateStore { + inner: Mutex, +} + +impl PluginStateStore { + fn lock(&self) -> MutexGuard<'_, StateInner> { + self.inner + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + #[must_use] + pub(crate) fn get( + &self, + plugin: &str, + scope: PluginStateScope, + identity: &PluginStateIdentity, + ) -> Option { + self.lock() + .values + .get(&StateKey::new(plugin, scope, identity)) + .cloned() + } + + pub(crate) fn replace( + &self, + plugin: &str, + scope: PluginStateScope, + identity: &PluginStateIdentity, + value: Value, + ) -> Result, PluginStateError> { + validate_value_size(&value)?; + let key = StateKey::new(plugin, scope, identity); + let mut inner = self.lock(); + validate_replacement(&inner, &identity.scope_identity(scope), &key, &value)?; + inner.managed.insert(key.clone()); + Ok(inner.values.insert(key, value)) + } + + pub(crate) fn remove( + &self, + plugin: &str, + scope: PluginStateScope, + identity: &PluginStateIdentity, + ) -> Result, PluginStateError> { + let key = StateKey::new(plugin, scope, identity); + let mut inner = self.lock(); + validate_removal(&inner, &identity.scope_identity(scope), &key)?; + inner.managed.insert(key.clone()); + Ok(inner.values.remove(&key)) + } + + pub(crate) fn hydrate( + &self, + identity: PluginStateIdentity, + mut snapshot: Option, + ) -> Result<(), PluginStateError> { + let entries = match snapshot.as_ref() { + Some(snapshot) => snapshot_entries(snapshot)?, + None => Vec::new(), + }; + if !identity.owns_scope(PluginStateScope::Root) + && let Some(snapshot) = snapshot.as_mut() + { + let root_plugins = snapshot + .plugin_names_with_scope(StoredStateScope::Root)? + .into_iter() + .map(str::to_owned) + .collect::>(); + for plugin in root_plugins { + snapshot.remove_plugin_state(&plugin, StoredStateScope::Root)?; + } + } + + let mut inner = self.lock(); + clear_identity_runtime(&mut inner, &identity, false); + for (plugin, scope, payload) in entries { + if !identity.owns_scope(scope) { + continue; + } + let key = StateKey::new(&plugin, scope, &identity); + inner.managed.insert(key.clone()); + inner.values.insert(key, payload); + } + if let Some(snapshot) = snapshot { + inner.bases.insert(identity, snapshot); + } else { + inner.bases.remove(&identity); + } + Ok(()) + } + + pub(crate) fn capture( + &self, + identity: &PluginStateIdentity, + revision: u64, + ) -> Result { + let mut inner = self.lock(); + let mut candidate = inner + .bases + .get(identity) + .cloned() + .unwrap_or_else(|| StoredSessionStateSnapshot::new(revision)); + candidate.set_state_revision(revision)?; + + let managed = inner + .managed + .iter() + .filter(|key| identity.owns_scope(key.scope) && key.owner == identity.owner(key.scope)) + .cloned() + .collect::>(); + for key in managed { + if let Some(value) = inner.values.get(&key) { + candidate.set_plugin_state( + &key.plugin, + PLUGIN_STATE_SCHEMA_VERSION, + key.scope.stored(), + value.clone(), + )?; + } else { + candidate.remove_plugin_state(&key.plugin, key.scope.stored())?; + } + } + + inner.bases.insert(identity.clone(), candidate.clone()); + Ok(candidate) + } + + pub(crate) fn reset(&self, identity: &PluginStateIdentity) { + clear_identity_runtime(&mut self.lock(), identity, true); + } + + pub(crate) fn drop_owner(&self, owner: n00nId) { + let mut inner = self.lock(); + inner.values.retain(|key, _| key.owner != owner); + inner.managed.retain(|key| key.owner != owner); + inner.bases.retain(|identity, _| { + identity.session_id != owner && identity.root_session_id != owner + }); + } + + pub(crate) fn drop_plugin(&self, plugin: &str) { + let mut inner = self.lock(); + let keys = inner + .values + .keys() + .filter(|key| key.plugin == plugin) + .cloned() + .collect::>(); + for key in keys { + inner.values.remove(&key); + inner.managed.insert(key); + } + } +} +fn validate_replacement( + inner: &StateInner, + identity: &PluginStateIdentity, + replacement_key: &StateKey, + replacement_value: &Value, +) -> Result<(), PluginStateError> { + let mut candidate = inner + .bases + .get(identity) + .cloned() + .unwrap_or_else(|| StoredSessionStateSnapshot::new(0)); + let mut replacement_applied = false; + for key in inner + .managed + .iter() + .filter(|key| identity.owns_scope(key.scope) && key.owner == identity.owner(key.scope)) + { + if key == replacement_key { + candidate.set_plugin_state( + &key.plugin, + PLUGIN_STATE_SCHEMA_VERSION, + key.scope.stored(), + replacement_value.clone(), + )?; + replacement_applied = true; + } else if let Some(value) = inner.values.get(key) { + candidate.set_plugin_state( + &key.plugin, + PLUGIN_STATE_SCHEMA_VERSION, + key.scope.stored(), + value.clone(), + )?; + } else { + candidate.remove_plugin_state(&key.plugin, key.scope.stored())?; + } + } + if !replacement_applied { + candidate.set_plugin_state( + &replacement_key.plugin, + PLUGIN_STATE_SCHEMA_VERSION, + replacement_key.scope.stored(), + replacement_value.clone(), + )?; + } + Ok(()) +} + +fn validate_value_size(value: &Value) -> Result<(), PluginStateError> { + let bytes = serde_json::to_vec(value)?.len(); + if bytes > MAX_PLUGIN_STATE_BYTES { + return Err(PluginStateError::ValueTooLarge { + bytes, + maximum: MAX_PLUGIN_STATE_BYTES, + }); + } + + Ok(()) +} + +fn validate_removal( + inner: &StateInner, + identity: &PluginStateIdentity, + removal_key: &StateKey, +) -> Result<(), PluginStateError> { + let mut candidate = inner + .bases + .get(identity) + .cloned() + .unwrap_or_else(|| StoredSessionStateSnapshot::new(0)); + for key in inner.managed.iter().filter(|key| { + identity.owns_scope(key.scope) + && key.owner == identity.owner(key.scope) + && *key != removal_key + }) { + if let Some(value) = inner.values.get(key) { + candidate.set_plugin_state( + &key.plugin, + PLUGIN_STATE_SCHEMA_VERSION, + key.scope.stored(), + value.clone(), + )?; + } else { + candidate.remove_plugin_state(&key.plugin, key.scope.stored())?; + } + } + candidate.remove_plugin_state(&removal_key.plugin, removal_key.scope.stored())?; + Ok(()) +} + +fn snapshot_entries( + snapshot: &StoredSessionStateSnapshot, +) -> Result, PluginStateError> { + snapshot + .plugin_entries_for_apply(PLUGIN_STATE_SCHEMA_VERSION)? + .into_iter() + .map(|entry| { + validate_value_size(entry.payload)?; + Ok(( + entry.plugin.to_owned(), + entry.scope.into(), + entry.payload.clone(), + )) + }) + .collect() +} + +fn clear_identity_runtime( + inner: &mut StateInner, + identity: &PluginStateIdentity, + mark_managed: bool, +) { + let keys = inner + .values + .keys() + .filter(|key| identity.owns_scope(key.scope) && key.owner == identity.owner(key.scope)) + .cloned() + .collect::>(); + inner + .values + .retain(|key, _| !identity.owns_scope(key.scope) || key.owner != identity.owner(key.scope)); + if mark_managed { + inner.managed.extend(keys); + } else { + inner.managed.retain(|key| { + !identity.owns_scope(key.scope) || key.owner != identity.owner(key.scope) + }); + } +} + +#[cfg(test)] +mod tests { + use super::{PluginStateError, PluginStateIdentity, PluginStateScope, PluginStateStore}; + use n00n_agent::tools::SessionIdentity; + use n00n_storage::{ + id::{SessionRef, n00nId}, + sessions::{ + MAX_PLUGIN_STATE_BYTES, SESSION_STATE_SCHEMA_VERSION, SessionMeta, + StoredSessionStateSnapshot, StoredStateScope, + }, + }; + use serde_json::json; + use std::{ + sync::{Arc, Barrier}, + thread, + }; + + fn identity() -> PluginStateIdentity { + PluginStateIdentity::from(&SessionIdentity::child( + SessionRef::generate(), + SessionRef::generate(), + )) + } + + #[test] + fn identity_uses_canonical_ids() { + let raw = "550e8400-e29b-41d4-a716-446655440000"; + let id = raw.parse::().unwrap(); + let legacy = SessionIdentity::root(raw.parse::().unwrap()); + let canonical = SessionIdentity::root(SessionRef::from_id(id)); + + assert_eq!( + PluginStateIdentity::from(&legacy), + PluginStateIdentity::from(&canonical) + ); + } + + #[test] + fn plugin_scope_and_owner_are_all_part_of_the_key() { + let store = PluginStateStore::default(); + let root = SessionRef::generate(); + let a = PluginStateIdentity::from(&SessionIdentity::child( + SessionRef::generate(), + root.clone(), + )); + let b = PluginStateIdentity::from(&SessionIdentity::child(SessionRef::generate(), root)); + store + .replace("one", PluginStateScope::Session, &a, json!(1)) + .unwrap(); + store + .replace("two", PluginStateScope::Session, &a, json!(2)) + .unwrap(); + store + .replace("one", PluginStateScope::Session, &b, json!(3)) + .unwrap(); + store + .replace("one", PluginStateScope::Root, &a, json!(4)) + .unwrap(); + + assert_eq!( + store.get("one", PluginStateScope::Session, &a), + Some(json!(1)) + ); + assert_eq!( + store.get("two", PluginStateScope::Session, &a), + Some(json!(2)) + ); + assert_eq!( + store.get("one", PluginStateScope::Session, &b), + Some(json!(3)) + ); + assert_eq!(store.get("one", PluginStateScope::Root, &b), Some(json!(4))); + } + + #[test] + fn oversized_replace_is_atomic() { + let store = PluginStateStore::default(); + let identity = identity(); + store + .replace( + "plugin", + PluginStateScope::Session, + &identity, + json!("kept"), + ) + .unwrap(); + + let error = store + .replace( + "plugin", + PluginStateScope::Session, + &identity, + json!("x".repeat(MAX_PLUGIN_STATE_BYTES)), + ) + .unwrap_err(); + + assert!(matches!(error, PluginStateError::ValueTooLarge { .. })); + assert_eq!( + store.get("plugin", PluginStateScope::Session, &identity), + Some(json!("kept")) + ); + } + + #[test] + fn capture_preserves_opaque_data_and_advances_revision() { + let store = PluginStateStore::default(); + let identity = identity(); + let snapshot = serde_json::from_value(json!({ + "schema_version": SESSION_STATE_SCHEMA_VERSION, + "state_revision": 3, + "future": {"kept": true}, + "plugins": {"unknown": {"future": {"raw": "kept"}}} + })) + .unwrap(); + store.hydrate(identity.clone(), Some(snapshot)).unwrap(); + store + .replace( + "known", + PluginStateScope::Session, + &identity, + json!({"v": 1}), + ) + .unwrap(); + + let captured = store.capture(&identity, 4).unwrap(); + let raw = serde_json::to_value(captured).unwrap(); + assert_eq!(raw["state_revision"], json!(4)); + assert_eq!(raw["future"], json!({"kept": true})); + assert_eq!(raw["plugins"]["unknown"]["future"], json!({"raw": "kept"})); + } + + #[test] + fn future_and_malformed_hydrates_leave_current_state_unchanged() { + let store = PluginStateStore::default(); + let identity = identity(); + store + .replace( + "plugin", + PluginStateScope::Session, + &identity, + json!("kept"), + ) + .unwrap(); + let future = serde_json::from_value(json!({ + "schema_version": SESSION_STATE_SCHEMA_VERSION + 1, + "state_revision": 7, + "opaque": true + })) + .unwrap(); + assert!(store.hydrate(identity.clone(), Some(future)).is_err()); + assert_eq!( + store.get("plugin", PluginStateScope::Session, &identity), + Some(json!("kept")) + ); + + let meta: SessionMeta = serde_json::from_value(json!({ + "state_snapshot": { + "schema_version": SESSION_STATE_SCHEMA_VERSION, + "plugins": {} + } + })) + .unwrap(); + assert!( + store + .hydrate(identity.clone(), meta.state_snapshot) + .is_err() + ); + assert_eq!( + store.get("plugin", PluginStateScope::Session, &identity), + Some(json!("kept")) + ); + } + + #[test] + fn hydrate_replaces_pending_runtime_removals() { + let store = PluginStateStore::default(); + let identity = identity(); + store + .replace("plugin", PluginStateScope::Session, &identity, json!("old")) + .unwrap(); + store + .remove("plugin", PluginStateScope::Session, &identity) + .unwrap(); + + let mut snapshot = StoredSessionStateSnapshot::new(2); + snapshot + .set_plugin_state("plugin", 1, StoredStateScope::Session, json!("hydrated")) + .unwrap(); + store.hydrate(identity.clone(), Some(snapshot)).unwrap(); + + let captured = store.capture(&identity, 3).unwrap(); + assert_eq!( + captured + .plugin_payload_for_apply("plugin", 1, StoredStateScope::Session) + .unwrap(), + Some(&json!("hydrated")) + ); + } + + #[test] + fn failed_revision_regression_does_not_replace_base() { + let store = PluginStateStore::default(); + let identity = identity(); + let snapshot = StoredSessionStateSnapshot::new(5); + store.hydrate(identity.clone(), Some(snapshot)).unwrap(); + + assert!(store.capture(&identity, 4).is_err()); + assert_eq!( + store.capture(&identity, 6).unwrap().state_revision(), + Some(6) + ); + } + + #[test] + fn reset_removes_managed_state_but_preserves_opaque_base() { + let store = PluginStateStore::default(); + let identity = identity(); + let mut snapshot = StoredSessionStateSnapshot::new(1); + snapshot + .set_plugin_state("known", 1, StoredStateScope::Session, json!("old")) + .unwrap(); + snapshot + .set_plugin_state("future", 2, StoredStateScope::Session, json!("opaque")) + .unwrap(); + store.hydrate(identity.clone(), Some(snapshot)).unwrap(); + + store.reset(&identity); + let captured = store.capture(&identity, 2).unwrap(); + assert_eq!( + captured + .plugin_payload_for_apply("known", 1, StoredStateScope::Session) + .unwrap(), + None + ); + assert_eq!( + captured + .plugin_payload_for_apply("future", 2, StoredStateScope::Session) + .unwrap(), + Some(&json!("opaque")) + ); + } + + #[test] + fn replacement_rejects_uncapturable_names_and_entry_counts_atomically() { + let store = PluginStateStore::default(); + let identity = identity(); + assert!( + store + .replace( + "bad/name", + PluginStateScope::Session, + &identity, + json!(true), + ) + .is_err() + ); + assert_eq!( + store.get("bad/name", PluginStateScope::Session, &identity), + None + ); + + for index in 0..64 { + store + .replace( + &format!("plugin_{index}"), + PluginStateScope::Session, + &identity, + json!(index), + ) + .unwrap(); + } + assert!( + store + .replace("plugin_64", PluginStateScope::Session, &identity, json!(64),) + .is_err() + ); + assert_eq!( + store.get("plugin_64", PluginStateScope::Session, &identity), + None + ); + } + + #[test] + fn child_hydration_cannot_clobber_live_root_state() { + let store = PluginStateStore::default(); + let root_ref = SessionRef::generate(); + let root = PluginStateIdentity::from(&SessionIdentity::root(root_ref.clone())); + let child = + PluginStateIdentity::from(&SessionIdentity::child(SessionRef::generate(), root_ref)); + store + .replace("plugin", PluginStateScope::Root, &root, json!("live")) + .unwrap(); + let mut child_snapshot = StoredSessionStateSnapshot::new(1); + child_snapshot + .set_plugin_state("plugin", 1, StoredStateScope::Root, json!("stale")) + .unwrap(); + + store.hydrate(child.clone(), Some(child_snapshot)).unwrap(); + + assert_eq!( + store.get("plugin", PluginStateScope::Root, &child), + Some(json!("live")) + ); + } + + #[test] + fn child_root_replacement_uses_root_snapshot_limits() { + let store = PluginStateStore::default(); + let root_ref = SessionRef::generate(); + let root = PluginStateIdentity::from(&SessionIdentity::root(root_ref.clone())); + let child = + PluginStateIdentity::from(&SessionIdentity::child(SessionRef::generate(), root_ref)); + for index in 0..64 { + store + .replace( + &format!("plugin_{index}"), + PluginStateScope::Root, + &root, + json!(index), + ) + .unwrap(); + } + + assert!( + store + .replace("plugin_64", PluginStateScope::Root, &child, json!(64),) + .is_err() + ); + assert_eq!(store.get("plugin_64", PluginStateScope::Root, &root), None); + } + + #[test] + fn removal_from_malformed_container_is_atomic() { + let store = PluginStateStore::default(); + let root_ref = SessionRef::generate(); + let identity = PluginStateIdentity::from(&SessionIdentity::root(root_ref)); + let snapshot: StoredSessionStateSnapshot = serde_json::from_value(json!({ + "schema_version": SESSION_STATE_SCHEMA_VERSION, + "state_revision": 1, + "plugins": {"plugin": null} + })) + .unwrap(); + store.hydrate(identity.clone(), Some(snapshot)).unwrap(); + + assert!( + store + .remove("plugin", PluginStateScope::Root, &identity) + .is_err() + ); + let captured = store.capture(&identity, 2).unwrap(); + assert_eq!( + serde_json::to_value(captured).unwrap()["plugins"]["plugin"], + json!(null) + ); + } + + #[test] + fn concurrent_child_mutations_remain_isolated() { + let store = Arc::new(PluginStateStore::default()); + let root_ref = SessionRef::generate(); + let root = PluginStateIdentity::from(&SessionIdentity::root(root_ref.clone())); + let first = PluginStateIdentity::from(&SessionIdentity::child( + SessionRef::generate(), + root_ref.clone(), + )); + let second = + PluginStateIdentity::from(&SessionIdentity::child(SessionRef::generate(), root_ref)); + store + .replace("plugin", PluginStateScope::Root, &root, json!("root")) + .unwrap(); + let barrier = Arc::new(Barrier::new(2)); + + let workers = [first.clone(), second.clone()].map(|identity| { + let store = Arc::clone(&store); + let barrier = Arc::clone(&barrier); + thread::spawn(move || { + barrier.wait(); + for value in 0..100 { + store + .replace("plugin", PluginStateScope::Session, &identity, json!(value)) + .unwrap(); + } + }) + }); + for worker in workers { + worker.join().unwrap(); + } + + assert_eq!( + store.get("plugin", PluginStateScope::Session, &first), + Some(json!(99)) + ); + assert_eq!( + store.get("plugin", PluginStateScope::Session, &second), + Some(json!(99)) + ); + assert_eq!( + store.get("plugin", PluginStateScope::Root, &first), + Some(json!("root")) + ); + } + + #[test] + fn dropping_child_owner_preserves_root_and_sibling_state() { + let store = PluginStateStore::default(); + let root_ref = SessionRef::generate(); + let root = PluginStateIdentity::from(&SessionIdentity::root(root_ref.clone())); + let child = PluginStateIdentity::from(&SessionIdentity::child( + SessionRef::generate(), + root_ref.clone(), + )); + let sibling = + PluginStateIdentity::from(&SessionIdentity::child(SessionRef::generate(), root_ref)); + store + .replace("plugin", PluginStateScope::Root, &root, json!("root")) + .unwrap(); + store + .replace("plugin", PluginStateScope::Session, &child, json!("child")) + .unwrap(); + store + .replace( + "plugin", + PluginStateScope::Session, + &sibling, + json!("sibling"), + ) + .unwrap(); + + store.drop_owner(child.session_id); + + assert_eq!(store.get("plugin", PluginStateScope::Session, &child), None); + assert_eq!( + store.get("plugin", PluginStateScope::Session, &sibling), + Some(json!("sibling")) + ); + assert_eq!( + store.get("plugin", PluginStateScope::Root, &child), + Some(json!("root")) + ); + } + + #[test] + fn child_capture_does_not_emit_supported_root_state() { + let store = PluginStateStore::default(); + let root_ref = SessionRef::generate(); + let child = + PluginStateIdentity::from(&SessionIdentity::child(SessionRef::generate(), root_ref)); + let mut snapshot = StoredSessionStateSnapshot::new(1); + snapshot + .set_plugin_state("plugin", 1, StoredStateScope::Root, json!("stale")) + .unwrap(); + snapshot + .set_plugin_state("plugin", 1, StoredStateScope::Session, json!("session")) + .unwrap(); + snapshot + .set_plugin_state("future", 2, StoredStateScope::Root, json!("opaque")) + .unwrap(); + store.hydrate(child.clone(), Some(snapshot)).unwrap(); + + let captured = store.capture(&child, 2).unwrap(); + assert_eq!( + captured + .plugin_payload_for_apply("plugin", 1, StoredStateScope::Root) + .unwrap(), + None + ); + assert_eq!( + captured + .plugin_payload_for_apply("future", 2, StoredStateScope::Root) + .unwrap(), + None + ); + assert_eq!( + captured + .plugin_payload_for_apply("plugin", 1, StoredStateScope::Session) + .unwrap(), + Some(&json!("session")) + ); + } +} diff --git a/n00n-lua/tests/plugin_host.rs b/n00n-lua/tests/plugin_host.rs index 3e85c99f8..0818ac0ee 100644 --- a/n00n-lua/tests/plugin_host.rs +++ b/n00n-lua/tests/plugin_host.rs @@ -12,8 +12,8 @@ use std::time::Duration; use n00n_agent::template::env_vars; use n00n_agent::tools::{ - ActiveTools, DescriptionContext, ToolAudience, ToolFilter, ToolRegistry, ToolSource, - timeout_annotation, + ActiveTools, DescriptionContext, SessionIdentity, ToolAudience, ToolFilter, ToolRegistry, + ToolSource, timeout_annotation, }; use n00n_config::{AlwaysThinking, PluginsConfig, ToolOutputLines}; use n00n_lua::{PluginError, PluginHost, WARM_TOOL_CAP}; @@ -4061,6 +4061,428 @@ fn session_rejects_nonempty_lua_tools_object() { assert!(error.contains(TOOLS_MUST_BE_ARRAY_ERR), "got: {error}"); } +#[test] +fn lua_session_rejects_missing_identity() { + smol::block_on(async { + let reg = fresh_registry(); + let host = PluginHost::new(Arc::clone(®)).unwrap(); + let src = format!( + r#"n00n.api.register_tool({{ + name = "missing_identity_probe", + description = "test", + schema = {MINIMAL_SCHEMA}, + audiences = {{ "main" }}, + handler = function(input, ctx) + local sess, err = n00n.agent.session(ctx, {{}}) + if sess ~= nil then return "unexpected session" end + return err or "no error" + end + }})"# + ); + host.load_source("missing_identity_plugin", &src).unwrap(); + let entry = reg.get("missing_identity_probe").unwrap(); + let invocation = entry.tool.parse(&serde_json::json!({})).unwrap(); + let mut ctx = n00n_agent::tools::test_support::stub_ctx(&n00n_agent::AgentMode::Build); + ctx.identity = None; + + let result = invocation.execute(&ctx).await.output.unwrap(); + let n00n_agent::ToolOutput::Plain(output) = result else { + panic!("expected plain output"); + }; + assert_eq!(output.text, "session identity is unavailable"); + }); +} + +#[test] +fn plugin_state_capture_waits_for_inflight_handler_callbacks() { + let reg = fresh_registry(); + let host = PluginHost::new(Arc::clone(®)).unwrap(); + let source = format!( + r#" + n00n.api.register_tool({{ + name = "delayed_state", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) + local buf = n00n.ui.buf() + buf:set_lines({{ "waiting" }}) + ctx:live_buf(buf) + n00n.async.run(function() + local id = n00n.fn.jobstart("sleep 0.05") + n00n.fn.jobwait(id) + return "finished" + end, function(err) + if err then + ctx:finish(err) + return + end + local _, state_err = ctx:state_replace("session", {{ value = "finished" }}) + ctx:finish(state_err or "done") + end) + return nil + end, + }}) + "# + ); + host.load_source("delayed_state", &source).unwrap(); + let identity = SessionIdentity::root(SessionRef::generate()); + let entry = reg.get("delayed_state").unwrap(); + let invocation = entry.tool.parse(&serde_json::json!({})).unwrap(); + let (event_tx, event_rx) = flume::unbounded(); + let sender = n00n_agent::EventSender::new(event_tx, 0); + let mut ctx = n00n_agent::tools::test_support::stub_ctx_with( + &n00n_agent::AgentMode::Build, + Some(&sender), + Some("delayed-state"), + ); + ctx.identity = Some(identity.clone()); + let worker = std::thread::spawn(move || smol::block_on(invocation.execute(&ctx))); + + loop { + let event = event_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + if matches!( + event.event, + n00n_agent::AgentEvent::LiveToolBuf { ref id, .. } if id == "delayed-state" + ) { + break; + } + } + + let snapshot = host + .event_handle() + .unwrap() + .capture_state(&identity, 1) + .unwrap(); + assert_eq!( + snapshot + .plugin_payload_for_apply( + "delayed_state", + 1, + n00n_storage::sessions::StoredStateScope::Session, + ) + .unwrap(), + Some(&serde_json::json!({"value": "finished"})) + ); + assert_eq!(worker.join().unwrap().output.unwrap().as_text(), "done"); +} + +#[test] +fn plugin_state_lifecycle_methods_reject_dead_host() { + let reg = fresh_registry(); + let host = PluginHost::new(reg).unwrap(); + let handle = host.event_handle().unwrap(); + let identity = SessionIdentity::root(SessionRef::generate()); + drop(host); + + assert!(matches!( + handle.capture_state(&identity, 1), + Err(PluginError::HostDead) + )); + assert!(matches!( + handle.hydrate_state(&identity, None), + Err(PluginError::HostDead) + )); + assert!(matches!( + handle.reset_state(&identity), + Err(PluginError::HostDead) + )); + assert!(matches!( + handle.drop_state_owner(identity.session_id().id()), + Err(PluginError::HostDead) + )); +} + +#[test] +fn plugin_state_capture_services_nested_lua_tool_calls() { + let reg = fresh_registry(); + let host = PluginHost::new(Arc::clone(®)).unwrap(); + let source = format!( + r#" + n00n.api.register_tool({{ + name = "nested_state_writer", description = "test", schema = {MINIMAL_SCHEMA}, + header = function() return "nested writer" end, + handler = function(input, ctx) + local _, err = ctx:state_replace("session", {{ value = "nested" }}) + return err or "written" + end, + }}) + n00n.api.register_tool({{ + name = "nested_state_parent", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) + local buf = n00n.ui.buf() + buf:set_lines({{ "waiting" }}) + ctx:live_buf(buf) + local id = n00n.fn.jobstart("sleep 0.5") + n00n.fn.jobwait(id) + local result, err = n00n.agent.call_tool(ctx, "nested_state_writer", {{}}) + return err or result + end, + }}) + "# + ); + host.load_source("nested_state", &source).unwrap(); + let identity = SessionIdentity::root(SessionRef::generate()); + let entry = reg.get("nested_state_parent").unwrap(); + let invocation = entry.tool.parse(&serde_json::json!({})).unwrap(); + let (event_tx, event_rx) = flume::unbounded(); + let sender = n00n_agent::EventSender::new(event_tx, 0); + let mut ctx = n00n_agent::tools::test_support::stub_ctx_with( + &n00n_agent::AgentMode::Build, + Some(&sender), + Some("nested-state"), + ); + ctx.identity = Some(identity.clone()); + ctx.registry = Arc::clone(®); + let worker = std::thread::spawn(move || smol::block_on(invocation.execute(&ctx))); + + loop { + let event = event_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + if matches!( + event.event, + n00n_agent::AgentEvent::LiveToolBuf { ref id, .. } if id == "nested-state" + ) { + break; + } + } + + let snapshot = host + .event_handle() + .unwrap() + .capture_state(&identity, 1) + .unwrap(); + assert_eq!( + snapshot + .plugin_payload_for_apply( + "nested_state", + 1, + n00n_storage::sessions::StoredStateScope::Session, + ) + .unwrap(), + Some(&serde_json::json!({"value": "nested"})) + ); + assert_eq!(worker.join().unwrap().output.unwrap().as_text(), "written"); +} + +#[test] +fn plugin_state_isolates_namespaces_and_session_scope_while_sharing_root_scope() { + let reg = fresh_registry(); + let host = PluginHost::new(Arc::clone(®)).unwrap(); + let plugin_a = format!( + r#" + local function read(ctx, scope) + local value, err = ctx:state_get(scope) + if err then return err end + return value and value.name or "none" + end + n00n.api.register_tool({{ + name = "a_write_root", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) + local _, err = ctx:state_replace("root", {{ name = "root-a" }}) + return err or "ok" + end, + }}) + n00n.api.register_tool({{ + name = "a_write_session", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) + local _, err = ctx:state_replace("session", {{ name = "session-a" }}) + return err or "ok" + end, + }}) + n00n.api.register_tool({{ + name = "a_read_root", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) return read(ctx, "root") end, + }}) + n00n.api.register_tool({{ + name = "a_read_session", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) return read(ctx, "session") end, + }}) + "# + ); + let plugin_b = format!( + r#"n00n.api.register_tool({{ + name = "b_read_root", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) + local value, err = ctx:state_get("root") + if err then return err end + return value and value.name or "none" + end, + }})"# + ); + host.load_source("plugin_a", &plugin_a).unwrap(); + host.load_source("plugin_b", &plugin_b).unwrap(); + + let root = SessionIdentity::root(SessionRef::generate()); + let child = SessionIdentity::child(SessionRef::generate(), root.root_session_id().clone()); + let execute = |name: &str, identity: &SessionIdentity| { + let entry = reg.get(name).unwrap(); + let invocation = entry.tool.parse(&serde_json::json!({})).unwrap(); + let mut ctx = n00n_agent::tools::test_support::stub_ctx(&n00n_agent::AgentMode::Build); + ctx.identity = Some(identity.clone()); + let output = smol::block_on(async { invocation.execute(&ctx).await }) + .output + .unwrap(); + let n00n_agent::ToolOutput::Plain(output) = output else { + panic!("expected plain output"); + }; + output.text + }; + + assert_eq!(execute("a_write_root", &root), "ok"); + assert_eq!(execute("a_write_session", &root), "ok"); + assert_eq!(execute("a_read_root", &child), "root-a"); + assert_eq!(execute("a_read_session", &child), "none"); + assert_eq!(execute("b_read_root", &root), "none"); + + let handle = host.event_handle().unwrap(); + let captured = handle.capture_state(&root, 7).unwrap(); + assert_eq!(captured.state_revision(), Some(7)); + assert_eq!( + captured + .plugin_payload_for_apply( + "plugin_a", + 1, + n00n_storage::sessions::StoredStateScope::Root + ) + .unwrap(), + Some(&serde_json::json!({"name": "root-a"})) + ); + handle.reset_state(&root).unwrap(); + let reset = handle.capture_state(&root, 8).unwrap(); + assert!( + reset + .plugin_payload_for_apply( + "plugin_a", + 1, + n00n_storage::sessions::StoredStateScope::Root + ) + .unwrap() + .is_none() + ); + handle.hydrate_state(&root, Some(captured)).unwrap(); + assert_eq!(execute("a_read_root", &root), "root-a"); + assert_eq!(execute("a_read_session", &root), "session-a"); + + host.unload("plugin_a").unwrap(); + let unloaded = handle.capture_state(&root, 8).unwrap(); + assert!( + unloaded + .plugin_payload_for_apply( + "plugin_a", + 1, + n00n_storage::sessions::StoredStateScope::Root, + ) + .unwrap() + .is_none() + ); +} + +#[test] +fn plugin_state_rejects_context_reuse_after_handler_finishes() { + let reg = fresh_registry(); + let host = PluginHost::new(Arc::clone(®)).unwrap(); + let source = format!( + r#" + local saved + n00n.api.register_tool({{ + name = "save_state_ctx", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) + saved = ctx + local _, err = ctx:state_replace("session", {{ value = "original" }}) + return err or "saved" + end, + }}) + n00n.api.register_tool({{ + name = "save_dispatch_ctx", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function(input, ctx) + saved = ctx + return "saved" + end, + }}) + n00n.api.register_tool({{ + name = "reuse_state_ctx", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function() + local _, err = saved:state_replace("session", {{ value = "stale" }}) + return err or "unexpected success" + end, + }}) + n00n.api.register_tool({{ + name = "reuse_state_ctx_deadline", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function() + local _, err = saved:set_deadline(1) + return err or "unexpected success" + end, + }}) + n00n.api.register_tool({{ + name = "reuse_state_ctx_dispatch", description = "test", schema = {MINIMAL_SCHEMA}, + handler = function() + local _, err = n00n.agent.call_tool(saved, "missing", {{}}) + return err or "unexpected success" + end, + }}) + "# + ); + host.load_source("stale_ctx", &source).unwrap(); + let identity = SessionIdentity::root(SessionRef::generate()); + let execute = |name: &str| { + let entry = reg.get(name).unwrap(); + let invocation = entry.tool.parse(&serde_json::json!({})).unwrap(); + let mut ctx = n00n_agent::tools::test_support::stub_ctx(&n00n_agent::AgentMode::Build); + ctx.identity = Some(identity.clone()); + let output = smol::block_on(async { invocation.execute(&ctx).await }) + .output + .unwrap(); + let n00n_agent::ToolOutput::Plain(output) = output else { + panic!("expected plain output"); + }; + output.text + }; + + assert_eq!(execute("save_state_ctx"), "saved"); + assert_eq!( + execute("reuse_state_ctx"), + "state context is no longer active" + ); + assert_eq!( + execute("reuse_state_ctx_dispatch"), + "state context is no longer active" + ); + assert_eq!( + execute("reuse_state_ctx_deadline"), + "state context is no longer active" + ); + + let execute_without_identity = |name: &str| { + let entry = reg.get(name).unwrap(); + let invocation = entry.tool.parse(&serde_json::json!({})).unwrap(); + let ctx = n00n_agent::tools::test_support::stub_ctx(&n00n_agent::AgentMode::Build); + let output = smol::block_on(async { invocation.execute(&ctx).await }) + .output + .unwrap(); + let n00n_agent::ToolOutput::Plain(output) = output else { + panic!("expected plain output"); + }; + output.text + }; + assert_eq!(execute_without_identity("save_dispatch_ctx"), "saved"); + assert_eq!( + execute_without_identity("reuse_state_ctx_dispatch"), + "state context is no longer active" + ); + + let snapshot = host + .event_handle() + .unwrap() + .capture_state(&identity, 1) + .unwrap(); + assert_eq!( + snapshot + .plugin_payload_for_apply( + "stale_ctx", + 1, + n00n_storage::sessions::StoredStateScope::Session, + ) + .unwrap(), + Some(&serde_json::json!({"value": "original"})) + ); +} #[test] fn lua_sessions_under_one_parent_use_unique_identity_everywhere() { diff --git a/n00n-storage/src/sessions.rs b/n00n-storage/src/sessions.rs index 61cbda1dc..34da3876c 100644 --- a/n00n-storage/src/sessions.rs +++ b/n00n-storage/src/sessions.rs @@ -51,7 +51,7 @@ const MAX_FIRST_MESSAGE_BYTES: usize = 256 * 1024; pub const SESSION_STATE_SCHEMA_VERSION: u32 = 1; const MAX_PLUGIN_STATE_ENTRIES: usize = 64; const MAX_PLUGIN_STATE_NAME_BYTES: usize = 128; -const MAX_PLUGIN_STATE_BYTES: usize = 256 * 1024; +pub const MAX_PLUGIN_STATE_BYTES: usize = 256 * 1024; const MAX_SESSION_STATE_BYTES: usize = 1024 * 1024; const META_RECORD_PREFIX: &str = r#"{"t":"meta""#; const MSG_RECORD_PREFIX: &str = r#"{"t":"msg""#; @@ -277,11 +277,38 @@ impl Default for StoredPluginScopes { } impl StoredPluginScopes { - fn insert(&mut self, scope: StoredStateScope, state: StoredPluginState) { - self.scopes - .get_or_insert_with(BTreeMap::new) - .insert(scope.as_str().to_owned(), state); - self.malformed = None; + fn set( + &mut self, + plugin: &str, + scope: StoredStateScope, + schema_version: u32, + payload: serde_json::Value, + ) -> Result<(), SessionStateError> { + let Some(scopes) = self.scopes.as_mut() else { + return Err(SessionStateError::InvalidPluginContainer { + plugin: plugin.to_owned(), + }); + }; + if let Some(state) = scopes.get_mut(scope.as_str()) { + let serde_json::Value::Object(fields) = &mut state.raw else { + return Err(SessionStateError::InvalidPluginState { + plugin: plugin.to_owned(), + scope, + }); + }; + fields.insert( + "schema_version".to_owned(), + serde_json::Value::from(schema_version), + ); + fields.insert("payload".to_owned(), payload); + state.schema_version = Some(u64::from(schema_version)); + } else { + scopes.insert( + scope.as_str().to_owned(), + StoredPluginState::new(schema_version, payload), + ); + } + Ok(()) } } @@ -338,6 +365,9 @@ enum StoredSessionStateSnapshotInner { schema_version: u64, raw: serde_json::Value, }, + Malformed { + raw: serde_json::Value, + }, } #[derive(Debug, Clone, PartialEq)] @@ -345,6 +375,13 @@ pub struct StoredSessionStateSnapshot { inner: StoredSessionStateSnapshotInner, } +#[derive(Debug, Clone, PartialEq)] +pub struct StoredPluginStateEntry<'a> { + pub plugin: &'a str, + pub scope: StoredStateScope, + pub payload: &'a serde_json::Value, +} + impl Default for StoredSessionStateSnapshot { fn default() -> Self { Self::new(0) @@ -358,7 +395,8 @@ impl Serialize for StoredSessionStateSnapshot { { match &self.inner { StoredSessionStateSnapshotInner::Supported(snapshot) => snapshot.serialize(serializer), - StoredSessionStateSnapshotInner::Unsupported { raw, .. } => raw.serialize(serializer), + StoredSessionStateSnapshotInner::Unsupported { raw, .. } + | StoredSessionStateSnapshotInner::Malformed { raw } => raw.serialize(serializer), } } } @@ -405,6 +443,8 @@ impl<'de> Deserialize<'de> for StoredSessionStateSnapshot { pub enum SessionStateError { #[error("unsupported session-state schema version {found} (expected {expected})")] UnsupportedSchemaVersion { found: u64, expected: u32 }, + #[error("session-state snapshot envelope is malformed")] + InvalidEnvelope, #[error("session state has {found} plugin entries (maximum {maximum})")] TooManyPlugins { found: usize, maximum: usize }, #[error("invalid session-state plugin name {plugin:?}")] @@ -434,12 +474,8 @@ pub enum SessionStateError { found: u64, expected: u32, }, - #[error("state scope mismatch for plugin {plugin:?}: found {found:?}, expected {expected:?}")] - ScopeMismatch { - plugin: String, - found: StoredStateScope, - expected: StoredStateScope, - }, + #[error("session-state revision cannot regress from {current} to {requested}")] + StateRevisionRegression { current: u64, requested: u64 }, #[error("failed to measure serialized session state: {0}")] Serialize(#[from] serde_json::Error), } @@ -461,15 +497,37 @@ impl StoredSessionStateSnapshot { pub fn state_revision(&self) -> Option { match &self.inner { StoredSessionStateSnapshotInner::Supported(snapshot) => Some(snapshot.state_revision), - StoredSessionStateSnapshotInner::Unsupported { .. } => None, + StoredSessionStateSnapshotInner::Unsupported { .. } + | StoredSessionStateSnapshotInner::Malformed { .. } => None, } } - /// Adds or replaces one scope of a plugin's state after enforcing snapshot bounds. + /// Advances the state revision without allowing regression. /// /// # Errors - /// Returns a typed error for invalid names or exceeded bounds. - pub fn insert_plugin_state( + /// Returns a typed error for unsupported envelopes, revision regression, or exceeded bounds. + pub fn set_state_revision(&mut self, state_revision: u64) -> Result<(), SessionStateError> { + let StoredSessionStateSnapshotInner::Supported(snapshot) = &self.inner else { + return Err(self.unsupported_schema_error()); + }; + if state_revision < snapshot.state_revision { + return Err(SessionStateError::StateRevisionRegression { + current: snapshot.state_revision, + requested: state_revision, + }); + } + let mut candidate = snapshot.clone(); + candidate.state_revision = state_revision; + validate_supported_snapshot(&candidate)?; + self.inner = StoredSessionStateSnapshotInner::Supported(candidate); + Ok(()) + } + + /// Adds or replaces one exact plugin scope after enforcing snapshot bounds. + /// + /// # Errors + /// Returns a typed error for unsupported envelopes, invalid names, or exceeded bounds. + pub fn set_plugin_state( &mut self, plugin: &str, schema_version: u32, @@ -480,16 +538,132 @@ impl StoredSessionStateSnapshot { return Err(self.unsupported_schema_error()); }; let mut candidate = snapshot.clone(); + validate_plugin_name(plugin)?; candidate .plugins .entry(plugin.to_owned()) .or_default() - .insert(scope, StoredPluginState::new(schema_version, payload)); + .set(plugin, scope, schema_version, payload)?; + validate_supported_snapshot(&candidate)?; + self.inner = StoredSessionStateSnapshotInner::Supported(candidate); + Ok(()) + } + + /// Adds or replaces one exact plugin scope after enforcing snapshot bounds. + /// + /// # Errors + /// Returns a typed error for unsupported envelopes, invalid names, or exceeded bounds. + pub fn insert_plugin_state( + &mut self, + plugin: &str, + schema_version: u32, + scope: StoredStateScope, + payload: serde_json::Value, + ) -> Result<(), SessionStateError> { + self.set_plugin_state(plugin, schema_version, scope, payload) + } + + /// Removes one exact plugin scope while preserving every sibling entry. + /// + /// # Errors + /// Returns a typed error for unsupported envelopes, invalid names, or exceeded bounds. + pub fn remove_plugin_state( + &mut self, + plugin: &str, + scope: StoredStateScope, + ) -> Result<(), SessionStateError> { + let StoredSessionStateSnapshotInner::Supported(snapshot) = &self.inner else { + return Err(self.unsupported_schema_error()); + }; + validate_plugin_name(plugin)?; + let mut candidate = snapshot.clone(); + if candidate + .plugins + .get(plugin) + .is_some_and(|stored_scopes| stored_scopes.scopes.is_none()) + { + return Err(SessionStateError::InvalidPluginContainer { + plugin: plugin.to_owned(), + }); + } + let remove_plugin = candidate + .plugins + .get_mut(plugin) + .and_then(|stored_scopes| stored_scopes.scopes.as_mut()) + .is_some_and(|scopes| { + scopes.remove(scope.as_str()); + scopes.is_empty() + }); + if remove_plugin { + candidate.plugins.remove(plugin); + } validate_supported_snapshot(&candidate)?; self.inner = StoredSessionStateSnapshotInner::Supported(candidate); Ok(()) } + /// Returns plugin names that contain the exact stored scope, including opaque state versions. + /// + /// # Errors + /// Returns a typed error when the snapshot envelope version is unsupported. + pub fn plugin_names_with_scope( + &self, + scope: StoredStateScope, + ) -> Result, SessionStateError> { + let StoredSessionStateSnapshotInner::Supported(snapshot) = &self.inner else { + return Err(self.unsupported_schema_error()); + }; + Ok(snapshot + .plugins + .iter() + .filter_map(|(plugin, stored_scopes)| { + stored_scopes + .scopes + .as_ref() + .is_some_and(|scopes| scopes.contains_key(scope.as_str())) + .then_some(plugin.as_str()) + }) + .collect()) + } + + /// Enumerates well-formed entries matching a supported plugin state version. + /// + /// Malformed entries, unknown scopes, and other plugin state versions remain preserved but are + /// omitted from the result. + /// + /// # Errors + /// Returns a typed error when the snapshot envelope version is unsupported. + pub fn plugin_entries_for_apply( + &self, + schema_version: u32, + ) -> Result>, SessionStateError> { + let StoredSessionStateSnapshotInner::Supported(snapshot) = &self.inner else { + return Err(self.unsupported_schema_error()); + }; + let mut entries = Vec::new(); + for (plugin, stored_scopes) in &snapshot.plugins { + let Some(scopes) = stored_scopes.scopes.as_ref() else { + continue; + }; + for (scope_name, state) in scopes { + let Some(scope) = StoredStateScope::from_stored(scope_name) else { + continue; + }; + let (Some(found), Some(payload)) = (state.schema_version, state.payload()) else { + continue; + }; + if found == u64::from(schema_version) { + entries.push(StoredPluginStateEntry { + plugin, + scope, + payload, + }); + } + } + } + Ok(entries) + } + /// Validates persisted state before any plugin can apply it. /// /// These limits bound state accepted by plugins. The session reader separately caps each @@ -502,7 +676,8 @@ impl StoredSessionStateSnapshot { StoredSessionStateSnapshotInner::Supported(snapshot) => { validate_supported_snapshot(snapshot) } - StoredSessionStateSnapshotInner::Unsupported { .. } => { + StoredSessionStateSnapshotInner::Unsupported { .. } + | StoredSessionStateSnapshotInner::Malformed { .. } => { Err(self.unsupported_schema_error()) } } @@ -531,29 +706,10 @@ impl StoredSessionStateSnapshot { }); }; let Some(state) = scopes.get(scope.as_str()) else { - let Some(stored_scope) = scopes.keys().next() else { - return Err(SessionStateError::MissingPluginState { - plugin: plugin.to_owned(), - }); - }; - let Some(found) = StoredStateScope::from_stored(stored_scope) else { - return Err(SessionStateError::InvalidStoredScope { - plugin: plugin.to_owned(), - scope: stored_scope.clone(), - }); - }; - return Err(SessionStateError::ScopeMismatch { - plugin: plugin.to_owned(), - found, - expected: scope, - }); + return Ok(None); }; + validate_stored_plugin_state_size(plugin, state)?; let payload = state.payload(); - let measured = match payload { - Some(payload) => payload, - None => &state.raw, - }; - validate_plugin_state_size(plugin, measured)?; let (Some(found), Some(payload)) = (state.schema_version, payload) else { return Err(SessionStateError::InvalidPluginState { plugin: plugin.to_owned(), @@ -576,6 +732,9 @@ impl StoredSessionStateSnapshot { u64::from(snapshot.schema_version) } StoredSessionStateSnapshotInner::Unsupported { schema_version, .. } => *schema_version, + StoredSessionStateSnapshotInner::Malformed { .. } => { + return SessionStateError::InvalidEnvelope; + } }; SessionStateError::UnsupportedSchemaVersion { found, @@ -613,11 +772,7 @@ fn validate_supported_snapshot( continue; }; for state in scopes.values() { - let measured = match state.payload() { - Some(payload) => payload, - None => &state.raw, - }; - validate_plugin_state_size(plugin, measured)?; + validate_stored_plugin_state_size(plugin, state)?; } } let bytes = serde_json::to_vec(snapshot)?.len(); @@ -648,7 +803,37 @@ fn validate_plugin_state_size( plugin: &str, value: &serde_json::Value, ) -> Result<(), SessionStateError> { - let bytes = serde_json::to_vec(value)?.len(); + validate_plugin_state_bytes(plugin, serde_json::to_vec(value)?.len()) +} + +fn validate_stored_plugin_state_size( + plugin: &str, + state: &StoredPluginState, +) -> Result<(), SessionStateError> { + let serde_json::Value::Object(fields) = &state.raw else { + return validate_plugin_state_size(plugin, &state.raw); + }; + let Some(payload) = fields.get("payload") else { + return validate_plugin_state_size(plugin, &state.raw); + }; + let payload_bytes = serde_json::to_vec(payload)?.len(); + let opaque = fields + .iter() + .filter(|(name, _)| { + name.as_str() != "payload" + && (name.as_str() != "schema_version" || state.schema_version.is_none()) + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect::>(); + let opaque_bytes = if opaque.is_empty() { + 0 + } else { + serde_json::to_vec(&opaque)?.len() + }; + validate_plugin_state_bytes(plugin, payload_bytes.saturating_add(opaque_bytes)) +} + +fn validate_plugin_state_bytes(plugin: &str, bytes: usize) -> Result<(), SessionStateError> { if bytes > MAX_PLUGIN_STATE_BYTES { return Err(SessionStateError::PluginStateTooLarge { plugin: plugin.to_owned(), @@ -669,14 +854,27 @@ where let Some(raw) = raw else { return Ok(None); }; - match serde_json::from_value(raw) { + let bytes = serde_json::to_vec(&raw) + .map_err(serde::de::Error::custom)? + .len(); + if bytes > MAX_SESSION_STATE_BYTES { + return Err(serde::de::Error::custom( + SessionStateError::SnapshotTooLarge { + bytes, + maximum: MAX_SESSION_STATE_BYTES, + }, + )); + } + match serde_json::from_value(raw.clone()) { Ok(snapshot) => Ok(Some(snapshot)), Err(error) => { warn!( category = ?error.classify(), - "ignoring malformed session-state snapshot" + "quarantining malformed session-state snapshot" ); - Ok(None) + Ok(Some(StoredSessionStateSnapshot { + inner: StoredSessionStateSnapshotInner::Malformed { raw }, + })) } } } @@ -4717,7 +4915,7 @@ mod tests { } #[test] - fn session_state_snapshot_rejects_scope_mismatch() { + fn session_state_snapshot_absent_requested_scope_is_none() { let mut snapshot = super::StoredSessionStateSnapshot::new(1); snapshot .insert_plugin_state( @@ -4727,26 +4925,260 @@ mod tests { serde_json::json!({ "todos": [] }), ) .unwrap(); + assert_eq!( + snapshot + .plugin_payload_for_apply("todo_write", 1, super::StoredStateScope::Session) + .unwrap(), + None + ); + } + + #[test] + fn session_state_snapshot_enumerates_only_valid_supported_entries() { + let snapshot: super::StoredSessionStateSnapshot = + serde_json::from_value(serde_json::json!({ + "schema_version": 1, + "state_revision": 3, + "plugins": { + "plugin": { + "root": { "schema_version": 1, "payload": { "valid": true } }, + "session": { "schema_version": 2, "payload": { "future": true } }, + "future_scope": { "schema_version": 1, "payload": { "opaque": true } } + }, + "malformed": { + "root": { "schema_version": "invalid", "payload": null } + }, + "malformed_container": null + } + })) + .unwrap(); + + let entries = snapshot.plugin_entries_for_apply(1).unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].plugin, "plugin"); + assert_eq!(entries[0].scope, super::StoredStateScope::Root); + assert_eq!(entries[0].payload, &serde_json::json!({ "valid": true })); + } + + #[test] + fn session_state_snapshot_mutations_preserve_opaque_data() { + let raw = serde_json::json!({ + "schema_version": 1, + "state_revision": 3, + "future_envelope": { "kept": true }, + "plugins": { + "target": { + "root": { "schema_version": 1, "payload": "old", "state_extra": 7 }, + "future_scope": { "opaque": true }, + "session": { "malformed": true } + }, + "malformed_sibling": null, + "unknown_sibling": { "other": [1, 2, 3] } + } + }); + let mut snapshot: super::StoredSessionStateSnapshot = + serde_json::from_value(raw.clone()).unwrap(); + + snapshot.set_state_revision(4).unwrap(); + snapshot + .set_plugin_state( + "target", + 1, + super::StoredStateScope::Root, + serde_json::json!("new"), + ) + .unwrap(); + let after_set = serde_json::to_value(&snapshot).unwrap(); + assert_eq!(after_set["state_revision"], 4); + assert_eq!( + after_set["plugins"]["target"]["root"]["state_extra"], + raw["plugins"]["target"]["root"]["state_extra"] + ); + assert_eq!(after_set["future_envelope"], raw["future_envelope"]); + assert_eq!( + after_set["plugins"]["target"]["future_scope"], + raw["plugins"]["target"]["future_scope"] + ); + assert_eq!( + after_set["plugins"]["target"]["session"], + raw["plugins"]["target"]["session"] + ); + assert_eq!( + after_set["plugins"]["malformed_sibling"], + raw["plugins"]["malformed_sibling"] + ); + assert_eq!( + after_set["plugins"]["unknown_sibling"], + raw["plugins"]["unknown_sibling"] + ); + + snapshot + .remove_plugin_state("target", super::StoredStateScope::Root) + .unwrap(); + let after_remove = serde_json::to_value(snapshot).unwrap(); + assert!(after_remove["plugins"]["target"].get("root").is_none()); + assert_eq!( + after_remove["plugins"]["target"]["future_scope"], + raw["plugins"]["target"]["future_scope"] + ); + assert_eq!( + after_remove["plugins"]["target"]["session"], + raw["plugins"]["target"]["session"] + ); + } + + #[test] + fn session_state_snapshot_rejects_mutating_malformed_plugin_container() { + let raw = serde_json::json!({ + "schema_version": 1, + "state_revision": 3, + "plugins": {"target": null} + }); + let mut snapshot: super::StoredSessionStateSnapshot = + serde_json::from_value(raw.clone()).unwrap(); + assert!(matches!( - snapshot.plugin_payload_for_apply("todo_write", 1, super::StoredStateScope::Session), - Err(super::SessionStateError::ScopeMismatch { .. }) + snapshot.set_plugin_state( + "target", + 1, + super::StoredStateScope::Root, + serde_json::json!({"new": true}), + ), + Err(super::SessionStateError::InvalidPluginContainer { .. }) )); + assert_eq!(serde_json::to_value(snapshot).unwrap(), raw); + } + + #[test] + fn session_state_snapshot_counts_opaque_state_fields_toward_entry_limit() { + let raw = serde_json::json!({ + "schema_version": 1, + "state_revision": 3, + "plugins": { + "target": { + "root": { + "schema_version": 1, + "payload": null, + "opaque": "x".repeat(super::MAX_PLUGIN_STATE_BYTES) + } + } + } + }); + + assert!(serde_json::from_value::(raw).is_err()); + } + + #[test] + fn session_state_snapshot_counts_malformed_version_as_opaque_state() { + let raw = serde_json::json!({ + "schema_version": 1, + "state_revision": 3, + "plugins": { + "target": { + "root": { + "schema_version": "x".repeat(super::MAX_PLUGIN_STATE_BYTES), + "payload": null + } + } + } + }); + + assert!(serde_json::from_value::(raw).is_err()); + } + + #[test] + fn session_state_snapshot_mutations_are_atomic() { + let mut snapshot = super::StoredSessionStateSnapshot::new(5); + snapshot + .set_plugin_state( + "kept", + 1, + super::StoredStateScope::Root, + serde_json::json!({ "value": true }), + ) + .unwrap(); + let before = serde_json::to_value(&snapshot).unwrap(); + + assert!(matches!( + snapshot.set_state_revision(4), + Err(super::SessionStateError::StateRevisionRegression { .. }) + )); + assert_eq!(serde_json::to_value(&snapshot).unwrap(), before); + + assert!(matches!( + snapshot.set_plugin_state( + "oversized", + 1, + super::StoredStateScope::Session, + Value::String("x".repeat(super::MAX_PLUGIN_STATE_BYTES)), + ), + Err(super::SessionStateError::PluginStateTooLarge { .. }) + )); + assert_eq!(serde_json::to_value(&snapshot).unwrap(), before); + + assert!(matches!( + snapshot.remove_plugin_state("bad/name", super::StoredStateScope::Root), + Err(super::SessionStateError::InvalidPluginName { .. }) + )); + assert_eq!(serde_json::to_value(snapshot).unwrap(), before); + } + + #[test] + fn session_state_snapshot_unsupported_envelope_mutations_fail_unchanged() { + let raw = serde_json::json!({ + "schema_version": 2, + "state_revision": 9, + "opaque": { "kept": true } + }); + let mut snapshot: super::StoredSessionStateSnapshot = + serde_json::from_value(raw.clone()).unwrap(); + + assert!(matches!( + snapshot.plugin_entries_for_apply(1), + Err(super::SessionStateError::UnsupportedSchemaVersion { found: 2, .. }) + )); + assert!(matches!( + snapshot.set_state_revision(10), + Err(super::SessionStateError::UnsupportedSchemaVersion { found: 2, .. }) + )); + assert!(matches!( + snapshot.set_plugin_state( + "plugin", + 1, + super::StoredStateScope::Root, + serde_json::json!(null), + ), + Err(super::SessionStateError::UnsupportedSchemaVersion { found: 2, .. }) + )); + assert!(matches!( + snapshot.remove_plugin_state("plugin", super::StoredStateScope::Root), + Err(super::SessionStateError::UnsupportedSchemaVersion { found: 2, .. }) + )); + assert_eq!(serde_json::to_value(snapshot).unwrap(), raw); } #[test] fn malformed_session_state_does_not_invalidate_session_meta() { - let json = r#"{ - "mode":"plan", - "plan_written":true, - "state_snapshot":{ - "schema_version":1, - "plugins":{"todo_write":{"root":{"schema_version":1}}} + let raw = serde_json::json!({ + "mode": "plan", + "plan_written": true, + "state_snapshot": { + "schema_version": 1, + "plugins": {"todo_write": {"root": {"schema_version": 1}}} } - }"#; - let meta: super::SessionMeta = serde_json::from_str(json).unwrap(); + }); + let meta: super::SessionMeta = serde_json::from_value(raw.clone()).unwrap(); assert_eq!(meta.mode, Some(super::StoredMode::Plan)); assert!(meta.plan_written); - assert!(meta.state_snapshot.is_none()); + let snapshot = meta.state_snapshot.as_ref().unwrap(); + assert!(matches!( + snapshot.validate_for_apply(), + Err(super::SessionStateError::InvalidEnvelope) + )); + assert_eq!( + serde_json::to_value(meta).unwrap()["state_snapshot"], + raw["state_snapshot"] + ); } #[test] @@ -4767,13 +5199,15 @@ mod tests { } #[test] - fn session_state_snapshot_quarantines_missing_plugin_scope() { + fn session_state_snapshot_empty_plugin_scope_is_absent() { let json = r#"{"schema_version":1,"state_revision":1,"plugins":{"p":{}}}"#; let snapshot: super::StoredSessionStateSnapshot = serde_json::from_str(json).unwrap(); - assert!(matches!( - snapshot.plugin_payload_for_apply("p", 1, super::StoredStateScope::Root), - Err(super::SessionStateError::MissingPluginState { .. }) - )); + assert_eq!( + snapshot + .plugin_payload_for_apply("p", 1, super::StoredStateScope::Root) + .unwrap(), + None + ); } #[test] diff --git a/n00n-ui/src/agent/agent_loop.rs b/n00n-ui/src/agent/agent_loop.rs index b5b0f089e..022f5cccc 100644 --- a/n00n-ui/src/agent/agent_loop.rs +++ b/n00n-ui/src/agent/agent_loop.rs @@ -9,7 +9,8 @@ use n00n_agent::permissions::PermissionManager; use n00n_agent::template; use n00n_agent::template::Vars; use n00n_agent::tools::{ - DescriptionContext, FileReadTracker, ToolAudience, ToolFilter, ToolRegistry, ToolsSnapshot, + DescriptionContext, FileReadTracker, SessionIdentity, ToolAudience, ToolFilter, ToolRegistry, + ToolsSnapshot, }; use n00n_agent::{ Agent, AgentConfig, AgentEvent, AgentInput, AgentParams, AgentRunParams, CancelMap, @@ -55,7 +56,7 @@ pub(super) struct AgentLoop { agent_tx: flume::Sender, answer_rx: Arc>>, queue: Arc, - session_id: Option, + identity: Option, timeouts: n00n_providers::Timeouts, openai_options: OpenAiOptions, lua_handle: Option, @@ -144,7 +145,7 @@ impl AgentLoop { agent_tx, answer_rx: Arc::new(async_lock::Mutex::new(answer_rx)), queue, - session_id, + identity: session_id.map(SessionIdentity::root), timeouts, openai_options, lua_handle, @@ -342,7 +343,7 @@ impl AgentLoop { config: Arc::new(self.config.clone()), tool_output_lines: self.tool_output_lines, permissions: Arc::clone(&self.permissions), - session_id: self.session_id.clone(), + identity: self.identity.clone(), timeouts: self.timeouts, file_tracker: Arc::clone(&self.file_tracker), prompt_slots: Arc::new(prompt_slots), From e244668b69b961df464e01c390c8b5cc89e68f30 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 3 Aug 2026 19:45:32 -0400 Subject: [PATCH 06/24] fix(agent,lua,storage,ui): harden agent spawn and session recursion --- n00n-agent/src/agent/run.rs | 32 ++- n00n-agent/src/tools/mod.rs | 2 - n00n-lua/src/api/session.rs | 117 ++++++++- n00n-lua/src/api/util/command.rs | 9 + n00n-lua/src/api/util/ctx.rs | 22 +- n00n-lua/src/lib.rs | 4 +- n00n-lua/src/runtime.rs | 44 ++-- n00n-storage/src/sessions.rs | 33 +++ n00n-ui/src/agent/agent_loop.rs | 68 +++++- n00n-ui/src/agent/mod.rs | 15 +- n00n-ui/src/agent/shared_queue.rs | 35 ++- n00n-ui/src/event_loop.rs | 388 +++++++++++++++++++++++++----- n00n-ui/src/session_lineage.rs | 145 +++++++++-- plugins/lib/n00n/subagent.lua | 26 +- plugins/task/init.lua | 18 +- plugins/team/init.lua | 14 +- 16 files changed, 799 insertions(+), 173 deletions(-) diff --git a/n00n-agent/src/agent/run.rs b/n00n-agent/src/agent/run.rs index a98ea1d71..1de03b48f 100644 --- a/n00n-agent/src/agent/run.rs +++ b/n00n-agent/src/agent/run.rs @@ -240,10 +240,10 @@ impl<'h> Agent<'h> { let supports_tool_examples = params.model.supports_tool_examples(); let fusion_enabled = params.config.fusion.enabled; let admission_scope = params - .session_id + .identity .as_ref() - .map_or_else(crate::tools::ToolAdmission::new_scope, |id| { - Arc::::from(id.to_string()) + .map_or_else(crate::tools::ToolAdmission::new_scope, |identity| { + Arc::::from(identity.session_id().to_string()) }); let fusion_state = if fusion_enabled { Some(FusionState::new()) @@ -368,6 +368,23 @@ impl<'h> Agent<'h> { self.total_cost } + pub async fn run_tool(&self, id: String, name: &str, input: &Value) -> ToolDoneEvent { + let ctx = self.tool_context(); + let done = tool_dispatch::run( + &self.registry, + self.mcp.as_ref(), + id, + name, + input, + &ctx, + tool_dispatch::Emit::Notify, + ) + .await; + self.event_tx + .try_send(AgentEvent::ToolDone(Box::new(done.clone()))); + done + } + /// Runs the agent loop with the given input. /// /// # Errors @@ -881,7 +898,6 @@ impl<'h> Agent<'h> { event_tx: self.event_tx.clone(), mode: Arc::clone(&self.mode), tool_use_id: None, - session_id: self.session_id.clone(), user_response_rx: self.user_response_rx.clone(), loaded_instructions: self.loaded_instructions.clone(), cancel: self.cancel.clone(), @@ -2174,15 +2190,15 @@ mod tests { } #[test] - fn tool_context_preserves_agent_session_id() { + fn tool_context_preserves_agent_session_identity() { let mut history = History::new(Vec::new()); let (mut agent, _event_rx) = make_agent(MockProvider::new(Vec::new()), &mut history); - let session_id = SessionRef::generate(); - agent.session_id = Some(session_id.clone()); + let identity = SessionIdentity::root(SessionRef::generate()); + agent.identity = Some(identity.clone()); let ctx = agent.tool_context(); - assert_eq!(ctx.session_id, Some(session_id)); + assert_eq!(ctx.identity, Some(identity)); } fn make_agent_with_config( diff --git a/n00n-agent/src/tools/mod.rs b/n00n-agent/src/tools/mod.rs index 3c53d2c5a..c44bbc34a 100644 --- a/n00n-agent/src/tools/mod.rs +++ b/n00n-agent/src/tools/mod.rs @@ -346,7 +346,6 @@ pub struct ToolContext { pub event_tx: EventSender, pub mode: Arc, pub tool_use_id: Option, - pub session_id: Option, pub user_response_rx: Option>>>, pub loaded_instructions: LoadedInstructions, pub cancel: CancelToken, @@ -642,7 +641,6 @@ pub fn interpreter_ctx( event_tx: event_tx.clone(), mode: Arc::new(mode.clone()), tool_use_id: None, - session_id: None, user_response_rx, loaded_instructions: LoadedInstructions::new(), cancel, diff --git a/n00n-lua/src/api/session.rs b/n00n-lua/src/api/session.rs index a9ca75b28..b94bd20b6 100644 --- a/n00n-lua/src/api/session.rs +++ b/n00n-lua/src/api/session.rs @@ -5,9 +5,9 @@ use mlua::{Lua, Result as LuaResult, Table, Value}; use n00n_lua_macro::{lua_fn, lua_table}; -use crate::api::util::command::{SessionReply, SessionRequest, UiAction}; -use crate::api::util::convert::json_to_lua; -use crate::runtime::active_session_id; +use crate::api::util::command::{SessionBootstrap, SessionReply, SessionRequest, UiAction}; +use crate::api::util::convert::{json_to_lua, lua_to_json}; +use crate::runtime::active_session_identity; const NO_UI_ERR: &str = "no interactive UI attached"; @@ -118,7 +118,8 @@ async fn delete( /// /// @param opts table? Optional fields: prompt (string) first user message /// to submit right away; focus (boolean) switch the UI to the new session; -/// parent_id (string?) session that spawned this session. +/// parent_id (string?) session that spawned this session; tool (string), +/// input (table), and title (string?) for a direct host-executed bootstrap. /// @return (string|nil, string|nil) New session id, or nil and an error. /// @example /// local id, err = n00n.session.new({ prompt = "fix the tests", focus = true }) @@ -128,14 +129,37 @@ async fn new( #[ctx] tx: Option>, opts: Option
, ) -> LuaResult { - let caller_id = active_session_id(&lua); - let (prompt, focus, parent_id) = match opts { + let caller_id = active_session_identity(&lua).map(|identity| identity.session_id().clone()); + let (prompt, focus, parent_id, tool, input, title) = match opts { Some(opts) => ( opts.get("prompt")?, opts.get("focus").unwrap_or_else(|_| false), opts.get("parent_id")?, + opts.get::>("tool")?, + opts.get::>("input")?, + opts.get::>("title")?, ), - None => (None, false, None), + None => (None, false, None, None, None, None), + }; + let bootstrap = match tool { + Some(tool) => { + if prompt.is_some() { + return Ok(err_pair("direct session bootstrap cannot include prompt")); + } + let input = match input { + Some(input) => input, + None => Value::Table(lua.create_table()?), + }; + Some(SessionBootstrap { + tool, + input: lua_to_json(&lua, &input)?, + title, + }) + } + None if input.is_some() || title.is_some() => { + return Ok(err_pair("session bootstrap input/title requires tool")); + } + None => None, }; roundtrip( lua, @@ -145,6 +169,7 @@ async fn new( focus, parent_id, caller_id, + bootstrap, }, ) .await @@ -170,7 +195,7 @@ async fn prompt( text: String, opts: Option
, ) -> LuaResult { - let caller_id = active_session_id(&lua); + let caller_id = active_session_identity(&lua).map(|identity| identity.session_id().clone()); let (id, steer, control) = match opts { Some(opts) => ( opts.get("session")?, @@ -203,7 +228,8 @@ async fn cancel( #[ctx] tx: Option>, id: String, ) -> LuaResult { - roundtrip(lua, tx, SessionRequest::Cancel { id }).await + let caller_id = active_session_identity(&lua).map(|identity| identity.session_id().clone()); + roundtrip(lua, tx, SessionRequest::Cancel { id, caller_id }).await } /// Renames a session, live or stored. @@ -239,7 +265,7 @@ lua_table! { #[cfg(test)] mod tests { use super::*; - use n00n_agent::cancel::CancelToken; + use n00n_agent::{cancel::CancelToken, tools::SessionIdentity}; use n00n_storage::id::SessionRef; use serde_json::json; use test_case::test_case; @@ -260,7 +286,12 @@ mod tests { let lua = lua_with_session(Some(tx)); let _scope = TaskScope::new( &lua, - TaskCell::new(CancelToken::none(), None, None, Some(caller_id.clone())), + TaskCell::new( + CancelToken::none(), + None, + None, + Some(SessionIdentity::root(caller_id.clone())), + ), ); let expected_caller_id = caller_id.clone(); let checker = std::thread::spawn(move || { @@ -310,6 +341,67 @@ mod tests { assert_eq!(prompt_status, "queued"); } + #[test] + fn direct_bootstrap_forwards_tool_input_title_and_runtime_identity() { + let (tx, rx) = flume::unbounded::(); + let caller_id = SessionRef::generate(); + let lua = lua_with_session(Some(tx)); + let _scope = TaskScope::new( + &lua, + TaskCell::new( + CancelToken::none(), + None, + None, + Some(SessionIdentity::root(caller_id.clone())), + ), + ); + let expected_caller_id = caller_id.clone(); + let checker = std::thread::spawn(move || { + let Ok(UiAction::Session { + req: + SessionRequest::New { + prompt, + focus, + parent_id, + caller_id, + bootstrap: Some(bootstrap), + }, + reply_tx, + }) = rx.recv() + else { + panic!("expected direct bootstrap request"); + }; + assert_eq!(prompt, None); + assert!(!focus); + assert_eq!(parent_id, None); + assert_eq!(caller_id.as_ref(), Some(&expected_caller_id)); + assert_eq!(bootstrap.tool, "task"); + assert_eq!( + bootstrap.input, + json!({ "prompt": "inspect", "background": false }) + ); + assert_eq!(bootstrap.title.as_deref(), Some("task: inspect")); + reply_tx.send(Ok(json!("child"))).unwrap(); + }); + + let (child_id, error): (String, Option) = smol::block_on( + lua.load( + r#" + return session.new({ + tool = "task", + input = { prompt = "inspect", background = false }, + title = "task: inspect", + }) + "#, + ) + .eval_async(), + ) + .unwrap(); + checker.join().unwrap(); + assert_eq!(child_id, "child"); + assert_eq!(error, None); + } + #[test] fn live_without_ui_returns_error_pair() { let lua = lua_with_session(None); @@ -369,13 +461,14 @@ mod tests { let lua = lua_with_session(Some(tx)); let checker = std::thread::spawn(move || { let Ok(UiAction::Session { - req: SessionRequest::Cancel { id }, + req: SessionRequest::Cancel { id, caller_id }, reply_tx, }) = rx.recv() else { panic!("expected cancel request"); }; assert_eq!(id, "abc"); + assert_eq!(caller_id, None); reply_tx.send(Ok(json!(true))).unwrap(); }); let (val, err): (bool, Option) = diff --git a/n00n-lua/src/api/util/command.rs b/n00n-lua/src/api/util/command.rs index cd0f08c99..3a18c697e 100644 --- a/n00n-lua/src/api/util/command.rs +++ b/n00n-lua/src/api/util/command.rs @@ -404,6 +404,13 @@ pub enum WinCommand { Close, } +#[derive(Debug)] +pub struct SessionBootstrap { + pub tool: String, + pub input: serde_json::Value, + pub title: Option, +} + #[derive(Debug)] pub enum SessionRequest { List, @@ -417,6 +424,7 @@ pub enum SessionRequest { focus: bool, parent_id: Option, caller_id: Option, + bootstrap: Option, }, Prompt { id: Option, @@ -427,6 +435,7 @@ pub enum SessionRequest { }, Cancel { id: String, + caller_id: Option, }, Focus { id: String, diff --git a/n00n-lua/src/api/util/ctx.rs b/n00n-lua/src/api/util/ctx.rs index f6d168545..25cbb9488 100644 --- a/n00n-lua/src/api/util/ctx.rs +++ b/n00n-lua/src/api/util/ctx.rs @@ -10,10 +10,9 @@ use mlua::{LuaSerdeExt, MultiValue, UserData, UserDataMethods, Value as LuaValue use n00n_agent::agent::LoadedInstructions; use n00n_agent::cancel::CancelToken; use n00n_agent::tools::{ - Deadline, FileReadTracker, LocalTools, ToolAudience, ToolContext, ToolLive, + Deadline, FileReadTracker, LocalTools, SessionIdentity, ToolAudience, ToolContext, ToolLive, }; use n00n_config::{AgentConfig, ToolOutputLines}; -use n00n_storage::id::SessionRef; use crate::api::tool::ToolCallReply; use crate::api::ui::buf::BufHandle; @@ -114,7 +113,7 @@ enum Caps { config: Arc, workflow: bool, audience: ToolAudience, - session_id: Option, + identity: Option, }, Restore { state: Option, @@ -150,7 +149,7 @@ impl LuaCtx { config: Arc::clone(&ctx.config), workflow: ctx.workflow, audience: ctx.audience, - session_id: ctx.session_id.clone(), + identity: ctx.identity.clone(), }, ) } @@ -177,10 +176,10 @@ impl LuaCtx { } } - pub(crate) fn session_id(&self) -> Option { + pub(crate) fn session_identity(&self) -> Option { match &self.caps { - Caps::Handler { agent, .. } => agent.session_id.clone(), - Caps::Start { session_id, .. } => session_id.clone(), + Caps::Handler { agent, .. } => agent.identity.clone(), + Caps::Start { identity, .. } => identity.clone(), Caps::Restore { .. } => None, } } @@ -589,7 +588,6 @@ mod tests { fn populated_ctx() -> ToolContext { let mut ctx = stub_ctx_with(&AgentMode::Build, None, Some(TOOL_USE_ID)); - ctx.session_id = Some(n00n_storage::id::SessionRef::generate()); ctx.deadline = Deadline::after(Duration::from_mins(1)); ctx.tool_output_lines = ToolOutputLines { bash: 999, @@ -610,12 +608,12 @@ mod tests { } #[test] - fn agent_context_keeps_tool_use_id_and_session_id_and_resets_per_call_state() { + fn agent_context_keeps_tool_use_id_and_identity_and_resets_per_call_state() { let ctx = populated_ctx(); - let expected_session_id = ctx.session_id.clone(); + let expected_identity = ctx.identity.clone(); let agent = AgentContext::from(&ctx); assert_eq!(agent.tool_use_id.as_deref(), Some(TOOL_USE_ID)); - assert_eq!(agent.session_id, expected_session_id); + assert_eq!(agent.identity, expected_identity); assert!(matches!(agent.deadline, Deadline::None)); assert_eq!(agent.tool_output_lines, ToolOutputLines::default()); assert!(agent.local_tools.is_empty()); @@ -636,7 +634,7 @@ mod tests { ); let inner = agent.to_tool_context(); assert_eq!(inner.tool_use_id, None); - assert_eq!(inner.session_id, agent.session_id); + assert_eq!(inner.identity, agent.identity); assert!(inner.live_sink.is_none(), "sink must not be inherited"); assert_eq!(agent.tool_use_id.as_deref(), Some(TOOL_USE_ID)); } diff --git a/n00n-lua/src/lib.rs b/n00n-lua/src/lib.rs index b9e424fd5..4238cde31 100644 --- a/n00n-lua/src/lib.rs +++ b/n00n-lua/src/lib.rs @@ -16,8 +16,8 @@ pub use api::keymap::{KeymapEntry, KeymapReader, KeymapSnapshot}; pub use api::options::{OptionSpec, OptionType, PluginOptionSpecs}; pub use api::util::command::{ Anchor, Axis, Border, Dimension, Edge, FloatConfig, FloatConfigPatch, HintReader, HintSnapshot, - LuaCommandInfo, LuaCommandReader, SessionReply, SessionRequest, Split, TitlePos, UiAction, - WinCommand, WinEvent, + LuaCommandInfo, LuaCommandReader, SessionBootstrap, SessionReply, SessionRequest, Split, + TitlePos, UiAction, WinCommand, WinEvent, }; pub use docs::{DocKind, FnDoc, ModuleDoc, ParamDoc, api_docs}; pub use error::PluginError; diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index 4668c47d5..a2191826a 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -16,13 +16,13 @@ use n00n_agent::cancel::CancelToken; use n00n_agent::prompt::{PromptId, ResolvedSlots, Slot, SlotEntry}; use n00n_agent::tools::tool_search::{LoadNamespace, ToolSearch}; use n00n_agent::tools::{ - HeaderResult, PermissionScopes, RegistryError, Tool, ToolLive, ToolRegistry, ToolSource, + HeaderResult, PermissionScopes, RegistryError, SessionIdentity, Tool, ToolLive, ToolRegistry, + ToolSource, }; use n00n_agent::{BufferSnapshot, SharedBuf, SnapshotLine, SnapshotSpan, SpanStyle}; use serde_json::Value; use n00n_config::RawConfig; -use n00n_storage::id::SessionRef; use crate::api::autocmd::AutocmdStore; use crate::api::create_n00n_global; @@ -333,7 +333,7 @@ pub(crate) struct TaskCell { /// Forwards live bufs and annotations to a parent /// `n00n.agent.call_tool(on_live_buf/on_annotation)`. pub(crate) live_sink: Option>, - pub(crate) session_id: Option, + pub(crate) identity: Option, /// When `Some`, `n00n.async.run` tasks queue here instead of the global /// `SpawnQueue` so restore can run them inline before snapshotting. pub(crate) inline_spawn: Option>, @@ -351,7 +351,7 @@ impl TaskCell { cancel: CancelToken, deadline: Option, live: Option, - session_id: Option, + identity: Option, ) -> Self { Self { cancel, @@ -362,7 +362,7 @@ impl TaskCell { live, root_buf: None, live_sink: None, - session_id, + identity, inline_spawn: None, bufs_claim: Weak::new(), async_tasks: Cell::new(0), @@ -630,9 +630,9 @@ pub(crate) fn active_task(lua: &Lua) -> TaskHandle { ) } -pub(crate) fn active_session_id(lua: &Lua) -> Option { +pub(crate) fn active_session_identity(lua: &Lua) -> Option { let handle = lua.app_data_ref::()?; - lock_cell(&handle).session_id.clone() + lock_cell(&handle).identity.clone() } pub(crate) fn with_task_jobs(lua: &Lua, f: impl FnOnce(&mut JobStore) -> R) -> R { @@ -651,14 +651,14 @@ pub(crate) fn with_live_ctx(lua: &Lua, f: impl FnOnce(&LiveCtx) -> R) -> Opti pub(crate) fn enqueue_async_task(lua: &Lua, work_fn: RegistryKey) -> Result<(), mlua::Error> { let handle = lua.app_data_ref::(); - let (cancel, live_ctx, parent_deadline, session_id) = match &handle { + let (cancel, live_ctx, parent_deadline, identity) = match &handle { Some(h) => { let cell = lock_cell(h); ( cell.cancel.clone(), cell.live.clone(), cell.deadline.get(), - cell.session_id.clone(), + cell.identity.clone(), ) } None => (CancelToken::none(), None, None, None), @@ -672,7 +672,7 @@ pub(crate) fn enqueue_async_task(lua: &Lua, work_fn: RegistryKey) -> Result<(), cancel, deadline, live_ctx, - session_id, + identity, owner: None, parent: None, }; @@ -871,7 +871,7 @@ pub(crate) struct PendingAsyncTask { pub cancel: CancelToken, pub deadline: Option, pub live_ctx: Option, - pub session_id: Option, + pub identity: Option, pub owner: Option>, /// Parent task that spawned this `noon.async.run` task, if any. /// Used to decrement the parent's `async_tasks` counter on completion. @@ -979,7 +979,7 @@ fn spawn_async_task( task.cancel.clone(), task.deadline, task.live_ctx.clone(), - task.session_id.clone(), + task.identity.clone(), ), ); let result = scope @@ -2317,10 +2317,10 @@ async fn run_tool_start( ctx: Box, ) { let _context_liveness = ContextLivenessGuard(ctx.context_liveness()); - let session_id = ctx.session_id(); + let identity = ctx.session_identity(); let scope = TaskScope::new( lua, - TaskCell::new(ctx.cancel.clone(), None, Some(live), session_id), + TaskCell::new(ctx.cancel.clone(), None, Some(live), identity), ); let run = async { let input_lua = json_to_lua(lua, &input)?; @@ -2373,7 +2373,7 @@ async fn run_tool_call( let (finish_tx, finish_rx) = flume::bounded::(1); ctx.finish_tx = Some(finish_tx); let cancel = ctx.cancel.clone(); - let session_id = ctx.session_id(); + let identity = ctx.session_identity(); let input_lua = match json_to_lua(&lua, &input) { Ok(v) => v, @@ -2390,7 +2390,7 @@ async fn run_tool_call( Err(e) => return ToolCallReply::err(strip_traceback(&e)), }; let live_id = live.as_ref().map(|l| l.tool_use_id.clone()); - let mut cell = TaskCell::new(cancel, deadline, live, session_id.clone()); + let mut cell = TaskCell::new(cancel, deadline, live, identity.clone()); cell.live_sink = live_sink; let scope = TaskScope::new(&lua, cell); let handle = Arc::clone(scope.handle()); @@ -2473,7 +2473,7 @@ async fn run_tool_call( // A fresh cell, because the original's cancel token and // deadline are stale: the watchdog interrupt would use them to // kill warm clicks. - let mut cell = TaskCell::new(CancelToken::none(), None, None, session_id); + let mut cell = TaskCell::new(CancelToken::none(), None, None, identity); cell.root_buf = Some(root); let mut warm = warm_tools.borrow_mut(); warm.push_back(WarmTool { @@ -3228,18 +3228,18 @@ mod tests { } #[test] - fn enqueue_async_task_inherits_session_id() { + fn enqueue_async_task_inherits_session_identity() { let lua = enqueue_test_lua(); - let session_id = n00n_storage::id::SessionRef::generate(); + let identity = SessionIdentity::root(n00n_storage::id::SessionRef::generate()); let _h = set_active( &lua, - TaskCell::new(CancelToken::none(), None, None, Some(session_id.clone())), + TaskCell::new(CancelToken::none(), None, None, Some(identity.clone())), ); enqueue_async_task(&lua, enqueue_dummy(&lua)).unwrap(); let queue = lua.app_data_ref::().unwrap(); let queued = queue.rx.try_recv().unwrap(); - assert_eq!(queued.session_id, Some(session_id)); + assert_eq!(queued.identity, Some(identity)); } #[test] @@ -3334,7 +3334,7 @@ mod tests { cancel, deadline, live_ctx: None, - session_id: None, + identity: None, owner: None, parent: None, } diff --git a/n00n-storage/src/sessions.rs b/n00n-storage/src/sessions.rs index 34da3876c..cbab15747 100644 --- a/n00n-storage/src/sessions.rs +++ b/n00n-storage/src/sessions.rs @@ -154,6 +154,35 @@ pub enum StoredDelivery { Immediate, } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum StoredSessionLifecycle { + Queued, + Bootstrapping, + Running, + WaitingInput, + Paused, + Succeeded, + Failed, + Cancelled, + #[default] + Idle, +} + +impl StoredSessionLifecycle { + #[must_use] + pub fn is_active(self) -> bool { + matches!( + self, + Self::Queued | Self::Bootstrapping | Self::Running | Self::WaitingInput + ) + } + + #[must_use] + pub fn is_idle(&self) -> bool { + matches!(self, Self::Idle) + } +} #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct StoredQueuedMessage { pub text: String, @@ -883,6 +912,10 @@ where pub struct SessionMeta { #[serde(default, skip_serializing_if = "Option::is_none")] pub parent_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub root_session_id: Option, + #[serde(default, skip_serializing_if = "StoredSessionLifecycle::is_idle")] + pub lifecycle: StoredSessionLifecycle, #[serde(default)] pub mode: Option, #[serde(default)] diff --git a/n00n-ui/src/agent/agent_loop.rs b/n00n-ui/src/agent/agent_loop.rs index 022f5cccc..5d59fa9e8 100644 --- a/n00n-ui/src/agent/agent_loop.rs +++ b/n00n-ui/src/agent/agent_loop.rs @@ -19,7 +19,6 @@ use n00n_agent::{ }; use n00n_lua::EventHandle; use n00n_providers::{AgentError, Message, Model, OpenAiOptions, System, TokenUsage}; -use n00n_storage::id::SessionRef; use n00n_storage::sessions::TranscriptEntry; use serde_json::Value; use tracing::{error, warn}; @@ -92,7 +91,7 @@ pub(super) struct AgentLoopInit { pub(super) queue: Arc, pub(super) cancel_map: Arc, pub(super) init_cancel: CancelToken, - pub(super) session_id: Option, + pub(super) identity: Option, pub(super) timeouts: n00n_providers::Timeouts, pub(super) openai_options: OpenAiOptions, pub(super) lua_handle: Option, @@ -118,7 +117,7 @@ impl AgentLoop { queue, cancel_map, init_cancel, - session_id, + identity, timeouts, openai_options, lua_handle, @@ -145,7 +144,7 @@ impl AgentLoop { agent_tx, answer_rx: Arc::new(async_lock::Mutex::new(answer_rx)), queue, - identity: session_id.map(SessionIdentity::root), + identity, timeouts, openai_options, lua_handle, @@ -195,6 +194,10 @@ impl AgentLoop { .await } QueueItem::Compact { .. } => self.do_compact(&event_tx).await, + QueueItem::DirectTool { tool, input, .. } => { + self.do_direct_tool_run(&event_tx, run_id, &tool, &input) + .await + } }; if let Err(e) = result { @@ -224,6 +227,63 @@ impl AgentLoop { !self.init_cancel.is_cancelled() } + async fn do_direct_tool_run( + &mut self, + event_tx: &EventSender, + run_id: u64, + tool: &str, + input: &serde_json::Value, + ) -> Result<(), AgentError> { + let slot = self.model_slot.load(); + self.rebuild_tools(&slot.model, false); + let (trigger, cancel) = CancelToken::new(); + self.set_cancel_trigger(run_id, trigger); + let agent = Agent::new( + AgentParams { + provider: Arc::clone(&slot.provider), + model: slot.model.clone(), + config: Arc::new(self.config.clone()), + tool_output_lines: self.tool_output_lines, + permissions: Arc::clone(&self.permissions), + identity: self.identity.clone(), + timeouts: self.timeouts, + openai_options: self.openai_options, + file_tracker: Arc::clone(&self.file_tracker), + prompt_slots: Arc::new(n00n_agent::prompt::ResolvedSlots::default()), + subagent_cancels: Arc::clone(&self.subagent_cancels), + registry: Arc::clone(ToolRegistry::global_arc()), + audience: ToolAudience::MAIN, + }, + AgentRunParams { + history: &mut self.history, + system: System::default(), + event_tx: event_tx.clone(), + tools: self.tools.clone(), + tool_filter: self.tool_filter.clone(), + }, + ) + .with_cancel(cancel) + .with_mcp(self.mcp.clone()); + let result = agent + .run_tool(format!("bootstrap-{run_id}"), tool, input) + .await; + drop(agent); + self.clear_cancel_trigger(run_id); + if result.is_error { + return Err(AgentError::Tool { + tool: tool.to_owned(), + message: result.output.as_text(), + }); + } + event_tx.send(AgentEvent::Done { + usage: TokenUsage::default(), + num_turns: 1, + stop_reason: None, + fusion: None, + })?; + Ok(()) + } + async fn do_compact(&mut self, event_tx: &EventSender) -> Result<(), AgentError> { let slot = self.model_slot.load(); let (provider, model) = agent::resolve_compaction_model( diff --git a/n00n-ui/src/agent/mod.rs b/n00n-ui/src/agent/mod.rs index 6f6d8cc3c..6bc361d5d 100644 --- a/n00n-ui/src/agent/mod.rs +++ b/n00n-ui/src/agent/mod.rs @@ -13,10 +13,9 @@ use arc_swap::ArcSwap; use n00n_agent::permissions::PermissionManager; use n00n_agent::{ AgentConfig, CancelMap, CancelToken, Envelope, McpCommand, McpConfigErrors, McpHandle, - McpSnapshotReader, ToolOutput, ToolOutputLines, + McpSnapshotReader, ToolOutput, ToolOutputLines, tools::SessionIdentity, }; use n00n_lua::EventHandle; -use n00n_storage::id::SessionRef; use n00n_storage::sessions::TranscriptEntry; use self::cancel_map::new_run_cancel_map; @@ -55,6 +54,7 @@ pub(crate) struct AgentHandles { pub(crate) queue: QueueSender, pub(crate) timeouts: n00n_providers::Timeouts, openai_options: OpenAiOptions, + identity: Option, task: smol::Task<()>, } @@ -70,7 +70,7 @@ impl AgentHandles { config: AgentConfig, tool_output_lines: ToolOutputLines, permissions: &Arc, - session_id: Option, + identity: Option, timeouts: n00n_providers::Timeouts, openai_options: OpenAiOptions, lua_handle: Option, @@ -87,7 +87,7 @@ impl AgentHandles { permissions, mcp_handle, mcp_config_errors, - session_id, + identity, timeouts, openai_options, lua_handle, @@ -152,7 +152,7 @@ impl AgentHandles { permissions, self.mcp_handle.clone(), self.mcp_config_errors.clone(), - Some(SessionRef::from(app.state.session.id)), + self.identity.clone(), self.timeouts, self.openai_options, lua_handle, @@ -216,7 +216,7 @@ fn spawn_agent_internal( permissions: &Arc, mcp_handle: Option, mcp_config_errors: McpConfigErrors, - session_id: Option, + identity: Option, timeouts: n00n_providers::Timeouts, openai_options: OpenAiOptions, lua_handle: Option, @@ -270,7 +270,7 @@ fn spawn_agent_internal( queue: queue_rx, cancel_map, init_cancel, - session_id, + identity: identity.clone(), timeouts, openai_options, lua_handle, @@ -293,6 +293,7 @@ fn spawn_agent_internal( queue: queue_tx, timeouts, openai_options, + identity, task, } } diff --git a/n00n-ui/src/agent/shared_queue.rs b/n00n-ui/src/agent/shared_queue.rs index 942395981..194e172d8 100644 --- a/n00n-ui/src/agent/shared_queue.rs +++ b/n00n-ui/src/agent/shared_queue.rs @@ -73,12 +73,19 @@ pub(crate) enum QueueItem { Compact { run_id: u64, }, + DirectTool { + run_id: u64, + tool: String, + input: serde_json::Value, + }, } impl QueueItem { pub(crate) fn run_id(&self) -> u64 { match self { - Self::Message { run_id, .. } | Self::Compact { run_id } => *run_id, + Self::Message { run_id, .. } + | Self::Compact { run_id } + | Self::DirectTool { run_id, .. } => *run_id, } } @@ -99,13 +106,18 @@ impl QueueItem { .fg .unwrap_or_else(|| theme::current().foreground), }, + Self::DirectTool { tool, .. } => QueueEntry { + text: Cow::Owned(tool.clone()), + color: theme::current().foreground, + }, } } - fn into_extracted_command(self) -> ExtractedCommand { + fn into_extracted_command(self) -> Option { match self { - Self::Message { input, run_id, .. } => ExtractedCommand::Interrupt(input, run_id), - Self::Compact { run_id } => ExtractedCommand::Compact(run_id), + Self::Message { input, run_id, .. } => Some(ExtractedCommand::Interrupt(input, run_id)), + Self::Compact { run_id } => Some(ExtractedCommand::Compact(run_id)), + Self::DirectTool { .. } => None, } } @@ -116,13 +128,14 @@ impl QueueItem { match self { Self::Message { displayed, .. } => !displayed, Self::Compact { .. } => true, + Self::DirectTool { .. } => false, } } fn is_ready(&self) -> bool { match self { Self::Message { ready, .. } => ready.load(Ordering::Acquire), - Self::Compact { .. } => true, + Self::Compact { .. } | Self::DirectTool { .. } => true, } } @@ -168,7 +181,7 @@ impl QueueSender { let mut items = lock(&self.items); let submission_id = match &entry { QueueItem::Message { submission_id, .. } => *submission_id, - QueueItem::Compact { .. } => return, + QueueItem::Compact { .. } | QueueItem::DirectTool { .. } => return, }; if items.iter().any(|item| { matches!(item, QueueItem::Message { submission_id: id, .. } if *id == submission_id) @@ -279,7 +292,7 @@ impl QueueSender { .filter(|item| item.visible_in_panel()) .filter_map(|item| match item { QueueItem::Message { text, .. } => Some(text.clone()), - QueueItem::Compact { .. } => None, + QueueItem::Compact { .. } | QueueItem::DirectTool { .. } => None, }) .collect() } @@ -291,7 +304,7 @@ impl QueueSender { QueueItem::Message { input, delivery, .. } => Some((input.clone(), *delivery)), - QueueItem::Compact { .. } => None, + QueueItem::Compact { .. } | QueueItem::DirectTool { .. } => None, }) .collect() } @@ -324,6 +337,7 @@ impl QueueReceiver { delivery: Delivery::TurnEnd, .. } | QueueItem::Compact { .. } + | QueueItem::DirectTool { .. } ) { None @@ -356,9 +370,12 @@ impl InterruptSource for QueueReceiver { Delivery::Immediate => Some(index), }, QueueItem::Compact { .. } => (point == InterruptPoint::Safe).then_some(index), + QueueItem::DirectTool { .. } => None, } })?; - items.remove(index).map(QueueItem::into_extracted_command) + items + .remove(index) + .and_then(QueueItem::into_extracted_command) } } diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index 325d5130d..02855601d 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -7,6 +7,7 @@ //! waits on every event source at once and wakes the moment a plugin action, //! agent event, or keypress arrives instead of sleeping in `event::poll`. +use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -20,7 +21,9 @@ use crossterm::event::{ }; use n00n_agent::command::CustomCommand; use n00n_agent::permissions::PermissionManager; -use n00n_agent::{AgentConfig, CancelToken, McpCommand, McpConfigErrors, McpHandle, mcp}; +use n00n_agent::{ + AgentConfig, CancelToken, McpCommand, McpConfigErrors, McpHandle, mcp, tools::SessionIdentity, +}; use n00n_config::UiConfig; use n00n_lua::{ EventHandle, HintReader, KeymapReader, LuaCommandReader, SessionReply, SessionRequest, UiAction, @@ -34,7 +37,9 @@ use n00n_providers::{ContentBlock, Message, Model, OpenAiOptions}; use n00n_storage::StateDir; use n00n_storage::StorageError; use n00n_storage::id::{SessionRef, n00nId, n00nIdParseError}; -use n00n_storage::sessions::{SessionError, TranscriptEntry, normalize_title}; +use n00n_storage::sessions::{ + SessionError, StoredSessionLifecycle, TranscriptEntry, normalize_title, +}; use serde_json::{Value, json}; use tracing::warn; @@ -48,7 +53,7 @@ use crate::components::{ Action, DisplayMessage, DisplayRole, ExitRequest, Status, SubmissionDispatch, }; use crate::input::InputReader; -use crate::session_lineage::{LineageLimits, LiveSession, SessionLineageGuard}; +use crate::session_lineage::{LineageError, LineageLimits, LiveSession, SessionLineageGuard}; use crate::color_compat; use crate::storage_writer::StorageWriter; @@ -128,6 +133,40 @@ fn parse_session_id(id: &str) -> Result { id.parse().map_err(|e: n00nIdParseError| e.to_string()) } +fn caller_session_id(caller: Option) -> Result { + caller + .map(|session| session.id()) + .ok_or_else(|| "authoritative caller session identity is unavailable".to_owned()) +} + +fn live_session(session: &AppSession) -> std::result::Result { + let root_session_id = match (session.meta.parent_id, session.meta.root_session_id) { + (Some(_), None) => return Err(LineageError::MissingRoot(session.id)), + (_, Some(root_session_id)) => root_session_id, + (None, None) => session.id, + }; + Ok(LiveSession { + id: session.id, + root_session_id, + parent_id: session.meta.parent_id, + runtime_present: true, + execution_active: session.meta.lifecycle.is_active(), + }) +} + +fn session_identity(session: &AppSession) -> std::result::Result { + let live = live_session(session)?; + let session_id = SessionRef::from(live.id); + if live.id == live.root_session_id { + Ok(SessionIdentity::root(session_id)) + } else { + Ok(SessionIdentity::child( + session_id, + SessionRef::from(live.root_session_id), + )) + } +} + fn paused_team_run(history: &[Message]) -> Option { let (user_index, last_user) = history .iter() @@ -215,8 +254,15 @@ struct SpawnCtx { } impl SpawnCtx { - fn spawn_runtime(&self, session: AppSession) -> SessionRuntime { + fn spawn_runtime(&self, session: AppSession) -> Result { let resumed = crate::app::session_has_content(&session); + let identity = session_identity(&session) + .map_err(|error| eyre!("invalid session identity: {error}"))?; + if let Some(handle) = &self.lua_event_handle { + handle + .hydrate_state(&identity, session.meta.state_snapshot.clone()) + .map_err(|error| eyre!("failed to hydrate plugin session state: {error}"))?; + } let permissions = Arc::new(self.permissions.fork()); let initial_plan_path = session.meta.plan_path.as_ref().map(PathBuf::from); let handles = AgentHandles::spawn( @@ -227,7 +273,7 @@ impl SpawnCtx { self.config.clone(), self.ui_config.tool_output_lines, &permissions, - Some(SessionRef::from(session.id)), + Some(identity), self.timeouts, self.openai_options, self.lua_event_handle.clone(), @@ -257,13 +303,13 @@ impl SpawnCtx { restore_session(&mut app, &handles); } let (shell_tx, shell_rx) = flume::unbounded::(); - SessionRuntime { + Ok(SessionRuntime { app, handles, shell_tx, shell_rx, last_status: SessionStatus::Idle, - } + }) } } @@ -489,11 +535,27 @@ impl<'t> EventLoop<'t> { let picker = Arc::new(terminal_image::picker()); + let runtime_ids: HashSet<_> = sessions.iter().map(|session| session.id).collect(); + let mut live_sessions = sessions + .iter() + .map(live_session) + .collect::, _>>() + .map_err(|error| eyre!("invalid live session lineage: {error}"))?; + let stored = AppSession::list(&cwd.to_string_lossy(), &storage) + .map_err(|error| eyre!("failed to reconstruct stored session lineage: {error}"))?; + for summary in stored { + if runtime_ids.contains(&summary.id) { + continue; + } + let session = AppSession::load(summary.id, &storage) + .map_err(|error| eyre!("failed to load stored session lineage node: {error}"))?; + let mut node = live_session(&session) + .map_err(|error| eyre!("invalid stored session lineage: {error}"))?; + node.runtime_present = false; + live_sessions.push(node); + } let lineage = SessionLineageGuard::from_live( - sessions.iter().map(|session| LiveSession { - id: session.id, - parent_id: session.meta.parent_id, - }), + live_sessions, LineageLimits { max_depth: config.max_depth, max_total_descendants: config.max_total_descendants, @@ -526,7 +588,7 @@ impl<'t> EventLoop<'t> { let mut runtimes: Vec = sessions .into_iter() .map(|session| ctx.spawn_runtime(session)) - .collect(); + .collect::>>()?; if runtimes.is_empty() { return Err(eyre!("event loop needs at least one session")); } @@ -724,19 +786,82 @@ impl<'t> EventLoop<'t> { if self.last_save.elapsed() < PERIODIC_SAVE_INTERVAL { return; } - for rt in &mut self.sessions { - if should_save_periodically(&rt.app.status) { - rt.app.save_session(); + for idx in 0..self.sessions.len() { + if should_save_periodically(&self.sessions[idx].app.status) { + if let Err(error) = self.capture_plugin_state(idx) { + warn!(session_id = %self.sessions[idx].id(), error = %error, "failed to capture plugin session state"); + } + self.sessions[idx].app.save_session(); } } self.last_save = Instant::now(); } fn handle_agent(&mut self, idx: usize, envelope: Box) { + let lifecycle = match &envelope.event { + n00n_agent::AgentEvent::Done { .. } => Some(StoredSessionLifecycle::Succeeded), + n00n_agent::AgentEvent::Error { .. } => Some(StoredSessionLifecycle::Failed), + n00n_agent::AgentEvent::PermissionRequest { .. } + | n00n_agent::AgentEvent::AuthRequired + | n00n_agent::AgentEvent::SubagentInputRequired { .. } => { + Some(StoredSessionLifecycle::WaitingInput) + } + n00n_agent::AgentEvent::ToolStart(_) + | n00n_agent::AgentEvent::TextDelta { .. } + | n00n_agent::AgentEvent::ThinkingDelta { .. } => Some(StoredSessionLifecycle::Running), + _ => None, + }; + let capture = matches!( + &envelope.event, + n00n_agent::AgentEvent::Done { .. } + | n00n_agent::AgentEvent::Error { .. } + | n00n_agent::AgentEvent::CompactionDone + ); + if capture && let Err(error) = self.capture_plugin_state(idx) { + warn!(session_id = %self.sessions[idx].id(), error = %error, "failed to capture plugin session state"); + } + let terminal = matches!( + lifecycle, + Some(StoredSessionLifecycle::Succeeded | StoredSessionLifecycle::Failed) + ); let actions = self.sessions[idx].app.update(Msg::Agent(envelope)); + if let Some(lifecycle) = lifecycle { + self.sessions[idx].app.state.session.meta.lifecycle = lifecycle; + if terminal { + let id = self.sessions[idx].id(); + if let Err(error) = self.lineage.set_execution_active(id, false) { + warn!(session_id = %id, error = %error, "failed to release session lineage activity"); + } + self.sessions[idx].app.save_session(); + } + } self.dispatch(idx, actions); } + fn capture_plugin_state(&mut self, idx: usize) -> std::result::Result<(), String> { + let Some(handle) = &self.ctx.lua_event_handle else { + return Ok(()); + }; + let session = &self.sessions[idx].app.state.session; + let identity = session_identity(session).map_err(|error| error.to_string())?; + let revision = match session + .meta + .state_snapshot + .as_ref() + .and_then(n00n_storage::sessions::StoredSessionStateSnapshot::state_revision) + { + Some(revision) => revision + .checked_add(1) + .ok_or_else(|| "plugin state revision exhausted".to_owned())?, + None => 1, + }; + let snapshot = handle + .capture_state(&identity, revision) + .map_err(|error| error.to_string())?; + self.sessions[idx].app.state.session.meta.state_snapshot = Some(snapshot); + Ok(()) + } + fn drain_channels(&mut self) -> Result<()> { // Leftovers beyond the budget are picked up right after the next draw. let mut scheduler = DrainScheduler::default(); @@ -944,63 +1069,174 @@ impl<'t> EventLoop<'t> { prompt, focus, parent_id, - caller_id: _, + caller_id, + bootstrap, } => { - let mut session = { - let slot = self.ctx.model_slot.load(); - let cwd = std::env::current_dir().unwrap_or_else(|_| ".".into()); - AppSession::new(&slot.model.spec(), &cwd.to_string_lossy()) - }; - let parent_id = match parent_id { - Some(id) => match parse_session_id(&id) { - Ok(id) => Some(id), + let reply = (|| { + let caller = caller_session_id(caller_id)?; + let explicit_parent = parent_id.as_deref().map(parse_session_id).transpose()?; + let reservation = self + .lineage + .reserve_new(caller, explicit_parent) + .map_err(|error| error.to_string())?; + let caller_lineage = match self.lineage.lineage(caller) { + Ok(lineage) => lineage, Err(error) => { - let _ = reply_tx.send(Err(error)); - return; + let _ = self.lineage.release(reservation); + return Err(error.to_string()); } - }, - None => None, - }; - session.meta.parent_id = parent_id; - let idx = self.push_runtime(self.ctx.spawn_runtime(session)); - let id = self.sessions[idx].id(); - if let Some(prompt) = prompt { - let _ = self.submit_text(idx, prompt, false, false); - } - if focus { - self.set_focus(idx); - } - let _ = reply_tx.send(Ok(json!(id))); + }; + let mut session = { + let slot = self.ctx.model_slot.load(); + let cwd = std::env::current_dir().unwrap_or_else(|_| ".".into()); + AppSession::new(&slot.model.spec(), &cwd.to_string_lossy()) + }; + session.meta.parent_id = Some(caller); + session.meta.root_session_id = Some(caller_lineage.root); + session.meta.lifecycle = StoredSessionLifecycle::Queued; + if let Some(bootstrap) = &bootstrap + && let Some(title) = &bootstrap.title + { + session.title = normalize_title(title); + } + let runtime = match self.ctx.spawn_runtime(session) { + Ok(runtime) => runtime, + Err(error) => { + let _ = self.lineage.release(reservation); + return Err(error.to_string()); + } + }; + let id = runtime.id(); + if let Err(error) = self.lineage.commit_new(reservation, id) { + runtime.handles.cancel(); + return Err(error.to_string()); + } + let idx = self.push_runtime(runtime); + let start_result = if let Some(bootstrap) = bootstrap { + let run_id = { + let app = &mut self.sessions[idx].app; + app.run_id += 1; + app.status = Status::Streaming; + app.state.session.meta.lifecycle = + StoredSessionLifecycle::Bootstrapping; + app.run_id + }; + self.sessions[idx] + .handles + .queue + .push(QueueItem::DirectTool { + run_id, + tool: bootstrap.tool, + input: bootstrap.input, + }); + Ok(json!("started")) + } else if let Some(prompt) = prompt { + self.submit_text(idx, prompt, false, false) + } else { + self.sessions[idx].app.state.session.meta.lifecycle = + StoredSessionLifecycle::Idle; + let _ = self.lineage.set_execution_active(id, false); + Ok(json!("idle")) + }; + if let Err(error) = start_result { + let runtime = self.remove_runtime(idx); + runtime.handles.cancel(); + let _ = self.lineage.rollback_new(id); + return Err(error); + } + self.sessions[idx].app.save_session(); + if focus { + self.set_focus(idx); + } + Ok(json!(id)) + })(); + let _ = reply_tx.send(reply); } SessionRequest::Prompt { id, text, steer, control, - caller_id: _, + caller_id, } => { - let idx = match id { - None => Ok(self.focused), - Some(id) => parse_session_id(&id).and_then(|id| { - self.position(id) - .ok_or_else(|| format!("{NOT_LIVE_ERR}: {id}")) - }), - }; - let _ = - reply_tx.send(idx.and_then(|idx| self.submit_text(idx, text, steer, control))); - } - SessionRequest::Cancel { id } => { - let reply = parse_session_id(&id).and_then(|id| { + let reply = (|| { + let caller = caller_session_id(caller_id)?; + let explicit_target = id.as_deref().map(parse_session_id).transpose()?; + let target = self + .lineage + .authorize_prompt(caller, explicit_target) + .map_err(|error| error.to_string())?; let idx = self - .position(id) - .ok_or_else(|| format!("{NOT_LIVE_ERR}: {id}"))?; - if SessionStatus::of(&self.sessions[idx].app) == SessionStatus::Idle { - return Err(format!("session is idle: {id}")); + .position(target) + .ok_or_else(|| format!("{NOT_LIVE_ERR}: {target}"))?; + let activated = self + .lineage + .begin_execution(target) + .map_err(|error| error.to_string())?; + match self.submit_text(idx, text, steer, control) { + Ok(state) => { + self.sessions[idx].app.state.session.meta.lifecycle = + StoredSessionLifecycle::Running; + Ok(state) + } + Err(error) => { + if activated { + let _ = self.lineage.set_execution_active(target, false); + } + Err(error) + } + } + })(); + let _ = reply_tx.send(reply); + } + SessionRequest::Cancel { id, caller_id } => { + let reply = (|| { + let caller = caller_session_id(caller_id)?; + let requested = parse_session_id(&id)?; + let target = self + .lineage + .authorize_prompt(caller, Some(requested)) + .map_err(|error| error.to_string())?; + let mut targets = self + .lineage + .descendants_of(target) + .map_err(|error| error.to_string())?; + targets.push(target); + let mut cancelled = false; + for session_id in targets { + let Some(idx) = self.position(session_id) else { + let mut session = AppSession::load(session_id, &self.ctx.storage) + .map_err(|error| error.to_string())?; + cancelled |= session.meta.lifecycle.is_active(); + session.meta.lifecycle = StoredSessionLifecycle::Cancelled; + session.updated_at = n00n_storage::now_epoch(); + self.ctx.storage_writer.send(Box::new(session)); + let _ = self.lineage.set_execution_active(session_id, false); + continue; + }; + if SessionStatus::of(&self.sessions[idx].app) != SessionStatus::Idle + || self.sessions[idx] + .app + .state + .session + .meta + .lifecycle + .is_active() + { + let actions = self.sessions[idx].app.cancel_current_run(); + self.dispatch(idx, actions); + cancelled = true; + } + self.sessions[idx].app.state.session.meta.lifecycle = + StoredSessionLifecycle::Cancelled; + self.sessions[idx].app.save_session(); + let _ = self.lineage.set_execution_active(session_id, false); + } + if !cancelled { + return Err(format!("session is idle: {target}")); } - let actions = self.sessions[idx].app.cancel_current_run(); - self.dispatch(idx, actions); Ok(json!(true)) - }); + })(); let _ = reply_tx.send(reply); } SessionRequest::Focus { id } => { @@ -1068,6 +1304,9 @@ impl<'t> EventLoop<'t> { fn remove_runtime(&mut self, idx: usize) -> SessionRuntime { debug_assert_ne!(idx, self.focused); let rt = self.sessions.remove(idx); + if let Err(error) = self.lineage.remove_runtime(rt.id()) { + warn!(session_id = %rt.id(), error = %error, "failed to remove session runtime from lineage"); + } if idx < self.focused { self.focused -= 1; } @@ -1095,15 +1334,20 @@ impl<'t> EventLoop<'t> { self.set_focus(i); return Ok(()); } - let focused = &mut self.sessions[self.focused]; - if SessionStatus::of(&focused.app) == SessionStatus::Idle && !focused.app.has_content() { - let actions = focused.app.load_session(id); - self.dispatch(self.focused, actions); - return Ok(()); - } let session = AppSession::load(id, &self.ctx.storage) .map_err(|e| format!("Failed to load session: {e}"))?; - let idx = self.push_runtime(self.ctx.spawn_runtime(session)); + let live = live_session(&session).map_err(|error| error.to_string())?; + self.lineage + .activate_runtime(live) + .map_err(|error| error.to_string())?; + let runtime = match self.ctx.spawn_runtime(session) { + Ok(runtime) => runtime, + Err(error) => { + let _ = self.lineage.remove_runtime(id); + return Err(error.to_string()); + } + }; + let idx = self.push_runtime(runtime); self.set_focus(idx); Ok(()) } @@ -1510,6 +1754,11 @@ impl<'t> EventLoop<'t> { for rt in &self.sessions { let _ = rt.handles.cmd_tx.try_send(AgentCommand::CancelAll); } + for idx in 0..self.sessions.len() { + if let Err(error) = self.capture_plugin_state(idx) { + warn!(session_id = %self.sessions[idx].id(), error = %error, "failed to capture plugin session state during shutdown"); + } + } let mut tabs = Vec::with_capacity(self.sessions.len()); let mut agent_tasks = Vec::with_capacity(self.sessions.len()); for rt in self.sessions.drain(..) { @@ -1522,6 +1771,13 @@ impl<'t> EventLoop<'t> { tabs.push(app.state.session); agent_tasks.push(handles.into_task()); } + if let Some(handle) = &self.ctx.lua_event_handle { + for session in &tabs { + if let Err(error) = handle.drop_state_owner(session.id) { + warn!(session_id = %session.id, error = %error, "failed to drop plugin session state owner"); + } + } + } if let Some(ref h) = self.ctx.mcp_handle { smol::block_on(h.shutdown()); } diff --git a/n00n-ui/src/session_lineage.rs b/n00n-ui/src/session_lineage.rs index 6f41d7c14..b6a6c1974 100644 --- a/n00n-ui/src/session_lineage.rs +++ b/n00n-ui/src/session_lineage.rs @@ -13,7 +13,10 @@ pub(crate) struct LineageLimits { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct LiveSession { pub(crate) id: n00nId, + pub(crate) root_session_id: n00nId, pub(crate) parent_id: Option, + pub(crate) runtime_present: bool, + pub(crate) execution_active: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -46,6 +49,16 @@ pub(crate) enum LineageError { UnknownSession(n00nId), #[error("session lineage parent must match caller")] ParentMismatch, + #[error("session lineage root is missing for descendant {0}")] + MissingRoot(n00nId), + #[error("session lineage parent {parent} is missing for {id}")] + MissingParent { id: n00nId, parent: n00nId }, + #[error("session lineage root mismatch for {id}: expected {expected}, found {found}")] + RootMismatch { + id: n00nId, + expected: n00nId, + found: n00nId, + }, #[error("session lineage contains a cycle at {0}")] Cycle(n00nId), #[error("session lineage depth limit exceeded: {limit}")] @@ -68,8 +81,10 @@ pub(crate) enum LineageError { #[derive(Debug, Clone, Copy)] struct SessionNode { + root_session_id: n00nId, parent_id: Option, - active: bool, + runtime_present: bool, + execution_active: bool, } #[derive(Debug, Clone, Copy)] @@ -104,8 +119,10 @@ impl SessionLineageGuard { .insert( session.id, SessionNode { + root_session_id: session.root_session_id, parent_id: session.parent_id, - active: true, + runtime_present: session.runtime_present, + execution_active: session.execution_active, }, ) .is_some() @@ -119,19 +136,26 @@ impl SessionLineageGuard { pub(crate) fn activate_runtime(&mut self, session: LiveSession) -> Result<(), LineageError> { if let Some(existing) = self.sessions.get(&session.id) { - if existing.active { + if existing.runtime_present { return Err(LineageError::DuplicateSession(session.id)); } - if existing.parent_id != session.parent_id { + if existing.parent_id != session.parent_id + || existing.root_session_id != session.root_session_id + { return Err(LineageError::ParentChanged { id: session.id }); } self.sessions .get_mut(&session.id) .ok_or(LineageError::UnknownSession(session.id))? - .active = true; + .runtime_present = true; + self.sessions + .get_mut(&session.id) + .ok_or(LineageError::UnknownSession(session.id))? + .execution_active = session.execution_active; if let Err(error) = self.validate_graph() { if let Some(node) = self.sessions.get_mut(&session.id) { - node.active = false; + node.runtime_present = false; + node.execution_active = false; } return Err(error); } @@ -141,8 +165,10 @@ impl SessionLineageGuard { self.sessions.insert( session.id, SessionNode { + root_session_id: session.root_session_id, parent_id: session.parent_id, - active: true, + runtime_present: true, + execution_active: session.execution_active, }, ); if let Err(error) = self.validate_graph() { @@ -157,16 +183,61 @@ impl SessionLineageGuard { .sessions .get_mut(&id) .ok_or(LineageError::UnknownSession(id))?; - node.active = false; + node.runtime_present = false; + node.execution_active = false; Ok(()) } + pub(crate) fn set_execution_active( + &mut self, + id: n00nId, + active: bool, + ) -> Result { + let node = self + .sessions + .get_mut(&id) + .ok_or(LineageError::UnknownSession(id))?; + if active && !node.runtime_present { + return Err(LineageError::TargetNotLive(id)); + } + let changed = node.execution_active != active; + node.execution_active = active; + Ok(changed) + } + + pub(crate) fn begin_execution(&mut self, id: n00nId) -> Result { + let lineage = self.lineage_for(id)?; + let node = self + .sessions + .get(&id) + .ok_or(LineageError::UnknownSession(id))?; + if !node.runtime_present { + return Err(LineageError::TargetNotLive(id)); + } + if node.execution_active { + return Ok(false); + } + if id != lineage.root { + let counts = self.descendant_counts(lineage.root)?; + if counts.active >= self.limits.max_active_descendants { + return Err(LineageError::ActiveDescendantsExceeded { + limit: self.limits.max_active_descendants, + }); + } + } + self.sessions + .get_mut(&id) + .ok_or(LineageError::UnknownSession(id))? + .execution_active = true; + Ok(true) + } + pub(crate) fn lineage(&self, caller: n00nId) -> Result { let node = self .sessions .get(&caller) .ok_or(LineageError::CallerNotLive(caller))?; - if !node.active { + if !node.runtime_present { return Err(LineageError::CallerNotLive(caller)); } self.lineage_for(caller).map_err(|error| match error { @@ -254,8 +325,10 @@ impl SessionLineageGuard { self.sessions.insert( child_id, SessionNode { + root_session_id: pending.root, parent_id: Some(pending.parent), - active: true, + runtime_present: true, + execution_active: true, }, ); if let Err(error) = self.validate_graph() { @@ -271,6 +344,34 @@ impl SessionLineageGuard { .map(|_| ()) .ok_or(LineageError::UnknownReservation) } + pub(crate) fn rollback_new(&mut self, id: n00nId) -> Result<(), LineageError> { + if !self.sessions.contains_key(&id) { + return Err(LineageError::UnknownSession(id)); + } + if self + .sessions + .values() + .any(|node| node.parent_id == Some(id)) + { + return Err(LineageError::ParentChanged { id }); + } + self.sessions.remove(&id); + Ok(()) + } + + pub(crate) fn descendants_of(&self, parent: n00nId) -> Result, LineageError> { + self.validate_graph()?; + if !self.sessions.contains_key(&parent) { + return Err(LineageError::UnknownSession(parent)); + } + let mut descendants = Vec::new(); + for &id in self.sessions.keys() { + if id != parent && self.path_from(id)?.contains(&parent) { + descendants.push(id); + } + } + Ok(descendants) + } pub(crate) fn authorize_prompt( &self, @@ -284,7 +385,7 @@ impl SessionLineageGuard { .sessions .get(&target) .ok_or(LineageError::UnknownSession(target))?; - if !target_node.active { + if !target_node.runtime_present { return Err(LineageError::TargetNotLive(target)); } let caller_path = self.path_from(caller)?; @@ -307,7 +408,7 @@ impl SessionLineageGuard { continue; } total += 1; - if node.active { + if node.execution_active { active += 1; } } @@ -340,6 +441,13 @@ impl SessionLineageGuard { .last() .copied() .ok_or(LineageError::UnknownSession(id))?; + if node.root_session_id != root { + return Err(LineageError::RootMismatch { + id, + expected: root, + found: node.root_session_id, + }); + } Ok(SessionLineage { caller: id, root, @@ -369,7 +477,10 @@ impl SessionLineageGuard { return Ok(path); }; if !self.sessions.contains_key(&parent) { - return Ok(path); + return Err(LineageError::MissingParent { + id: current, + parent, + }); } current = parent; } @@ -391,7 +502,13 @@ mod tests { } fn session(id: n00nId, parent_id: Option) -> LiveSession { - LiveSession { id, parent_id } + LiveSession { + id, + root_session_id: parent_id.unwrap_or(id), + parent_id, + runtime_present: true, + execution_active: parent_id.is_some(), + } } fn limits( diff --git a/plugins/lib/n00n/subagent.lua b/plugins/lib/n00n/subagent.lua index 065d2ceed..17abd7fcb 100644 --- a/plugins/lib/n00n/subagent.lua +++ b/plugins/lib/n00n/subagent.lua @@ -8,6 +8,26 @@ local route_tier = require("n00n.route_tier").route_tier local usage = require("n00n.usage") local structured_output = require("n00n.structured_output") +local ORCHESTRATION_TOOLS = { "task", "team", "workflow", "agent_control", "batch" } + +local function excluded_tools(opts) + local excluded = {} + local seen = {} + if not opts.allow_orchestration then + for _, name in ipairs(ORCHESTRATION_TOOLS) do + excluded[#excluded + 1] = name + seen[name] = true + end + end + for _, name in ipairs(opts.except_tools or {}) do + if not seen[name] then + excluded[#excluded + 1] = name + seen[name] = true + end + end + return excluded +end + -- Launch a subagent with the given options. -- Returns (result | nil, err, cost, usage, model_spec) -- @@ -25,6 +45,7 @@ local structured_output = require("n00n.structured_output") -- include_mcp: Include MCP tools (default: true) -- only_tools: Optional allowlist of tool names -- except_tools: Optional denylist of tool names +-- allow_orchestration: Expose recursive orchestration tools (default: false) -- system_append: Trusted instruction appended to the system prompt -- local_tools: Additional local tools to register -- preview: ActivityPreview object wrapping sess:prompt (optional) @@ -125,12 +146,13 @@ function M.launch(ctx, opts) end -- Get tool definitions + local excluded = excluded_tools(opts) local tool_defs, tools_err = n00n.agent.tools(ctx, { audience = audience, spec = model_spec, only = opts.only_tools, - except = opts.except_tools, + except = excluded, include_mcp = opts.include_mcp, }) @@ -178,7 +200,7 @@ function M.launch(ctx, opts) thinking = opts.thinking, mode = subagent_type, include_mcp = opts.include_mcp, - except = opts.except_tools, + except = excluded, }) if sess_err then return nil, sess_err, nil, nil, model_spec diff --git a/plugins/task/init.lua b/plugins/task/init.lua index 6caefc803..744c18013 100644 --- a/plugins/task/init.lua +++ b/plugins/task/init.lua @@ -15,6 +15,7 @@ local subagent = require("n00n.subagent") local DONE_NAME = "done" local DONE_DESCRIPTION = "Call when the task is complete with your final answer." local DONE_PROMPT_SUFFIX = "\n\nWhen finished, call the done tool with your final answer." +local ORCHESTRATION_TOOLS = { "task", "team", "workflow", "agent_control", "batch" } local BODY_INDENT_COLS = 4 local MIN_MD_WIDTH = 20 local DEFAULT_OUTPUT_LINES = 5 @@ -81,13 +82,13 @@ local function handler(input, ctx) forwarded[key] = value end forwarded.background = false - local forwarded_json, encode_err = n00n.json.encode(forwarded) - if encode_err then - return { llm_output = "failed to encode task input: " .. tostring(encode_err), is_error = true } - end - local prompt = "Use the task tool now with background=false. Do not only describe this request.\n\n" - .. forwarded_json - local id, err = n00n.session.new({ prompt = prompt, focus = false }) + local title = "task: " .. (input.description or input.prompt or "background task"):sub(1, 60) + local id, err = n00n.session.new({ + tool = "task", + input = forwarded, + title = title, + focus = false, + }) if not id then return { llm_output = err, is_error = true } end @@ -174,6 +175,7 @@ local function handler(input, ctx) output_schema = input.output_schema, preview = preview, activity_label = input.description or "task", + except_tools = ORCHESTRATION_TOOLS, }) if err then return { llm_output = err, is_error = true } @@ -212,6 +214,7 @@ local function handler(input, ctx) local tool_defs, tools_err = n00n.agent.tools(ctx, { audience = audience, spec = model.spec, + except = ORCHESTRATION_TOOLS, }) if tools_err then return { llm_output = tools_err, is_error = true } @@ -244,6 +247,7 @@ local function handler(input, ctx) name = input.description, thinking = input.thinking, mode = subagent_type, + except = ORCHESTRATION_TOOLS, }) if sess_err then return { llm_output = sess_err, is_error = true } diff --git a/plugins/team/init.lua b/plugins/team/init.lua index 8ee75af20..bdee3cc3b 100644 --- a/plugins/team/init.lua +++ b/plugins/team/init.lua @@ -802,15 +802,17 @@ local function run_team(input, ctx) forwarded[key] = value end forwarded.background = false - local prompt = "Use the team tool now. Do not only describe this request.\n\n" .. n00n.json.encode(forwarded) - local id, err = n00n.session.new({ prompt = prompt, focus = false }) + local title = "team: " .. (input.goal or ""):sub(1, 60) + local id, err = n00n.session.new({ + tool = "team", + input = forwarded, + title = title, + focus = false, + }) if not id then return { llm_output = err, is_error = true } end - local title = "team: " .. (input.goal or ""):sub(1, 60) - pcall(function() - n00n.session.set_title({ id = id, title = title }) - end) + return n00n.json.encode({ agent_id = id, status = "started", title = title }) end From 7a631e1c25f2c2b39248d9cb1db2367505435232 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Tue, 4 Aug 2026 15:33:47 -0400 Subject: [PATCH 07/24] fix(agent/lua): defer plugin state store read until after tool normalization --- n00n-lua/src/api/agent.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/n00n-lua/src/api/agent.rs b/n00n-lua/src/api/agent.rs index bceadf932..bf9a7b466 100644 --- a/n00n-lua/src/api/agent.rs +++ b/n00n-lua/src/api/agent.rs @@ -655,11 +655,6 @@ async fn session( opts: Table, ) -> LuaResult> { let agent_ctx = try_pair!(dispatch_ctx(&ctx, "session")).clone(); - let Some(parent_identity) = agent_ctx.identity.clone() else { - return Ok(err_pair("session identity is unavailable")); - }; - let plugin_state_store = try_pair!(ctx.plugin_state_store()); - drop(ctx); let model_spec: Option = opts.get("model_spec")?; let system: Option = opts.get("system")?; let tools_val: Option = opts.get("tools")?; @@ -716,7 +711,8 @@ async fn session( let explicit_tools = tools_val.is_some(); let (mut tools_json, mut tool_filter) = if let Some(val) = tools_val { - let tools = try_pair!(normalize_tool_definitions(lua_to_json(&lua, &val)?)); + let tools = + normalize_tool_definitions(lua_to_json(&lua, &val)?).map_err(mlua::Error::runtime)?; (tools, ToolFilter::All) } else { let vars = n00n_agent::template::Vars::new(); @@ -798,6 +794,12 @@ async fn session( None => agent_ctx.opts.thinking, }; + let Some(parent_identity) = agent_ctx.identity.clone() else { + return Ok(err_pair("session identity is unavailable")); + }; + let plugin_state_store = try_pair!(ctx.plugin_state_store()); + drop(ctx); + let session_id = n00nId::generate(); let child_id = session_id.to_string(); let parent_tool_use_id = child_id.clone(); From a2bd7799fd497be6399d187bf656b00f400de447 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Tue, 4 Aug 2026 19:09:39 -0400 Subject: [PATCH 08/24] fix(lint): resolve clippy warnings for agent spawn recursion PR --- n00n-lua/src/api/session.rs | 4 ++-- n00n-ui/src/app/session.rs | 3 +++ n00n-ui/src/session_lineage.rs | 16 +++++++++++++--- src/cmd/tui_bridge.rs | 10 ++++++++-- 4 files changed, 26 insertions(+), 7 deletions(-) diff --git a/n00n-lua/src/api/session.rs b/n00n-lua/src/api/session.rs index b94bd20b6..7cf1da582 100644 --- a/n00n-lua/src/api/session.rs +++ b/n00n-lua/src/api/session.rs @@ -293,7 +293,7 @@ mod tests { Some(SessionIdentity::root(caller_id.clone())), ), ); - let expected_caller_id = caller_id.clone(); + let expected_caller_id = caller_id; let checker = std::thread::spawn(move || { let Ok(UiAction::Session { req: @@ -355,7 +355,7 @@ mod tests { Some(SessionIdentity::root(caller_id.clone())), ), ); - let expected_caller_id = caller_id.clone(); + let expected_caller_id = caller_id; let checker = std::thread::spawn(move || { let Ok(UiAction::Session { req: diff --git a/n00n-ui/src/app/session.rs b/n00n-ui/src/app/session.rs index e89f045ec..04519d11c 100644 --- a/n00n-ui/src/app/session.rs +++ b/n00n-ui/src/app/session.rs @@ -145,6 +145,7 @@ fn restored_submission( } impl App { + #[allow(dead_code)] pub(crate) fn has_content(&self) -> bool { session_has_content(&self.state.session) } @@ -385,6 +386,7 @@ impl App { ))] } + #[allow(dead_code)] pub(crate) fn apply_loaded_session( &mut self, session: AppSession, @@ -406,6 +408,7 @@ impl App { self.loaded_session_snapshot() } + #[allow(dead_code)] pub(crate) fn load_session(&mut self, session_id: n00nId) -> Vec { let mut session = match AppSession::load(session_id, &self.storage) { Ok(s) => s, diff --git a/n00n-ui/src/session_lineage.rs b/n00n-ui/src/session_lineage.rs index b6a6c1974..cbc52a0be 100644 --- a/n00n-ui/src/session_lineage.rs +++ b/n00n-ui/src/session_lineage.rs @@ -4,6 +4,7 @@ use n00n_storage::id::n00nId; use thiserror::Error; #[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(clippy::struct_field_names)] pub(crate) struct LineageLimits { pub(crate) max_depth: usize, pub(crate) max_total_descendants: usize, @@ -253,7 +254,10 @@ impl SessionLineageGuard { ) -> Result { self.validate_graph()?; let caller_lineage = self.lineage(caller)?; - let parent = explicit_parent.unwrap_or(caller); + let parent = match explicit_parent { + Some(p) => p, + None => caller, + }; if parent != caller { return Err(LineageError::ParentMismatch); } @@ -380,7 +384,10 @@ impl SessionLineageGuard { ) -> Result { self.validate_graph()?; let caller_lineage = self.lineage(caller)?; - let target = explicit_target.unwrap_or(caller); + let target = match explicit_target { + Some(t) => t, + None => caller, + }; let target_node = self .sessions .get(&target) @@ -504,7 +511,10 @@ mod tests { fn session(id: n00nId, parent_id: Option) -> LiveSession { LiveSession { id, - root_session_id: parent_id.unwrap_or(id), + root_session_id: match parent_id { + Some(p) => p, + None => id, + }, parent_id, runtime_present: true, execution_active: parent_id.is_some(), diff --git a/src/cmd/tui_bridge.rs b/src/cmd/tui_bridge.rs index 9a2a62e61..66f00a14c 100644 --- a/src/cmd/tui_bridge.rs +++ b/src/cmd/tui_bridge.rs @@ -181,8 +181,14 @@ fn build_team_resume_prompt(run_info: &Value) -> ControlResult { } fn stop_one(tx: &flume::Sender, id: &str) -> ControlResult<()> { - session_call(tx, SessionRequest::Cancel { id: id.to_owned() }) - .map_err(|e| map_not_found(id, e))?; + session_call( + tx, + SessionRequest::Cancel { + id: id.to_owned(), + caller_id: None, + }, + ) + .map_err(|e| map_not_found(id, e))?; Ok(()) } From cfc94d28ccad76704970b85b8a2d720e15dec9f3 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Tue, 4 Aug 2026 19:14:15 -0400 Subject: [PATCH 09/24] fix(ui): drop unused checkpoint_session after merge --- n00n-ui/src/app/session.rs | 7 ------- 1 file changed, 7 deletions(-) diff --git a/n00n-ui/src/app/session.rs b/n00n-ui/src/app/session.rs index 1758e6849..4cd39c13d 100644 --- a/n00n-ui/src/app/session.rs +++ b/n00n-ui/src/app/session.rs @@ -210,13 +210,6 @@ impl App { ); } - pub(crate) fn checkpoint_session(&mut self) { - let snapshot = self.session_snapshot(); - if session_has_content(&snapshot) { - self.storage_writer.send(Box::new(snapshot)); - } - } - pub(crate) fn hydrate_plugin_state(&mut self) { let Some(handle) = &self.lua_event_handle else { return; From c22892d11fec96607c393dd778243c5283f01cc1 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Wed, 5 Aug 2026 02:25:22 -0400 Subject: [PATCH 10/24] chore: add changelog fragment --- changelog.d/327.fixed.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/327.fixed.md diff --git a/changelog.d/327.fixed.md b/changelog.d/327.fixed.md new file mode 100644 index 000000000..51e0f2b75 --- /dev/null +++ b/changelog.d/327.fixed.md @@ -0,0 +1 @@ +Hardened agent spawn recursion and session lineage handling to prevent runaway subagent loops. \ No newline at end of file From 2a60906735ed56e2ddb6ae902e0e1993debd9ba7 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Wed, 5 Aug 2026 05:03:06 -0400 Subject: [PATCH 11/24] fix(n00n-lua): do not discard plugin state on unload; regenerate docs --- n00n-lua/src/runtime.rs | 1 - n00n-lua/src/state.rs | 14 -------------- site/docs/content/configuration/_index.md | 3 +++ site/docs/content/lua-api/_index.md | 6 +++++- 4 files changed, 8 insertions(+), 16 deletions(-) diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index 9feadf0f0..bbabdf29b 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -2715,7 +2715,6 @@ pub fn spawn( break; } rt.clear_plugin(&plugin); - rt.state.drop_plugin(&plugin); let _ = reply.send(()); } Request::RunCommand { diff --git a/n00n-lua/src/state.rs b/n00n-lua/src/state.rs index 00261b2e0..bb780366a 100644 --- a/n00n-lua/src/state.rs +++ b/n00n-lua/src/state.rs @@ -238,20 +238,6 @@ impl PluginStateStore { identity.session_id != owner && identity.root_session_id != owner }); } - - pub(crate) fn drop_plugin(&self, plugin: &str) { - let mut inner = self.lock(); - let keys = inner - .values - .keys() - .filter(|key| key.plugin == plugin) - .cloned() - .collect::>(); - for key in keys { - inner.values.remove(&key); - inner.managed.insert(key); - } - } } fn validate_replacement( diff --git a/site/docs/content/configuration/_index.md b/site/docs/content/configuration/_index.md index ba533a5e0..e75449f67 100644 --- a/site/docs/content/configuration/_index.md +++ b/site/docs/content/configuration/_index.md @@ -107,6 +107,9 @@ How many lines of output to show per tool in the UI. All values are `usize` with | `max_output_bytes` | usize | `16384` | 1024 | Max tool output size (bytes) | | `max_output_lines` | usize | `500` | 10 | Max tool output lines | | `max_continuation_turns` | u32 | `3` | 1 | Max automatic continuation turns | +| `max_depth` | usize | `4` | 1 | Maximum session lineage depth | +| `max_total_descendants` | usize | `16` | 1 | Maximum total descendants per session lineage root | +| `max_active_descendants` | usize | `8` | 1 | Maximum active descendants per session lineage root | | `compaction_buffer` | u32 \| string | `20%` | - | Context reserved for compaction: token count or percent of the context window (e.g. "20%") | | `mcp_tool_desc_max_chars` | usize | `200` | 10 | Max MCP tool description length (characters) | diff --git a/site/docs/content/lua-api/_index.md b/site/docs/content/lua-api/_index.md index 7a70872d0..3b94bf386 100644 --- a/site/docs/content/lua-api/_index.md +++ b/site/docs/content/lua-api/_index.md @@ -2806,7 +2806,10 @@ Starts a new session in the current project. to submit right away; focus (boolean) switch the UI to the new session; - - `parent_id` (`string?`) session that spawned this session. + - `parent_id` (`string?`) session that spawned this session; tool (string), + + input (table), and title (string?) for a direct host-executed bootstrap. + **Returns:** (`string|nil`, `string|nil`) New session id, or nil and an error. @@ -6013,6 +6016,7 @@ function M.make_local_tool(schema, on_submit) -- include_mcp: Include MCP tools (default: true) -- only_tools: Optional allowlist of tool names -- except_tools: Optional denylist of tool names +-- allow_orchestration: Expose recursive orchestration tools (default: false) -- system_append: Trusted instruction appended to the system prompt -- local_tools: Additional local tools to register -- preview: ActivityPreview object wrapping sess:prompt (optional) From 866824f9c31c1fbe957769b27ed4c4c07cf75547 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Wed, 5 Aug 2026 05:30:18 -0400 Subject: [PATCH 12/24] fix(n00n-lua): drain non-nested tool requests during lifecycle barriers --- n00n-lua/src/runtime.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index bbabdf29b..161b523b5 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -1206,7 +1206,7 @@ async fn drain_runtime( match request_rx.try_recv() { Ok(request) => { if let Some(request) = - spawn_runtime_request(rt, ex, gate, lifecycle, request, true) + spawn_runtime_request(rt, ex, gate, lifecycle, request, false) { deferred.push_back(request); } @@ -1266,7 +1266,7 @@ async fn drain_runtime( RuntimeWake::Spawn(task) => spawn_async_task(&rt.lua, ex, gate, task), RuntimeWake::Request(request) => { if let Some(request) = - spawn_runtime_request(rt, ex, gate, lifecycle, *request, true) + spawn_runtime_request(rt, ex, gate, lifecycle, *request, false) { deferred.push_back(request); } From d26fa6968c46ee4b80d1096fdcd80d5c22b68dbb Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Sun, 9 Aug 2026 22:33:26 -0400 Subject: [PATCH 13/24] fix(orchestration): make background bootstraps durable --- n00n-agent/src/agent/run.rs | 15 +- n00n-docgen/src/gen_config.rs | 1 + n00n-lua/src/api/session.rs | 44 +++- n00n-lua/src/api/util/command.rs | 2 + n00n-storage/src/sessions.rs | 8 + n00n-ui/src/agent/agent_loop.rs | 2 +- n00n-ui/src/agent/shared_queue.rs | 25 ++ n00n-ui/src/app/queue.rs | 14 ++ n00n-ui/src/app/session.rs | 44 +++- n00n-ui/src/event_loop.rs | 282 ++++++++++++++++++---- n00n-ui/src/session_lineage.rs | 42 +++- plugins/lib/n00n/subagent.lua | 8 + plugins/task/init.lua | 4 +- plugins/team/init.lua | 8 +- site/docs/content/configuration/_index.md | 2 + site/docs/content/lua-api/_index.md | 1 + src/cmd/tui_bridge.rs | 34 ++- 17 files changed, 456 insertions(+), 80 deletions(-) diff --git a/n00n-agent/src/agent/run.rs b/n00n-agent/src/agent/run.rs index b204926e6..2e8bb21bc 100644 --- a/n00n-agent/src/agent/run.rs +++ b/n00n-agent/src/agent/run.rs @@ -355,7 +355,16 @@ impl<'h> Agent<'h> { self.total_cost } - pub async fn run_tool(&self, id: String, name: &str, input: &Value) -> ToolDoneEvent { + /// Runs one tool and emits its completion event. + /// + /// # Errors + /// Returns an error when the completion event cannot be delivered. + pub async fn run_tool( + &self, + id: String, + name: &str, + input: &Value, + ) -> Result { let ctx = self.tool_context(); let done = tool_dispatch::run( &self.registry, @@ -368,8 +377,8 @@ impl<'h> Agent<'h> { ) .await; self.event_tx - .try_send(AgentEvent::ToolDone(Box::new(done.clone()))); - done + .send(AgentEvent::ToolDone(Box::new(done.clone())))?; + Ok(done) } /// Runs the agent loop with the given input. diff --git a/n00n-docgen/src/gen_config.rs b/n00n-docgen/src/gen_config.rs index 1761e3595..0e245a078 100644 --- a/n00n-docgen/src/gen_config.rs +++ b/n00n-docgen/src/gen_config.rs @@ -221,6 +221,7 @@ All fields are optional. Typos in field names cause an error right away. write_theme_section(&mut out); write_tool_output_section(&mut out); write_section(&mut out, "[agent]", AgentConfig::FIELDS); + out.push_str("Keep `max_depth` and `max_active_descendants` at or below `max_total_descendants`. n00n reports a configuration error at startup if either value is higher.\n\n"); out.push_str("### `agent.fusion`\n\n"); out.push_str("| Field | Type | Default | Description |\n"); diff --git a/n00n-lua/src/api/session.rs b/n00n-lua/src/api/session.rs index 7cf1da582..41cf7c627 100644 --- a/n00n-lua/src/api/session.rs +++ b/n00n-lua/src/api/session.rs @@ -133,7 +133,17 @@ async fn new( let (prompt, focus, parent_id, tool, input, title) = match opts { Some(opts) => ( opts.get("prompt")?, - opts.get("focus").unwrap_or_else(|_| false), + match opts.get::("focus")? { + Value::Nil => false, + Value::Boolean(focus) => focus, + value => { + return Err(mlua::Error::FromLuaConversionError { + from: value.type_name(), + to: "Boolean".to_owned(), + message: Some("focus must be a boolean".to_owned()), + }); + } + }, opts.get("parent_id")?, opts.get::>("tool")?, opts.get::>("input")?, @@ -213,6 +223,7 @@ async fn prompt( steer, control, caller_id, + host_control: false, }, ) .await @@ -229,7 +240,16 @@ async fn cancel( id: String, ) -> LuaResult { let caller_id = active_session_identity(&lua).map(|identity| identity.session_id().clone()); - roundtrip(lua, tx, SessionRequest::Cancel { id, caller_id }).await + roundtrip( + lua, + tx, + SessionRequest::Cancel { + id, + caller_id, + host_control: false, + }, + ) + .await } /// Renames a session, live or stored. @@ -461,7 +481,12 @@ mod tests { let lua = lua_with_session(Some(tx)); let checker = std::thread::spawn(move || { let Ok(UiAction::Session { - req: SessionRequest::Cancel { id, caller_id }, + req: + SessionRequest::Cancel { + id, + caller_id, + host_control, + }, reply_tx, }) = rx.recv() else { @@ -469,6 +494,7 @@ mod tests { }; assert_eq!(id, "abc"); assert_eq!(caller_id, None); + assert!(!host_control); reply_tx.send(Ok(json!(true))).unwrap(); }); let (val, err): (bool, Option) = @@ -500,6 +526,7 @@ mod tests { steer, control, caller_id, + host_control, }, reply_tx, }) = rx.recv() @@ -511,6 +538,7 @@ mod tests { assert_eq!(steer, expected_steer); assert_eq!(control, expected_control); assert_eq!(caller_id, None); + assert!(!host_control); reply_tx.send(Ok(json!("queued"))).unwrap(); }); let (val, err): (String, Option) = @@ -520,6 +548,16 @@ mod tests { assert_eq!(val, "queued"); } + #[test] + fn new_focus_with_wrong_type_throws() { + let lua = lua_with_session(None); + let result: LuaResult = smol::block_on( + lua.load("return session.new({ focus = 'wrong' })") + .eval_async(), + ); + assert!(result.unwrap_err().to_string().contains("boolean")); + } + #[test] fn set_title_with_wrong_type_throws() { let lua = lua_with_session(None); diff --git a/n00n-lua/src/api/util/command.rs b/n00n-lua/src/api/util/command.rs index 3a18c697e..32bdc043f 100644 --- a/n00n-lua/src/api/util/command.rs +++ b/n00n-lua/src/api/util/command.rs @@ -432,10 +432,12 @@ pub enum SessionRequest { steer: bool, control: bool, caller_id: Option, + host_control: bool, }, Cancel { id: String, caller_id: Option, + host_control: bool, }, Focus { id: String, diff --git a/n00n-storage/src/sessions.rs b/n00n-storage/src/sessions.rs index 502f2dc7f..d948ff975 100644 --- a/n00n-storage/src/sessions.rs +++ b/n00n-storage/src/sessions.rs @@ -207,6 +207,12 @@ pub struct StoredQueuedMessage { pub prompt: Option, } +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct StoredDirectTool { + pub tool: String, + pub input: Value, +} + #[allow(clippy::trivially_copy_pass_by_ref)] // serde skip_serializing_if requires fn(&T) -> bool fn is_default_delivery(delivery: &StoredDelivery) -> bool { *delivery == StoredDelivery::TurnEnd @@ -927,6 +933,8 @@ pub struct SessionMeta { #[serde(default, skip_serializing_if = "Vec::is_empty")] pub queued_submissions: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub queued_direct_tools: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] pub subagents: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub thinking: Option, diff --git a/n00n-ui/src/agent/agent_loop.rs b/n00n-ui/src/agent/agent_loop.rs index 68661bc67..059426e8a 100644 --- a/n00n-ui/src/agent/agent_loop.rs +++ b/n00n-ui/src/agent/agent_loop.rs @@ -266,7 +266,7 @@ impl AgentLoop { .with_mcp(self.mcp.clone()); let result = agent .run_tool(format!("bootstrap-{run_id}"), tool, input) - .await; + .await?; drop(agent); self.clear_cancel_trigger(run_id); if result.is_error { diff --git a/n00n-ui/src/agent/shared_queue.rs b/n00n-ui/src/agent/shared_queue.rs index 194e172d8..48631e788 100644 --- a/n00n-ui/src/agent/shared_queue.rs +++ b/n00n-ui/src/agent/shared_queue.rs @@ -308,6 +308,15 @@ impl QueueSender { }) .collect() } + pub(crate) fn direct_tools(&self) -> Vec<(String, serde_json::Value)> { + lock(&self.items) + .iter() + .filter_map(|item| match item { + QueueItem::DirectTool { tool, input, .. } => Some((tool.clone(), input.clone())), + QueueItem::Message { .. } | QueueItem::Compact { .. } => None, + }) + .collect() + } pub(crate) fn panel_len(&self) -> usize { lock(&self.items) @@ -482,6 +491,22 @@ mod tests { ); } + #[test] + fn direct_tools_are_available_for_persistence() { + let (tx, _rx) = queue(); + tx.push(QueueItem::DirectTool { + run_id: 7, + tool: "task".into(), + input: serde_json::json!({"prompt": "ship"}), + }); + + assert_eq!( + tx.direct_tools(), + vec![("task".into(), serde_json::json!({"prompt": "ship"}))] + ); + assert!(tx.queued_inputs().is_empty()); + } + #[test] fn promoted_steering_is_available_at_safe_point() { let (tx, rx) = queue(); diff --git a/n00n-ui/src/app/queue.rs b/n00n-ui/src/app/queue.rs index b5057e2fa..6a55b1f6f 100644 --- a/n00n-ui/src/app/queue.rs +++ b/n00n-ui/src/app/queue.rs @@ -199,6 +199,20 @@ impl MessageQueue { ) } + pub(crate) fn direct_tools(&self) -> Vec<(String, serde_json::Value)> { + self.shared + .as_ref() + .map_or_else(Vec::new, QueueSender::direct_tools) + } + + pub(crate) fn push_direct_tool(&self, entry: QueueItem) -> bool { + let Some(shared) = &self.shared else { + return false; + }; + shared.push(entry); + true + } + fn clamp_focus(&mut self) { let len = self.len(); self.focus = match self.focus { diff --git a/n00n-ui/src/app/session.rs b/n00n-ui/src/app/session.rs index d36cceac3..ef7f1cf32 100644 --- a/n00n-ui/src/app/session.rs +++ b/n00n-ui/src/app/session.rs @@ -11,18 +11,29 @@ use n00n_agent::{AgentInput, AgentMode, McpPromptRef}; use n00n_providers::{Model, TokenUsage}; use n00n_storage::id::{SessionRef, n00nId}; use n00n_storage::sessions::{ - StoredDelivery, StoredImageMediaType, StoredImageSource, StoredMcpPrompt, StoredMode, - StoredQueuedMessage, StoredSessionStateSnapshot, StoredSubagent, StoredThinking, + StoredDelivery, StoredDirectTool, StoredImageMediaType, StoredImageSource, StoredMcpPrompt, + StoredMode, StoredQueuedMessage, StoredSessionStateSnapshot, StoredSubagent, StoredThinking, }; use crate::AppSession; use super::session_state::{SessionState, stored_to_rules}; use super::{App, Mode, PendingInput, PlanState}; +use crate::agent::shared_queue::QueueItem; use crate::agent::{Delivery, QueuedMessage}; const INITIAL_STATE_REVISION: u64 = 0; +fn plugin_state_identity(session: &AppSession) -> SessionIdentity { + let root_id = session.meta.root_session_id.map_or(session.id, |root| root); + let session_id = SessionRef::from_id(session.id); + if root_id == session.id { + SessionIdentity::root(session_id) + } else { + SessionIdentity::child(session_id, SessionRef::from_id(root_id)) + } +} + fn state_revision_or_initial(snapshot: Option<&StoredSessionStateSnapshot>) -> u64 { let Some(snapshot) = snapshot else { return INITIAL_STATE_REVISION; @@ -44,6 +55,7 @@ pub(crate) fn session_has_content(session: &AppSession) -> bool { || session.meta.input_draft.is_some() || !session.meta.queued_messages.is_empty() || !session.meta.queued_submissions.is_empty() + || !session.meta.queued_direct_tools.is_empty() || session.meta.mode != Some(n00n_storage::sessions::StoredMode::Build) || session.meta.plan_path.is_some() || session.meta.plan_written @@ -216,11 +228,7 @@ impl App { return; }; let session_id = self.state.session.id; - let identity = SessionIdentity::root(SessionRef::from_id(session_id)); - if let Err(error) = handle.drop_state_owner(session_id) { - tracing::warn!(%session_id, %error, "failed to clear stale plugin session state"); - return; - } + let identity = plugin_state_identity(&self.state.session); if let Err(error) = handle.hydrate_state(&identity, self.state.session.meta.state_snapshot.clone()) { @@ -233,7 +241,7 @@ impl App { return; }; let session_id = self.state.session.id; - let identity = SessionIdentity::root(SessionRef::from_id(session_id)); + let identity = plugin_state_identity(&self.state.session); let persisted_revision = state_revision_or_initial(self.state.session.meta.state_snapshot.as_ref()); let revision = self @@ -269,6 +277,17 @@ impl App { .into_iter() .map(|(input, delivery)| stored_message(input, delivery)) .collect(); + let queued_direct_tools: Vec<_> = self + .queue + .direct_tools() + .into_iter() + .map(|(tool, input)| StoredDirectTool { tool, input }) + .collect(); + if !queued_direct_tools.is_empty() { + self.state.session.meta.queued_direct_tools = queued_direct_tools; + } else if !self.state.session.meta.lifecycle.is_active() { + self.state.session.meta.queued_direct_tools.clear(); + } self.state.session.meta.subagents = self .chats @@ -370,6 +389,15 @@ impl App { for (msg, input, delivery) in queued { self.queue_restored_submission(msg, input, delivery); } + for bootstrap in self.state.session.meta.queued_direct_tools.clone() { + self.run_id += 1; + self.status = super::Status::Streaming; + self.queue.push_direct_tool(QueueItem::DirectTool { + run_id: self.run_id, + tool: bootstrap.tool, + input: bootstrap.input, + }); + } self.fire_restore_items(restore_items); diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index d2ec257be..e5406125b 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -7,7 +7,7 @@ //! waits on every event source at once and wakes the moment a plugin action, //! agent event, or keypress arrives instead of sleeping in `event::poll`. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -38,7 +38,7 @@ use n00n_storage::StateDir; use n00n_storage::StorageError; use n00n_storage::id::{SessionRef, n00nId, n00nIdParseError}; use n00n_storage::sessions::{ - SessionError, StoredSessionLifecycle, TranscriptEntry, normalize_title, + SessionError, StoredDirectTool, StoredSessionLifecycle, TranscriptEntry, normalize_title, }; use serde_json::{Value, json}; use tracing::warn; @@ -154,6 +154,32 @@ fn live_session(session: &AppSession) -> std::result::Result>, +) -> std::result::Result { + let mut current = start; + let mut seen = HashSet::new(); + loop { + if !seen.insert(current) { + return Err(LineageError::Cycle(current)); + } + let parent = parents + .get(¤t) + .ok_or(LineageError::UnknownSession(current))?; + let Some(parent) = parent else { + return Ok(current); + }; + if !parents.contains_key(parent) { + return Err(LineageError::MissingParent { + id: current, + parent: *parent, + }); + } + current = *parent; + } +} + fn session_identity(session: &AppSession) -> std::result::Result { let live = live_session(session)?; let session_id = SessionRef::from(live.id); @@ -167,6 +193,16 @@ fn session_identity(session: &AppSession) -> std::result::Result( + result: std::result::Result, + session_id: n00nId, + action: &'static str, +) { + if let Err(error) = result { + warn!(%session_id, %error, action, "session lineage cleanup failed"); + } +} + fn paused_team_run(history: &[Message]) -> Option { let (user_index, last_user) = history .iter() @@ -488,9 +524,9 @@ impl<'t> EventLoop<'t> { mut model, needs_login, commands, - sessions, + mut sessions, focused, - startup_warnings, + mut startup_warnings, storage, config, ui_config, @@ -537,22 +573,66 @@ impl<'t> EventLoop<'t> { let picker = Arc::new(terminal_image::picker()); let runtime_ids: HashSet<_> = sessions.iter().map(|session| session.id).collect(); + let stored = AppSession::list(&cwd.to_string_lossy(), &storage) + .map_err(|error| eyre!("failed to list stored session lineage: {error}"))?; + let mut stored_sessions = Vec::new(); + for summary in stored { + if runtime_ids.contains(&summary.id) { + continue; + } + match AppSession::load(summary.id, &storage) { + Ok(session) => stored_sessions.push(session), + Err(error) => startup_warnings.push(format!( + "Skipped unreadable stored session {}: {error}", + summary.id + )), + } + } + let parents: HashMap<_, _> = sessions + .iter() + .chain(&stored_sessions) + .map(|session| (session.id, session.meta.parent_id)) + .collect(); + for session in &mut sessions { + let root = resolved_root(session.id, &parents) + .map_err(|error| eyre!("invalid live session lineage: {error}"))?; + if session + .meta + .root_session_id + .is_some_and(|stored| stored != root) + { + return Err(eyre!( + "invalid live session lineage root for {}", + session.id + )); + } + session.meta.root_session_id = (session.meta.parent_id.is_some()).then_some(root); + } let mut live_sessions = sessions .iter() .map(live_session) .collect::, _>>() .map_err(|error| eyre!("invalid live session lineage: {error}"))?; - let stored = AppSession::list(&cwd.to_string_lossy(), &storage) - .map_err(|error| eyre!("failed to reconstruct stored session lineage: {error}"))?; - for summary in stored { - if runtime_ids.contains(&summary.id) { - continue; + for mut session in stored_sessions { + let root = match resolved_root(session.id, &parents) { + Ok(root) => root, + Err(error) => { + startup_warnings.push(format!( + "Skipped stored session {} with invalid lineage: {error}", + session.id + )); + continue; + } + }; + let migrated_root = session.meta.parent_id.map(|_| root); + if session.meta.root_session_id != migrated_root { + session.meta.root_session_id = migrated_root; + storage_writer.send(Box::new(session.clone())); } - let session = AppSession::load(summary.id, &storage) - .map_err(|error| eyre!("failed to load stored session lineage node: {error}"))?; let mut node = live_session(&session) .map_err(|error| eyre!("invalid stored session lineage: {error}"))?; node.runtime_present = false; + node.execution_active = false; live_sessions.push(node); } let lineage = SessionLineageGuard::from_live( @@ -800,6 +880,11 @@ impl<'t> EventLoop<'t> { } fn handle_agent(&mut self, idx: usize, envelope: Box) { + if envelope.run_id != self.sessions[idx].app.run_id { + let actions = self.sessions[idx].app.update(Msg::Agent(envelope)); + self.dispatch(idx, actions); + return; + } let lifecycle = match &envelope.event { n00n_agent::AgentEvent::Done { .. } => Some(StoredSessionLifecycle::Succeeded), n00n_agent::AgentEvent::Error { .. } => Some(StoredSessionLifecycle::Failed), @@ -830,6 +915,13 @@ impl<'t> EventLoop<'t> { if let Some(lifecycle) = lifecycle { self.sessions[idx].app.state.session.meta.lifecycle = lifecycle; if terminal { + self.sessions[idx] + .app + .state + .session + .meta + .queued_direct_tools + .clear(); let id = self.sessions[idx].id(); if let Err(error) = self.lineage.set_execution_active(id, false) { warn!(session_id = %id, error = %error, "failed to release session lineage activity"); @@ -973,6 +1065,7 @@ impl<'t> EventLoop<'t> { /// `List` replies from a background task (the scan can be slow); every /// other request is answered synchronously by the event loop, which owns /// the live runtimes. + #[allow(clippy::too_many_lines)] fn handle_session_request( &mut self, req: SessionRequest, @@ -1002,24 +1095,53 @@ impl<'t> EventLoop<'t> { return; } }; - if let Some(i) = self.position(id) { - if i == self.focused { - let _ = reply_tx.send(Err(DELETE_FOCUSED_ERR.into())); + let mut targets = match self.lineage.descendants_of(id) { + Ok(targets) => targets, + Err(LineageError::UnknownSession(_)) => Vec::new(), + Err(error) => { + let _ = reply_tx.send(Err(error.to_string())); return; } - let rt = self.remove_runtime(i); + }; + targets.push(id); + let focused_id = self.sessions[self.focused].id(); + if targets.contains(&focused_id) { + let _ = reply_tx.send(Err(DELETE_FOCUSED_ERR.into())); + return; + } + let mut runtime_indices: Vec<_> = targets + .iter() + .filter_map(|target| self.position(*target)) + .collect(); + runtime_indices.sort_unstable_by(|left, right| right.cmp(left)); + for index in runtime_indices { + let rt = self.remove_runtime(index); + let runtime_id = rt.id(); + rt.app.drop_plugin_state(runtime_id); rt.handles.cancel(); - rt.app.drop_plugin_state(id); } - self.ctx.storage_writer.delete(id, move |res| { - let reply = match res { - Ok(()) | Err(SessionError::Storage(StorageError::NotFound(_))) => { - Ok(json!(true)) + self.lineage.remove_sessions(&targets); + let (done_tx, done_rx) = flume::unbounded(); + let count = targets.len(); + for target in targets { + let done_tx = done_tx.clone(); + self.ctx.storage_writer.delete(target, move |result| { + let _ = done_tx.send(result); + }); + } + smol::spawn(async move { + let mut failure = None; + for _ in 0..count { + match done_rx.recv_async().await { + Ok(Ok(()) | Err(SessionError::Storage(StorageError::NotFound(_)))) => {} + Ok(Err(error)) => failure = Some(error.to_string()), + Err(error) => failure = Some(error.to_string()), } - Err(e) => Err(e.to_string()), - }; + } + let reply = failure.map_or_else(|| Ok(json!(true)), Err); let _ = reply_tx.send(reply); - }); + }) + .detach(); } SessionRequest::Live => { let list: Vec<_> = self @@ -1085,7 +1207,11 @@ impl<'t> EventLoop<'t> { let caller_lineage = match self.lineage.lineage(caller) { Ok(lineage) => lineage, Err(error) => { - let _ = self.lineage.release(reservation); + warn_lineage_cleanup( + self.lineage.release(reservation), + caller, + "release reservation", + ); return Err(error.to_string()); } }; @@ -1105,7 +1231,11 @@ impl<'t> EventLoop<'t> { let runtime = match self.ctx.spawn_runtime(session) { Ok(runtime) => runtime, Err(error) => { - let _ = self.lineage.release(reservation); + warn_lineage_cleanup( + self.lineage.release(reservation), + caller, + "release reservation", + ); return Err(error.to_string()); } }; @@ -1122,6 +1252,10 @@ impl<'t> EventLoop<'t> { app.status = Status::Streaming; app.state.session.meta.lifecycle = StoredSessionLifecycle::Bootstrapping; + app.state.session.meta.queued_direct_tools = vec![StoredDirectTool { + tool: bootstrap.tool.clone(), + input: bootstrap.input.clone(), + }]; app.run_id }; self.sessions[idx] @@ -1138,13 +1272,21 @@ impl<'t> EventLoop<'t> { } else { self.sessions[idx].app.state.session.meta.lifecycle = StoredSessionLifecycle::Idle; - let _ = self.lineage.set_execution_active(id, false); + warn_lineage_cleanup( + self.lineage.set_execution_active(id, false), + id, + "clear idle activity", + ); Ok(json!("idle")) }; if let Err(error) = start_result { let runtime = self.remove_runtime(idx); runtime.handles.cancel(); - let _ = self.lineage.rollback_new(id); + warn_lineage_cleanup( + self.lineage.rollback_new(id), + id, + "roll back new session", + ); return Err(error); } self.sessions[idx].app.save_session(); @@ -1161,14 +1303,19 @@ impl<'t> EventLoop<'t> { steer, control, caller_id, + host_control, } => { let reply = (|| { - let caller = caller_session_id(caller_id)?; let explicit_target = id.as_deref().map(parse_session_id).transpose()?; - let target = self - .lineage - .authorize_prompt(caller, explicit_target) - .map_err(|error| error.to_string())?; + let target = if host_control { + explicit_target + .ok_or_else(|| "host control requires a target session".to_owned())? + } else { + let caller = caller_session_id(caller_id)?; + self.lineage + .authorize_prompt(caller, explicit_target) + .map_err(|error| error.to_string())? + }; let idx = self .position(target) .ok_or_else(|| format!("{NOT_LIVE_ERR}: {target}"))?; @@ -1184,7 +1331,11 @@ impl<'t> EventLoop<'t> { } Err(error) => { if activated { - let _ = self.lineage.set_execution_active(target, false); + warn_lineage_cleanup( + self.lineage.set_execution_active(target, false), + target, + "roll back prompt activity", + ); } Err(error) } @@ -1192,14 +1343,21 @@ impl<'t> EventLoop<'t> { })(); let _ = reply_tx.send(reply); } - SessionRequest::Cancel { id, caller_id } => { + SessionRequest::Cancel { + id, + caller_id, + host_control, + } => { let reply = (|| { - let caller = caller_session_id(caller_id)?; let requested = parse_session_id(&id)?; - let target = self - .lineage - .authorize_prompt(caller, Some(requested)) - .map_err(|error| error.to_string())?; + let target = if host_control { + requested + } else { + let caller = caller_session_id(caller_id)?; + self.lineage + .authorize_prompt(caller, Some(requested)) + .map_err(|error| error.to_string())? + }; let mut targets = self .lineage .descendants_of(target) @@ -1212,9 +1370,14 @@ impl<'t> EventLoop<'t> { .map_err(|error| error.to_string())?; cancelled |= session.meta.lifecycle.is_active(); session.meta.lifecycle = StoredSessionLifecycle::Cancelled; + session.meta.queued_direct_tools.clear(); session.updated_at = n00n_storage::now_epoch(); self.ctx.storage_writer.send(Box::new(session)); - let _ = self.lineage.set_execution_active(session_id, false); + warn_lineage_cleanup( + self.lineage.set_execution_active(session_id, false), + session_id, + "clear cancelled activity", + ); continue; }; if SessionStatus::of(&self.sessions[idx].app) != SessionStatus::Idle @@ -1233,7 +1396,11 @@ impl<'t> EventLoop<'t> { self.sessions[idx].app.state.session.meta.lifecycle = StoredSessionLifecycle::Cancelled; self.sessions[idx].app.save_session(); - let _ = self.lineage.set_execution_active(session_id, false); + warn_lineage_cleanup( + self.lineage.set_execution_active(session_id, false), + session_id, + "clear cancelled activity", + ); } if !cancelled { return Err(format!("session is idle: {target}")); @@ -1347,7 +1514,11 @@ impl<'t> EventLoop<'t> { let runtime = match self.ctx.spawn_runtime(session) { Ok(runtime) => runtime, Err(error) => { - let _ = self.lineage.remove_runtime(id); + warn_lineage_cleanup( + self.lineage.remove_runtime(id), + id, + "roll back runtime activation", + ); return Err(error.to_string()); } }; @@ -1542,10 +1713,33 @@ impl<'t> EventLoop<'t> { }); } Action::CancelAgent { run_id } => { - let _ = self.sessions[idx] + let id = self.sessions[idx].id(); + match self.sessions[idx] .handles .cmd_tx - .try_send(AgentCommand::Cancel { run_id }); + .try_send(AgentCommand::Cancel { run_id }) + { + Ok(()) => { + self.sessions[idx].app.state.session.meta.lifecycle = + StoredSessionLifecycle::Cancelled; + self.sessions[idx] + .app + .state + .session + .meta + .queued_direct_tools + .clear(); + warn_lineage_cleanup( + self.lineage.set_execution_active(id, false), + id, + "clear keyboard-cancelled activity", + ); + self.sessions[idx].app.save_session(); + } + Err(error) => { + warn!(session_id = %id, %error, "failed to send agent cancellation"); + } + } } Action::CancelSubagent { tool_use_id } => { let _ = self.sessions[idx] diff --git a/n00n-ui/src/session_lineage.rs b/n00n-ui/src/session_lineage.rs index cbc52a0be..8e40827d3 100644 --- a/n00n-ui/src/session_lineage.rs +++ b/n00n-ui/src/session_lineage.rs @@ -86,6 +86,7 @@ struct SessionNode { parent_id: Option, runtime_present: bool, execution_active: bool, + deleted: bool, } #[derive(Debug, Clone, Copy)] @@ -124,6 +125,7 @@ impl SessionLineageGuard { parent_id: session.parent_id, runtime_present: session.runtime_present, execution_active: session.execution_active, + deleted: false, }, ) .is_some() @@ -137,6 +139,9 @@ impl SessionLineageGuard { pub(crate) fn activate_runtime(&mut self, session: LiveSession) -> Result<(), LineageError> { if let Some(existing) = self.sessions.get(&session.id) { + if existing.deleted { + return Err(LineageError::UnknownSession(session.id)); + } if existing.runtime_present { return Err(LineageError::DuplicateSession(session.id)); } @@ -170,6 +175,7 @@ impl SessionLineageGuard { parent_id: session.parent_id, runtime_present: true, execution_active: session.execution_active, + deleted: false, }, ); if let Err(error) = self.validate_graph() { @@ -198,7 +204,7 @@ impl SessionLineageGuard { .sessions .get_mut(&id) .ok_or(LineageError::UnknownSession(id))?; - if active && !node.runtime_present { + if active && (!node.runtime_present || node.deleted) { return Err(LineageError::TargetNotLive(id)); } let changed = node.execution_active != active; @@ -238,7 +244,7 @@ impl SessionLineageGuard { .sessions .get(&caller) .ok_or(LineageError::CallerNotLive(caller))?; - if !node.runtime_present { + if !node.runtime_present || node.deleted { return Err(LineageError::CallerNotLive(caller)); } self.lineage_for(caller).map_err(|error| match error { @@ -333,6 +339,7 @@ impl SessionLineageGuard { parent_id: Some(pending.parent), runtime_present: true, execution_active: true, + deleted: false, }, ); if let Err(error) = self.validate_graph() { @@ -377,6 +384,20 @@ impl SessionLineageGuard { Ok(descendants) } + pub(crate) fn remove_sessions(&mut self, ids: &[n00nId]) { + let removed: HashSet<_> = ids.iter().copied().collect(); + for id in &removed { + if let Some(node) = self.sessions.get_mut(id) { + node.runtime_present = false; + node.execution_active = false; + node.deleted = true; + } + } + self.reservations.retain(|_, reservation| { + !removed.contains(&reservation.caller) && !removed.contains(&reservation.parent) + }); + } + pub(crate) fn authorize_prompt( &self, caller: n00nId, @@ -392,15 +413,11 @@ impl SessionLineageGuard { .sessions .get(&target) .ok_or(LineageError::UnknownSession(target))?; - if !target_node.runtime_present { + if !target_node.runtime_present || target_node.deleted { return Err(LineageError::TargetNotLive(target)); } - let caller_path = self.path_from(caller)?; let target_path = self.path_from(target)?; - if caller_lineage.caller == target - || caller_path.contains(&target) - || target_path.contains(&caller) - { + if caller_lineage.caller == target || target_path.contains(&caller) { return Ok(target); } Err(LineageError::UnauthorizedTarget) @@ -410,6 +427,9 @@ impl SessionLineageGuard { let mut total = 0; let mut active = 0; for (&id, node) in &self.sessions { + if node.deleted { + continue; + } let lineage = self.lineage_for(id)?; if lineage.root != root || id == root { continue; @@ -495,7 +515,7 @@ impl SessionLineageGuard { } fn limit_reached(committed: usize, reserved: usize, limit: usize) -> bool { - committed >= limit || reserved >= limit.saturating_sub(committed) + committed.saturating_add(reserved) >= limit } #[cfg(test)] @@ -591,6 +611,10 @@ mod tests { guard.authorize_prompt(sibling, Some(foreign)), Err(LineageError::UnauthorizedTarget) )); + assert!(matches!( + guard.authorize_prompt(sibling, Some(root)), + Err(LineageError::UnauthorizedTarget) + )); } #[test] diff --git a/plugins/lib/n00n/subagent.lua b/plugins/lib/n00n/subagent.lua index 17abd7fcb..4bf86e102 100644 --- a/plugins/lib/n00n/subagent.lua +++ b/plugins/lib/n00n/subagent.lua @@ -10,6 +10,14 @@ local structured_output = require("n00n.structured_output") local ORCHESTRATION_TOOLS = { "task", "team", "workflow", "agent_control", "batch" } +function M.orchestration_tools() + local copy = {} + for index, name in ipairs(ORCHESTRATION_TOOLS) do + copy[index] = name + end + return copy +end + local function excluded_tools(opts) local excluded = {} local seen = {} diff --git a/plugins/task/init.lua b/plugins/task/init.lua index 744c18013..126c4c9f0 100644 --- a/plugins/task/init.lua +++ b/plugins/task/init.lua @@ -15,7 +15,7 @@ local subagent = require("n00n.subagent") local DONE_NAME = "done" local DONE_DESCRIPTION = "Call when the task is complete with your final answer." local DONE_PROMPT_SUFFIX = "\n\nWhen finished, call the done tool with your final answer." -local ORCHESTRATION_TOOLS = { "task", "team", "workflow", "agent_control", "batch" } +local ORCHESTRATION_TOOLS = subagent.orchestration_tools() local BODY_INDENT_COLS = 4 local MIN_MD_WIDTH = 20 local DEFAULT_OUTPUT_LINES = 5 @@ -82,7 +82,7 @@ local function handler(input, ctx) forwarded[key] = value end forwarded.background = false - local title = "task: " .. (input.description or input.prompt or "background task"):sub(1, 60) + local title = "task: " .. n00n.ui.truncate_text(input.description or input.prompt or "background task", 60).head local id, err = n00n.session.new({ tool = "task", input = forwarded, diff --git a/plugins/team/init.lua b/plugins/team/init.lua index bdee3cc3b..d71b553fe 100644 --- a/plugins/team/init.lua +++ b/plugins/team/init.lua @@ -802,7 +802,7 @@ local function run_team(input, ctx) forwarded[key] = value end forwarded.background = false - local title = "team: " .. (input.goal or ""):sub(1, 60) + local title = "team: " .. n00n.ui.truncate_text(input.goal or "", 60).head local id, err = n00n.session.new({ tool = "team", input = forwarded, @@ -813,7 +813,11 @@ local function run_team(input, ctx) return { llm_output = err, is_error = true } end - return n00n.json.encode({ agent_id = id, status = "started", title = title }) + local output, output_err = n00n.json.encode({ agent_id = id, status = "started", title = title }) + if output_err then + return { llm_output = "failed to encode team status: " .. tostring(output_err), is_error = true } + end + return output end local requested_mode = input.mode diff --git a/site/docs/content/configuration/_index.md b/site/docs/content/configuration/_index.md index b6ec40acf..104f512bc 100644 --- a/site/docs/content/configuration/_index.md +++ b/site/docs/content/configuration/_index.md @@ -113,6 +113,8 @@ How many lines of output to show per tool in the UI. All values are `usize` with | `compaction_buffer` | u32 \| string | `20%` | - | Context reserved for compaction: token count or percent of the context window (e.g. "20%") | | `mcp_tool_desc_max_chars` | usize | `200` | 10 | Max MCP tool description length (characters) | +Keep `max_depth` and `max_active_descendants` at or below `max_total_descendants`. n00n reports a configuration error at startup if either value is higher. + ### `agent.fusion` | Field | Type | Default | Description | diff --git a/site/docs/content/lua-api/_index.md b/site/docs/content/lua-api/_index.md index 00574cead..0b21609c0 100644 --- a/site/docs/content/lua-api/_index.md +++ b/site/docs/content/lua-api/_index.md @@ -6005,6 +6005,7 @@ function M.make_local_tool(schema, on_submit) -- Subagent launch helper module. -- Provides a unified interface for launching subagents with model resolution, -- system prompts, tool setup, and optional structured output validation. +M.ORCHESTRATION_TOOLS = { "task", "team", "workflow", "agent_control", "batch" } -- Launch a subagent with the given options. -- Returns (result | nil, err, cost, usage, model_spec) diff --git a/src/cmd/tui_bridge.rs b/src/cmd/tui_bridge.rs index 66f00a14c..6efea161f 100644 --- a/src/cmd/tui_bridge.rs +++ b/src/cmd/tui_bridge.rs @@ -12,6 +12,7 @@ use n00n_daemon::protocol::{AgentRecord, BackendKind, MessageOpts}; use n00n_daemon::registry::{ControlPlane, TuiCallbackBackend}; use n00n_daemon::server; use n00n_lua::{SessionRequest, UiAction}; +use n00n_storage::id::SessionRef; use serde_json::Value; const SESSION_ROUNDTRIP_TIMEOUT: Duration = Duration::from_secs(5); @@ -123,6 +124,8 @@ fn message_one( text: &str, opts: &MessageOpts, ) -> ControlResult { + id.parse::() + .map_err(|_| ControlError::InvalidId(id.to_owned()))?; session_call( tx, SessionRequest::Prompt { @@ -131,6 +134,7 @@ fn message_one( steer: opts.steer, control: opts.control, caller_id: None, + host_control: true, }, ) .map_err(|e| map_not_found(id, e))?; @@ -138,6 +142,8 @@ fn message_one( } fn resume_one(tx: &flume::Sender, id: &str) -> ControlResult<()> { + id.parse::() + .map_err(|_| ControlError::InvalidId(id.to_owned()))?; let value = session_call(tx, SessionRequest::Status { id: id.to_owned() }) .map_err(|e| map_not_found(id, e))?; let run_info = value.get("paused_team").ok_or_else(|| { @@ -152,6 +158,7 @@ fn resume_one(tx: &flume::Sender, id: &str) -> ControlResult<()> { steer: true, control: true, caller_id: None, + host_control: true, }, ) .map_err(|e| map_not_found(id, e))?; @@ -181,11 +188,14 @@ fn build_team_resume_prompt(run_info: &Value) -> ControlResult { } fn stop_one(tx: &flume::Sender, id: &str) -> ControlResult<()> { + id.parse::() + .map_err(|_| ControlError::InvalidId(id.to_owned()))?; session_call( tx, SessionRequest::Cancel { id: id.to_owned(), caller_id: None, + host_control: true, }, ) .map_err(|e| map_not_found(id, e))?; @@ -341,7 +351,7 @@ mod tests { respond_live( rx, json!([{ - "id": "sess-1", + "id": "00000000-0000-7000-8000-000000000001", "title": "t", "status": "working", "updated_at": 0, @@ -351,7 +361,7 @@ mod tests { let backend = tui_backend(tx); let agents = backend.list().map_err(|e| e.to_string())?; assert_eq!(agents.len(), 1); - assert_eq!(agents[0].id, "sess-1"); + assert_eq!(agents[0].id, "00000000-0000-7000-8000-000000000001"); assert_eq!(agents[0].backend, BackendKind::Tui); assert_eq!(agents[0].status, "working"); Ok(()) @@ -370,12 +380,14 @@ mod tests { steer, control, caller_id, + host_control, } => { - if id.as_deref() != Some("sess-1") + if id.as_deref() != Some("00000000-0000-7000-8000-000000000001") || text != "hi" || !steer || !control || caller_id.is_some() + || !host_control { let _ = reply_tx.send(Err(format!( "unexpected prompt id={id:?} text={text:?} steer={steer} control={control} caller_id={caller_id:?}" @@ -393,7 +405,7 @@ mod tests { let backend = tui_backend(tx); backend .message( - "sess-1", + "00000000-0000-7000-8000-000000000001", "hi", &MessageOpts { steer: true, @@ -413,10 +425,12 @@ mod tests { rx.recv_timeout(Duration::from_secs(2)) { match req { - SessionRequest::Status { id } if id == "sess-1" => { + SessionRequest::Status { id } + if id == "00000000-0000-7000-8000-000000000001" => + { saw_status = true; let _ = reply_tx.send(Ok(json!({ - "id": "sess-1", + "id": "00000000-0000-7000-8000-000000000001", "status": "paused", "paused_team": { "run_id": "run-abc", "mode": "swarm" }, })) as SessionReply); @@ -427,11 +441,13 @@ mod tests { steer, control, caller_id, + host_control, } => { - if id.as_deref() != Some("sess-1") + if id.as_deref() != Some("00000000-0000-7000-8000-000000000001") || !steer || !control || caller_id.is_some() + || !host_control { let _ = reply_tx.send(Err(format!( "unexpected prompt id={id:?} steer={steer} control={control} caller_id={caller_id:?}" @@ -456,7 +472,9 @@ mod tests { assert!(saw_status, "never received status request"); }); let backend = tui_backend(tx); - backend.resume("sess-1").map_err(|e| e.to_string())?; + backend + .resume("00000000-0000-7000-8000-000000000001") + .map_err(|e| e.to_string())?; Ok(()) } From 249a3fc15b73f344ca54e89f3df657e5572c6b80 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Sun, 9 Aug 2026 22:37:38 -0400 Subject: [PATCH 14/24] fix(lua): reconcile async task identity merge --- n00n-lua/src/runtime.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index 16f88ac3c..fb90bdf58 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -3597,6 +3597,7 @@ mod tests { CancelToken::none(), Some(Instant::now() + Duration::from_millis(10)), None, + None, ), ); @@ -3618,6 +3619,7 @@ mod tests { CancelToken::none(), Some(Instant::now() + Duration::from_millis(20)), None, + None, ))); let waiter = smol::spawn({ let handle = Arc::clone(&handle); @@ -3643,6 +3645,7 @@ mod tests { CancelToken::none(), Some(Instant::now() + Duration::from_secs(1)), None, + None, ))); let waiter = smol::spawn({ let handle = Arc::clone(&handle); @@ -3668,7 +3671,7 @@ mod tests { let handle = cancelled_handle(); let scope = TaskScope::new( &lua, - TaskCell::new(lock_cell(&handle).cancel.clone(), None, None), + TaskCell::new(lock_cell(&handle).cancel.clone(), None, None, None), ); assert_eq!(interrupt_reason(&lua), None); let cleanup_deadline = lock_cell(scope.handle()).interrupt_after.get().unwrap(); @@ -3895,6 +3898,7 @@ mod tests { cancel: token, deadline: None, live_ctx: None, + identity: None, owner: None, parent: None, }; @@ -4034,6 +4038,7 @@ mod tests { CancelToken::none(), Some(deadline), None, + None, ))); lock_cell(&handle).deadline_secs.set(Some(1)); @@ -4048,6 +4053,7 @@ mod tests { CancelToken::none(), Some(Instant::now()), None, + None, ))); let reply = timeout_reply(&handle, "test", "tool"); From 03eaea1694ebd6aa9abf04c3ccae3d135a54bd3d Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 10 Aug 2026 01:06:32 -0400 Subject: [PATCH 15/24] fix(orchestration): close background lifecycle gaps --- n00n-lua/src/loader.rs | 19 ++- n00n-lua/src/runtime.rs | 17 ++- n00n-lua/tests/plugin_host.rs | 7 +- n00n-ui/src/agent/agent_loop.rs | 6 +- n00n-ui/src/app/mod.rs | 7 +- n00n-ui/src/app/session.rs | 2 +- n00n-ui/src/event_loop.rs | 109 +++++++++++--- n00n-ui/src/session_lineage.rs | 216 +++++++++++++++++++++------- plugins/workflow/init.lua | 16 ++- site/docs/content/lua-api/_index.md | 2 +- 10 files changed, 322 insertions(+), 79 deletions(-) diff --git a/n00n-lua/src/loader.rs b/n00n-lua/src/loader.rs index 0848f99bd..d6ef4f66c 100644 --- a/n00n-lua/src/loader.rs +++ b/n00n-lua/src/loader.rs @@ -660,11 +660,18 @@ impl EventHandle { } } - pub fn run_command(&self, plugin: Arc, command: Arc, args: String) { + pub fn run_command( + &self, + plugin: Arc, + command: Arc, + args: String, + identity: Option, + ) { let _ = self.prio_tx.try_send(Request::RunCommand { plugin, command, args, + identity, }); } @@ -976,17 +983,25 @@ mod tests { prio_tx, state_leases: Arc::new(StateLeases::default()), }; - handle.run_command(Arc::from("myplugin"), Arc::from("/greet"), "world".into()); + let identity = SessionIdentity::root(n00n_storage::id::SessionRef::generate()); + handle.run_command( + Arc::from("myplugin"), + Arc::from("/greet"), + "world".into(), + Some(identity.clone()), + ); let req = prio_rx.try_recv().unwrap(); match req { Request::RunCommand { plugin, command, args, + identity: request_identity, } => { assert_eq!(plugin.as_ref(), "myplugin"); assert_eq!(command.as_ref(), "/greet"); assert_eq!(args, "world"); + assert_eq!(request_identity, Some(identity)); } _ => panic!("expected RunCommand"), } diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index fb90bdf58..0472ad5d5 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -183,6 +183,7 @@ pub enum Request { plugin: Arc, command: Arc, args: String, + identity: Option, }, CollectPromptSlots { reply: flume::Sender, @@ -688,7 +689,18 @@ impl TaskScope { /// /// [detached]: TaskScope::detached pub(crate) async fn run_detached(lua: &Lua, fut: F) -> F::Output { - let scope = TaskScope::detached(lua); + run_callback(lua, None, fut).await +} + +async fn run_callback( + lua: &Lua, + identity: Option, + fut: F, +) -> F::Output { + let scope = TaskScope::new( + lua, + TaskCell::new(CancelToken::none(), None, None, identity), + ); let handle = Arc::clone(scope.handle()); let pump = async { let mut event_buf = Vec::new(); @@ -3029,6 +3041,7 @@ pub fn spawn( plugin, command, args, + identity, } => { let handler_fn = rt.lua.app_data_ref::().and_then(|m| { @@ -3042,7 +3055,7 @@ pub fn spawn( let thread = lua.create_thread(func)?; thread.into_async::<()>(args)?.await }; - if let Err(e) = run_detached(&lua, run).await { + if let Err(e) = run_callback(&lua, identity, run).await { tracing::warn!(plugin = %plugin, command = %command, error = %e, "command handler failed"); } }) diff --git a/n00n-lua/tests/plugin_host.rs b/n00n-lua/tests/plugin_host.rs index 336e7d5b0..61a6ae31e 100644 --- a/n00n-lua/tests/plugin_host.rs +++ b/n00n-lua/tests/plugin_host.rs @@ -6103,6 +6103,7 @@ fn team_launcher_uses_native_model_picker_and_amp_labels() { Arc::from("team"), Arc::from("/team"), "fix the parser".into(), + None, ); let action = rx @@ -6187,7 +6188,7 @@ fn team_launcher_collects_goal_and_submits_configured_prompt() { let (_reg, host) = builtins_host(); let rx = host.ui_action_rx().unwrap(); let handle = host.event_handle().unwrap(); - handle.run_command(Arc::from("team"), Arc::from("/team"), String::new()); + handle.run_command(Arc::from("team"), Arc::from("/team"), String::new(), None); let action = rx .recv_timeout(Duration::from_secs(5)) @@ -6316,7 +6317,7 @@ fn async_run_from_parked_command_handler_runs_promptly() { .unwrap(); let rx = host.ui_action_rx().unwrap(); let handle = host.event_handle().unwrap(); - handle.run_command(Arc::from("p"), Arc::from("/park"), String::new()); + handle.run_command(Arc::from("p"), Arc::from("/park"), String::new(), None); let action = rx .recv_timeout(Duration::from_secs(5)) @@ -6347,7 +6348,7 @@ fn job_callbacks_fire_while_command_handler_parked() { .unwrap(); let rx = host.ui_action_rx().unwrap(); let handle = host.event_handle().unwrap(); - handle.run_command(Arc::from("p"), Arc::from("/stream"), String::new()); + handle.run_command(Arc::from("p"), Arc::from("/stream"), String::new(), None); let action = rx .recv_timeout(Duration::from_secs(5)) diff --git a/n00n-ui/src/agent/agent_loop.rs b/n00n-ui/src/agent/agent_loop.rs index 059426e8a..ab53148d2 100644 --- a/n00n-ui/src/agent/agent_loop.rs +++ b/n00n-ui/src/agent/agent_loop.rs @@ -238,6 +238,8 @@ impl AgentLoop { self.rebuild_tools(&slot.model, false); let (trigger, cancel) = CancelToken::new(); self.set_cancel_trigger(run_id, trigger); + while self.answer_rx.lock().await.try_recv().is_ok() {} + let agent = Agent::new( AgentParams { provider: Arc::clone(&slot.provider), @@ -263,6 +265,7 @@ impl AgentLoop { }, ) .with_cancel(cancel) + .with_user_response_rx(Arc::clone(&self.answer_rx)) .with_mcp(self.mcp.clone()); let result = agent .run_tool(format!("bootstrap-{run_id}"), tool, input) @@ -597,9 +600,8 @@ fn spawn_oauth_for_needs_auth(handle: &n00n_agent::mcp::McpHandle) { #[cfg(test)] mod tests { - use std::path::Path; - use n00n_agent::AgentMode; + use std::path::Path; use super::build_plan_path; diff --git a/n00n-ui/src/app/mod.rs b/n00n-ui/src/app/mod.rs index 933df0728..52a476d58 100644 --- a/n00n-ui/src/app/mod.rs +++ b/n00n-ui/src/app/mod.rs @@ -1961,7 +1961,12 @@ impl App { let Some(handle) = &self.lua_event_handle else { return; }; - handle.run_command(Arc::clone(&lua_cmd.plugin), Arc::clone(&lua_cmd.name), args); + handle.run_command( + Arc::clone(&lua_cmd.plugin), + Arc::clone(&lua_cmd.name), + args, + Some(session::plugin_state_identity(&self.state.session)), + ); } fn execute_mcp_prompt(&mut self, name: &str, args: &str) -> Vec { diff --git a/n00n-ui/src/app/session.rs b/n00n-ui/src/app/session.rs index ef7f1cf32..fdb73bf69 100644 --- a/n00n-ui/src/app/session.rs +++ b/n00n-ui/src/app/session.rs @@ -24,7 +24,7 @@ use crate::agent::{Delivery, QueuedMessage}; const INITIAL_STATE_REVISION: u64 = 0; -fn plugin_state_identity(session: &AppSession) -> SessionIdentity { +pub(super) fn plugin_state_identity(session: &AppSession) -> SessionIdentity { let root_id = session.meta.root_session_id.map_or(session.id, |root| root); let session_id = SessionRef::from_id(session.id); if root_id == session.id { diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index 38fddd946..bc2c4decf 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -7,6 +7,7 @@ //! waits on every event source at once and wakes the moment a plugin action, //! agent event, or keypress arrives instead of sleeping in `event::poll`. +use std::cell::RefCell; use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use std::sync::Arc; @@ -156,6 +157,12 @@ fn live_session(session: &AppSession) -> std::result::Result bool { + !session.meta.queued_messages.is_empty() + || !session.meta.queued_submissions.is_empty() + || !session.meta.queued_direct_tools.is_empty() +} + fn resolved_root( start: n00nId, parents: &HashMap>, @@ -259,6 +266,9 @@ struct SessionRuntime { shell_tx: flume::Sender, shell_rx: flume::Receiver, last_status: SessionStatus, + direct_bootstrap_active: bool, + direct_output: Option, + direct_output_is_error: bool, } impl SessionRuntime { @@ -289,17 +299,39 @@ struct SpawnCtx { available_models: Arc>>, storage_writer: Arc, picker: Arc, + hydrated_roots: RefCell>, } impl SpawnCtx { fn spawn_runtime(&self, session: AppSession) -> Result { let resumed = crate::app::session_has_content(&session); + let direct_bootstrap_active = !session.meta.queued_direct_tools.is_empty(); let identity = session_identity(&session) .map_err(|error| eyre!("invalid session identity: {error}"))?; if let Some(handle) = &self.lua_event_handle { - handle - .hydrate_state(&identity, session.meta.state_snapshot.clone()) - .map_err(|error| eyre!("failed to hydrate plugin session state: {error}"))?; + let root_id = session.meta.root_session_id.map_or(session.id, |root| root); + if !self.hydrated_roots.borrow().contains(&root_id) { + let root_snapshot = if root_id == session.id { + session.meta.state_snapshot.clone() + } else { + AppSession::load(root_id, &self.storage) + .map_err(|error| eyre!("failed to load root session state: {error}"))? + .meta + .state_snapshot + }; + handle + .hydrate_state( + &SessionIdentity::root(SessionRef::from_id(root_id)), + root_snapshot, + ) + .map_err(|error| eyre!("failed to hydrate root plugin state: {error}"))?; + self.hydrated_roots.borrow_mut().insert(root_id); + } + if root_id != session.id { + handle + .hydrate_state(&identity, session.meta.state_snapshot.clone()) + .map_err(|error| eyre!("failed to hydrate plugin session state: {error}"))?; + } } let permissions = Arc::new(self.permissions.fork()); let initial_plan_path = session.meta.plan_path.as_ref().map(PathBuf::from); @@ -341,13 +373,17 @@ impl SpawnCtx { if resumed { restore_session(&mut app, &handles); } + let last_status = SessionStatus::of(&app); let (shell_tx, shell_rx) = flume::unbounded::(); Ok(SessionRuntime { app, handles, shell_tx, shell_rx, - last_status: SessionStatus::Idle, + last_status, + direct_bootstrap_active, + direct_output: None, + direct_output_is_error: false, }) } } @@ -612,8 +648,12 @@ impl<'t> EventLoop<'t> { } let mut live_sessions = sessions .iter() - .map(live_session) - .collect::, _>>() + .map(|session| { + let mut live = live_session(session)?; + live.execution_active = has_restorable_work(session); + Ok(live) + }) + .collect::, LineageError>>() .map_err(|error| eyre!("invalid live session lineage: {error}"))?; for mut session in stored_sessions { let root = match resolved_root(session.id, &parents) { @@ -666,6 +706,7 @@ impl<'t> EventLoop<'t> { available_models: bg.available, storage_writer, picker, + hydrated_roots: RefCell::new(HashSet::new()), }; let mut runtimes: Vec = sessions @@ -895,6 +936,21 @@ impl<'t> EventLoop<'t> { self.dispatch(idx, actions); return; } + if self.sessions[idx].direct_bootstrap_active { + match &envelope.event { + n00n_agent::AgentEvent::ToolDone(done) => { + self.sessions[idx].direct_output = Some(done.output.as_text()); + self.sessions[idx].direct_output_is_error = done.is_error; + } + n00n_agent::AgentEvent::Error { message } + if self.sessions[idx].direct_output.is_none() => + { + self.sessions[idx].direct_output = Some(message.clone()); + self.sessions[idx].direct_output_is_error = true; + } + _ => {} + } + } let lifecycle = match &envelope.event { n00n_agent::AgentEvent::Done { .. } => Some(StoredSessionLifecycle::Succeeded), n00n_agent::AgentEvent::Error { .. } => Some(StoredSessionLifecycle::Failed), @@ -925,6 +981,7 @@ impl<'t> EventLoop<'t> { if let Some(lifecycle) = lifecycle { self.sessions[idx].app.state.session.meta.lifecycle = lifecycle; if terminal { + self.sessions[idx].direct_bootstrap_active = false; self.sessions[idx] .app .state @@ -1178,11 +1235,16 @@ impl<'t> EventLoop<'t> { .ok_or_else(|| format!("{NOT_LIVE_ERR}: {id}"))?; let rt = &self.sessions[idx]; let history = rt.handles.history.load(); - let output = history.iter().rev().find_map(|message| { + let assistant_output = history.iter().rev().find_map(|message| { matches!(message.role, n00n_providers::Role::Assistant) .then(|| message.first_text_content()) .flatten() }); + let output = assistant_output.or(rt.direct_output.as_deref()); + let direct_error = assistant_output + .is_none() + .then_some(rt.direct_output_is_error) + .filter(|_| rt.direct_output.is_some()); let paused_team = paused_team_run(&history); Ok(json!({ "id": rt.id(), @@ -1191,6 +1253,7 @@ impl<'t> EventLoop<'t> { "updated_at": rt.app.state.session.updated_at, "focused": idx == self.focused, "output": output, + "is_error": direct_error, "paused_team": paused_team, "cwd": rt.app.state.session.cwd, })) @@ -1257,16 +1320,18 @@ impl<'t> EventLoop<'t> { let idx = self.push_runtime(runtime); let start_result = if let Some(bootstrap) = bootstrap { let run_id = { - let app = &mut self.sessions[idx].app; - app.run_id += 1; - app.status = Status::Streaming; - app.state.session.meta.lifecycle = + let runtime = &mut self.sessions[idx]; + runtime.direct_bootstrap_active = true; + runtime.app.run_id += 1; + runtime.app.status = Status::Streaming; + runtime.app.state.session.meta.lifecycle = StoredSessionLifecycle::Bootstrapping; - app.state.session.meta.queued_direct_tools = vec![StoredDirectTool { - tool: bootstrap.tool.clone(), - input: bootstrap.input.clone(), - }]; - app.run_id + runtime.app.state.session.meta.queued_direct_tools = + vec![StoredDirectTool { + tool: bootstrap.tool.clone(), + input: bootstrap.input.clone(), + }]; + runtime.app.run_id }; self.sessions[idx] .handles @@ -1517,10 +1582,20 @@ impl<'t> EventLoop<'t> { } let session = AppSession::load(id, &self.ctx.storage) .map_err(|e| format!("Failed to load session: {e}"))?; - let live = live_session(&session).map_err(|error| error.to_string())?; + let restore_execution = has_restorable_work(&session); + let mut live = live_session(&session).map_err(|error| error.to_string())?; + live.execution_active = false; self.lineage .activate_runtime(live) .map_err(|error| error.to_string())?; + if restore_execution && let Err(error) = self.lineage.begin_execution(id) { + warn_lineage_cleanup( + self.lineage.remove_runtime(id), + id, + "roll back restored execution activation", + ); + return Err(error.to_string()); + } let runtime = match self.ctx.spawn_runtime(session) { Ok(runtime) => runtime, Err(error) => { diff --git a/n00n-ui/src/session_lineage.rs b/n00n-ui/src/session_lineage.rs index 8e40827d3..66440dbba 100644 --- a/n00n-ui/src/session_lineage.rs +++ b/n00n-ui/src/session_lineage.rs @@ -97,9 +97,17 @@ struct PendingReservation { depth: usize, } +#[derive(Debug, Clone, Copy)] +struct CachedLineage { + root: n00nId, + depth: usize, +} + pub(crate) struct SessionLineageGuard { limits: LineageLimits, sessions: HashMap, + children: HashMap>, + lineage_cache: HashMap, reservations: HashMap, next_reservation_id: u64, } @@ -112,6 +120,8 @@ impl SessionLineageGuard { let mut guard = Self { limits, sessions: HashMap::new(), + children: HashMap::new(), + lineage_cache: HashMap::new(), reservations: HashMap::new(), next_reservation_id: 1, }; @@ -133,7 +143,20 @@ impl SessionLineageGuard { return Err(LineageError::DuplicateSession(session.id)); } } - guard.validate_graph()?; + guard.rebuild_topology()?; + let roots = guard + .lineage_cache + .values() + .map(|lineage| lineage.root) + .collect::>(); + for root in roots { + let counts = guard.descendant_counts(root)?; + if counts.active > guard.limits.max_active_descendants { + return Err(LineageError::ActiveDescendantsExceeded { + limit: guard.limits.max_active_descendants, + }); + } + } Ok(guard) } @@ -158,13 +181,6 @@ impl SessionLineageGuard { .get_mut(&session.id) .ok_or(LineageError::UnknownSession(session.id))? .execution_active = session.execution_active; - if let Err(error) = self.validate_graph() { - if let Some(node) = self.sessions.get_mut(&session.id) { - node.runtime_present = false; - node.execution_active = false; - } - return Err(error); - } return Ok(()); } @@ -178,8 +194,9 @@ impl SessionLineageGuard { deleted: false, }, ); - if let Err(error) = self.validate_graph() { + if let Err(error) = self.rebuild_topology() { self.sessions.remove(&session.id); + self.rebuild_topology()?; return Err(error); } Ok(()) @@ -258,7 +275,6 @@ impl SessionLineageGuard { caller: n00nId, explicit_parent: Option, ) -> Result { - self.validate_graph()?; let caller_lineage = self.lineage(caller)?; let parent = match explicit_parent { Some(p) => p, @@ -342,8 +358,9 @@ impl SessionLineageGuard { deleted: false, }, ); - if let Err(error) = self.validate_graph() { + if let Err(error) = self.rebuild_topology() { self.sessions.remove(&child_id); + self.rebuild_topology()?; return Err(error); } Ok(()) @@ -367,18 +384,25 @@ impl SessionLineageGuard { return Err(LineageError::ParentChanged { id }); } self.sessions.remove(&id); + self.rebuild_topology()?; Ok(()) } pub(crate) fn descendants_of(&self, parent: n00nId) -> Result, LineageError> { - self.validate_graph()?; if !self.sessions.contains_key(&parent) { return Err(LineageError::UnknownSession(parent)); } + let mut pending = self + .children + .get(&parent) + .into_iter() + .flat_map(|children| children.iter().copied()) + .collect::>(); let mut descendants = Vec::new(); - for &id in self.sessions.keys() { - if id != parent && self.path_from(id)?.contains(&parent) { - descendants.push(id); + while let Some(id) = pending.pop() { + descendants.push(id); + if let Some(children) = self.children.get(&id) { + pending.extend(children.iter().copied()); } } Ok(descendants) @@ -403,7 +427,6 @@ impl SessionLineageGuard { caller: n00nId, explicit_target: Option, ) -> Result { - self.validate_graph()?; let caller_lineage = self.lineage(caller)?; let target = match explicit_target { Some(t) => t, @@ -416,8 +439,11 @@ impl SessionLineageGuard { if !target_node.runtime_present || target_node.deleted { return Err(LineageError::TargetNotLive(target)); } - let target_path = self.path_from(target)?; - if caller_lineage.caller == target || target_path.contains(&caller) { + let target_lineage = self.lineage_for(target)?; + if caller_lineage.caller == target + || (caller_lineage.root == target_lineage.root + && self.path_from(target)?.contains(&caller)) + { return Ok(target); } Err(LineageError::UnauthorizedTarget) @@ -426,14 +452,14 @@ impl SessionLineageGuard { pub(crate) fn descendant_counts(&self, root: n00nId) -> Result { let mut total = 0; let mut active = 0; - for (&id, node) in &self.sessions { + for id in self.descendants_of(root)? { + let node = self + .sessions + .get(&id) + .ok_or(LineageError::UnknownSession(id))?; if node.deleted { continue; } - let lineage = self.lineage_for(id)?; - if lineage.root != root || id == root { - continue; - } total += 1; if node.execution_active { active += 1; @@ -451,35 +477,52 @@ impl SessionLineageGuard { }) } - fn validate_graph(&self) -> Result<(), LineageError> { + fn rebuild_topology(&mut self) -> Result<(), LineageError> { + let mut children = HashMap::>::new(); + for (&id, node) in &self.sessions { + if let Some(parent) = node.parent_id { + if !self.sessions.contains_key(&parent) { + return Err(LineageError::MissingParent { id, parent }); + } + children.entry(parent).or_default().insert(id); + } + } + + let mut lineage_cache = HashMap::new(); for &id in self.sessions.keys() { - self.path_from(id)?; + resolve_cached_lineage(&self.sessions, &mut lineage_cache, id)?; } + for (&id, node) in &self.sessions { + let lineage = lineage_cache + .get(&id) + .ok_or(LineageError::UnknownSession(id))?; + if node.root_session_id != lineage.root { + return Err(LineageError::RootMismatch { + id, + expected: lineage.root, + found: node.root_session_id, + }); + } + } + self.children = children; + self.lineage_cache = lineage_cache; Ok(()) } fn lineage_for(&self, id: n00nId) -> Result { - let path = self.path_from(id)?; let node = self .sessions .get(&id) .ok_or(LineageError::UnknownSession(id))?; - let root = path - .last() - .copied() + let cached = self + .lineage_cache + .get(&id) .ok_or(LineageError::UnknownSession(id))?; - if node.root_session_id != root { - return Err(LineageError::RootMismatch { - id, - expected: root, - found: node.root_session_id, - }); - } Ok(SessionLineage { caller: id, - root, + root: cached.root, parent: node.parent_id, - depth: path.len() - 1, + depth: cached.depth, }) } @@ -488,12 +531,8 @@ impl SessionLineageGuard { return Err(LineageError::UnknownSession(start)); } let mut path = Vec::new(); - let mut seen = HashSet::new(); let mut current = start; loop { - if !seen.insert(current) { - return Err(LineageError::Cycle(current)); - } path.push(current); let parent = self .sessions @@ -503,17 +542,79 @@ impl SessionLineageGuard { let Some(parent) = parent else { return Ok(path); }; - if !self.sessions.contains_key(&parent) { - return Err(LineageError::MissingParent { - id: current, - parent, - }); - } current = parent; } } } +fn resolve_cached_lineage( + sessions: &HashMap, + cache: &mut HashMap, + start: n00nId, +) -> Result<(), LineageError> { + if cache.contains_key(&start) { + return Ok(()); + } + let mut trail = Vec::new(); + let mut seen = HashSet::new(); + let mut current = start; + loop { + if let Some(cached) = cache.get(¤t).copied() { + let mut depth = cached.depth; + for id in trail.into_iter().rev() { + depth = depth + .checked_add(1) + .ok_or(LineageError::DepthExceeded { limit: usize::MAX })?; + cache.insert( + id, + CachedLineage { + root: cached.root, + depth, + }, + ); + } + return Ok(()); + } + if !seen.insert(current) { + return Err(LineageError::Cycle(current)); + } + let node = sessions + .get(¤t) + .ok_or(LineageError::UnknownSession(current))?; + let Some(parent) = node.parent_id else { + cache.insert( + current, + CachedLineage { + root: current, + depth: 0, + }, + ); + let mut depth = 0usize; + for id in trail.into_iter().rev() { + depth = depth + .checked_add(1) + .ok_or(LineageError::DepthExceeded { limit: usize::MAX })?; + cache.insert( + id, + CachedLineage { + root: current, + depth, + }, + ); + } + return Ok(()); + }; + if !sessions.contains_key(&parent) { + return Err(LineageError::MissingParent { + id: current, + parent, + }); + } + trail.push(current); + current = parent; + } +} + fn limit_reached(committed: usize, reserved: usize, limit: usize) -> bool { committed.saturating_add(reserved) >= limit } @@ -660,6 +761,25 @@ mod tests { )); } + #[test] + fn restored_active_descendants_must_fit_limit() { + let root = id(1); + let first = id(2); + let second = id(3); + + assert!(matches!( + SessionLineageGuard::from_live( + [ + session(root, None), + session(first, Some(root)), + session(second, Some(root)), + ], + limits(4, 4, 1), + ), + Err(LineageError::ActiveDescendantsExceeded { limit: 1 }) + )); + } + #[test] fn reservation_release_is_exact_and_removal_releases_only_active_capacity() { let root = id(1); diff --git a/plugins/workflow/init.lua b/plugins/workflow/init.lua index be25aaa7c..0d971c7dd 100644 --- a/plugins/workflow/init.lua +++ b/plugins/workflow/init.lua @@ -402,14 +402,26 @@ local function new_run_id(script) return n00n.workflow.hash(script .. "\0" .. tostring(os.time()) .. "\0" .. tostring(run_seq)) end +local function utf8_prefix(text, limit) + local cut = math.min(#text, math.max(limit, 0)) + while cut > 0 do + local next_byte = text:byte(cut + 1) + if not next_byte or next_byte < 0x80 or next_byte >= 0xC0 then + break + end + cut = cut - 1 + end + return text:sub(1, cut) +end + local function bounded_text(text, limit) if #text <= limit then return text end if limit <= #RESULT_TRUNCATED_MARKER then - return RESULT_TRUNCATED_MARKER:sub(1, limit) + return utf8_prefix(RESULT_TRUNCATED_MARKER, limit) end - return text:sub(1, limit - #RESULT_TRUNCATED_MARKER) .. RESULT_TRUNCATED_MARKER + return utf8_prefix(text, limit - #RESULT_TRUNCATED_MARKER) .. RESULT_TRUNCATED_MARKER end local function parallel(fns, popts) diff --git a/site/docs/content/lua-api/_index.md b/site/docs/content/lua-api/_index.md index 0b21609c0..ed28cc3fa 100644 --- a/site/docs/content/lua-api/_index.md +++ b/site/docs/content/lua-api/_index.md @@ -6005,7 +6005,7 @@ function M.make_local_tool(schema, on_submit) -- Subagent launch helper module. -- Provides a unified interface for launching subagents with model resolution, -- system prompts, tool setup, and optional structured output validation. -M.ORCHESTRATION_TOOLS = { "task", "team", "workflow", "agent_control", "batch" } +function M.orchestration_tools() -- Launch a subagent with the given options. -- Returns (result | nil, err, cost, usage, model_spec) From 55365412043ab2aceb57d31ec97e40fd29692206 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 10 Aug 2026 01:09:59 -0400 Subject: [PATCH 16/24] test(control): validate TUI session IDs --- src/cmd/tui_bridge.rs | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/src/cmd/tui_bridge.rs b/src/cmd/tui_bridge.rs index 6efea161f..2462388b3 100644 --- a/src/cmd/tui_bridge.rs +++ b/src/cmd/tui_bridge.rs @@ -367,6 +367,20 @@ mod tests { Ok(()) } + #[test] + fn control_operations_reject_malformed_session_ids() { + let (tx, _rx) = flume::unbounded(); + let invalid = "live-a"; + let message_error = message_one(&tx, invalid, "hi", &MessageOpts::default()) + .expect_err("message must reject malformed IDs"); + let resume_error = resume_one(&tx, invalid).expect_err("resume must reject malformed IDs"); + let stop_error = stop_one(&tx, invalid).expect_err("stop must reject malformed IDs"); + + assert!(matches!(message_error, ControlError::InvalidId(id) if id == invalid)); + assert!(matches!(resume_error, ControlError::InvalidId(id) if id == invalid)); + assert!(matches!(stop_error, ControlError::InvalidId(id) if id == invalid)); + } + #[test] fn message_forwards_steer_and_control_opts() -> Result<(), String> { let (tx, rx) = flume::unbounded(); @@ -491,7 +505,7 @@ mod tests { respond_live( rx, json!([{ - "id": "live-a", + "id": "00000000-0000-7000-8000-000000000001", "title": "A", "status": "idle", "updated_at": 0, @@ -531,10 +545,11 @@ mod tests { .. } => { assert!( - agents - .iter() - .any(|a| a.id == "live-a" && a.backend == BackendKind::Tui), - "missing live-a: {agents:?}" + agents.iter().any(|a| { + a.id == "00000000-0000-7000-8000-000000000001" + && a.backend == BackendKind::Tui + }), + "missing live session: {agents:?}" ); Ok(()) } From caf8dfcb070b9a3b263b22d46ea91d14f6340bd8 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 10 Aug 2026 01:32:22 -0400 Subject: [PATCH 17/24] test(lua): stabilize deadline waiter timing --- n00n-lua/src/runtime.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index 0472ad5d5..9dd72a72a 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -3630,7 +3630,7 @@ mod tests { smol::block_on(async { let handle = Arc::new(Mutex::new(TaskCell::new( CancelToken::none(), - Some(Instant::now() + Duration::from_millis(20)), + Some(Instant::now() + Duration::from_millis(500)), None, None, ))); @@ -3641,7 +3641,7 @@ mod tests { smol::Timer::after(Duration::from_millis(10)).await; lock_cell(&handle).deadline.set(None); - smol::Timer::after(Duration::from_millis(30)).await; + smol::Timer::after(Duration::from_millis(550)).await; assert!(!waiter.is_finished()); lock_cell(&handle) From 7a9e74680dad6fcb3f93d192e1e86e33c9e1bccc Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 10 Aug 2026 03:13:40 -0400 Subject: [PATCH 18/24] fix(orchestration): persist queued session lifecycle --- n00n-agent/src/types.rs | 3 + n00n-storage/src/sessions.rs | 11 ++ n00n-ui/src/agent/agent_loop.rs | 8 + n00n-ui/src/agent/shared_queue.rs | 49 ++++- n00n-ui/src/app/mod.rs | 23 ++- n00n-ui/src/app/queue.rs | 1 - n00n-ui/src/app/session.rs | 1 + n00n-ui/src/chat.rs | 2 +- n00n-ui/src/event_loop.rs | 313 ++++++++++++++++++++++++------ n00n-ui/src/session_lineage.rs | 124 ++++++++++-- n00n-ui/src/storage_writer.rs | 45 +++++ src/print.rs | 1 + src/sdk_mode.rs | 1 + 13 files changed, 483 insertions(+), 99 deletions(-) diff --git a/n00n-agent/src/types.rs b/n00n-agent/src/types.rs index dd92b70b7..98b88ac67 100644 --- a/n00n-agent/src/types.rs +++ b/n00n-agent/src/types.rs @@ -1039,6 +1039,9 @@ pub enum AgentEvent { images: Vec, control: bool, }, + QueueDrained { + generation: u64, + }, Done { usage: TokenUsage, num_turns: u32, diff --git a/n00n-storage/src/sessions.rs b/n00n-storage/src/sessions.rs index 2d510a1db..0b12c2cbd 100644 --- a/n00n-storage/src/sessions.rs +++ b/n00n-storage/src/sessions.rs @@ -934,6 +934,10 @@ pub struct SessionMeta { pub queued_submissions: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub queued_direct_tools: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub direct_output: Option, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub direct_output_is_error: bool, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub subagents: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -4181,6 +4185,8 @@ mod tests { session.meta.input_draft = Some("draft line".into()); session.meta.queued_messages = vec!["queued".into()]; + session.meta.direct_output = Some("bootstrap output".into()); + session.meta.direct_output_is_error = true; session.title = "updated title".into(); session.updated_at = now_epoch() + 1; log.append(&session).unwrap(); @@ -4188,6 +4194,11 @@ mod tests { let loaded = TestSession::load_from(session.id, dir).unwrap(); assert_eq!(loaded.meta.input_draft.as_deref(), Some("draft line")); assert_eq!(loaded.meta.queued_messages, vec!["queued".to_string()]); + assert_eq!( + loaded.meta.direct_output.as_deref(), + Some("bootstrap output") + ); + assert!(loaded.meta.direct_output_is_error); assert_eq!(loaded.title, "updated title"); } diff --git a/n00n-ui/src/agent/agent_loop.rs b/n00n-ui/src/agent/agent_loop.rs index ab53148d2..511ec996a 100644 --- a/n00n-ui/src/agent/agent_loop.rs +++ b/n00n-ui/src/agent/agent_loop.rs @@ -160,12 +160,20 @@ impl AgentLoop { } while let Ok(()) = self.queue.recv_notify().await { + let mut last_run_id = None; while let Some(entry) = self.queue.pop() { if entry.run_id() < self.min_run_id { continue; } + last_run_id = Some(entry.run_id()); self.process_entry(entry).await; } + if let Some(run_id) = last_run_id + && let Some(generation) = self.queue.drain_generation() + { + let event_tx = EventSender::new(self.agent_tx.clone(), run_id); + let _ = event_tx.send(AgentEvent::QueueDrained { generation }); + } } } diff --git a/n00n-ui/src/agent/shared_queue.rs b/n00n-ui/src/agent/shared_queue.rs index 48631e788..e98dad38b 100644 --- a/n00n-ui/src/agent/shared_queue.rs +++ b/n00n-ui/src/agent/shared_queue.rs @@ -10,7 +10,7 @@ use std::borrow::Cow; use std::collections::VecDeque; use std::sync::{ Arc, Mutex, MutexGuard, PoisonError, - atomic::{AtomicBool, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, }; use n00n_agent::{ @@ -151,29 +151,40 @@ fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { #[derive(Clone)] pub(crate) struct QueueSender { items: Items, + generation: Arc, notify_tx: flume::Sender<()>, } pub(crate) struct QueueReceiver { items: Items, + generation: Arc, notify_rx: flume::Receiver<()>, } pub(crate) fn queue() -> (QueueSender, QueueReceiver) { let (notify_tx, notify_rx) = flume::bounded(1); let items: Items = Arc::new(Mutex::new(VecDeque::new())); + let generation = Arc::new(AtomicU64::new(0)); ( QueueSender { items: Arc::clone(&items), + generation: Arc::clone(&generation), notify_tx, }, - QueueReceiver { items, notify_rx }, + QueueReceiver { + items, + generation, + notify_rx, + }, ) } impl QueueSender { pub(crate) fn push(&self, entry: QueueItem) { - lock(&self.items).push_back(entry); + let mut items = lock(&self.items); + items.push_back(entry); + self.generation.fetch_add(1, Ordering::Release); + drop(items); let _ = self.notify_tx.try_send(()); } @@ -189,6 +200,7 @@ impl QueueSender { return; } items.push_front(entry); + self.generation.fetch_add(1, Ordering::Release); drop(items); let _ = self.notify_tx.try_send(()); } @@ -212,6 +224,7 @@ impl QueueSender { let mut items = lock(&self.items); let item_index = Self::panel_index(&items, index).unwrap_or_else(|| items.len()); items.insert(item_index, entry); + self.generation.fetch_add(1, Ordering::Release); } pub(crate) fn promote_latest_steering(&self) -> bool { @@ -228,14 +241,19 @@ impl QueueSender { return false; }; *delivery = Delivery::Immediate; + self.generation.fetch_add(1, Ordering::Release); true } - #[cfg(test)] pub(crate) fn is_empty(&self) -> bool { lock(&self.items).is_empty() } + pub(crate) fn is_drained(&self, generation: u64) -> bool { + let items = lock(&self.items); + items.is_empty() && self.generation.load(Ordering::Acquire) == generation + } + pub(crate) fn clear(&self) { lock(&self.items).clear(); } @@ -281,6 +299,7 @@ impl QueueSender { // Update input and ready flag atomically while holding the lock *queued_input = input; ready.store(true, Ordering::Release); + self.generation.fetch_add(1, Ordering::Release); drop(items); let _ = self.notify_tx.try_send(()); true @@ -357,6 +376,13 @@ impl QueueReceiver { items.remove(index) } + pub(crate) fn drain_generation(&self) -> Option { + let items = lock(&self.items); + items + .is_empty() + .then(|| self.generation.load(Ordering::Acquire)) + } + pub(crate) async fn recv_notify(&self) -> Result<(), flume::RecvError> { self.notify_rx.recv_async().await } @@ -491,6 +517,21 @@ mod tests { ); } + #[test] + fn stale_drain_generation_is_rejected_after_new_work() { + let (tx, rx) = queue(); + tx.push(msg(false)); + assert!(rx.pop().is_some()); + let drained = rx.drain_generation().expect("drained generation"); + assert!(tx.is_drained(drained)); + + tx.push(msg(false)); + assert!(rx.pop().is_some()); + assert!(!tx.is_drained(drained)); + let latest = rx.drain_generation().expect("latest generation"); + assert!(tx.is_drained(latest)); + } + #[test] fn direct_tools_are_available_for_persistence() { let (tx, _rx) = queue(); diff --git a/n00n-ui/src/app/mod.rs b/n00n-ui/src/app/mod.rs index 52a476d58..e55830681 100644 --- a/n00n-ui/src/app/mod.rs +++ b/n00n-ui/src/app/mod.rs @@ -1319,6 +1319,14 @@ impl App { } pub(crate) fn handle_submission_persistence_failure(&mut self, dispatch: &SubmissionDispatch) { + self.handle_submission_failure(dispatch, PERSISTENCE_FAILURE_MSG); + } + + pub(crate) fn handle_submission_failure( + &mut self, + dispatch: &SubmissionDispatch, + message: &str, + ) { self.queue.remove_submission(dispatch.submission_id); if dispatch.paint_required { let Some(pending) = self.pending_submission.as_ref() else { @@ -1331,7 +1339,7 @@ impl App { return; } if self.restore_pending_submission(dispatch.submission_id, dispatch.run_id) { - self.flash(PERSISTENCE_FAILURE_MSG.into()); + self.flash(message.into()); } return; } @@ -1339,15 +1347,10 @@ impl App { let relevant_run = dispatch.run_id == self.run_id; let _ = dispatch.gate.try_cancel(); if relevant_run && self.status == Status::Streaming { - self.status = Status::error(PERSISTENCE_FAILURE_MSG.into()); - self.main_chat().push(DisplayMessage::new( - DisplayRole::Error, - PERSISTENCE_FAILURE_MSG.into(), - )); - self.fire_session_autocmd( - "TurnError", - serde_json::json!({ "message": PERSISTENCE_FAILURE_MSG }), - ); + self.status = Status::error(message.into()); + self.main_chat() + .push(DisplayMessage::new(DisplayRole::Error, message.into())); + self.fire_session_autocmd("TurnError", serde_json::json!({ "message": message })); } } pub(crate) fn preserve_submission_for_shutdown(&mut self, dispatch: SubmissionDispatch) { diff --git a/n00n-ui/src/app/queue.rs b/n00n-ui/src/app/queue.rs index 6a55b1f6f..f58ed1fa9 100644 --- a/n00n-ui/src/app/queue.rs +++ b/n00n-ui/src/app/queue.rs @@ -30,7 +30,6 @@ impl MessageQueue { self.shared = Some(shared); } - #[cfg(test)] pub(crate) fn is_empty(&self) -> bool { self.shared .as_ref() diff --git a/n00n-ui/src/app/session.rs b/n00n-ui/src/app/session.rs index fdb73bf69..6217ccfda 100644 --- a/n00n-ui/src/app/session.rs +++ b/n00n-ui/src/app/session.rs @@ -56,6 +56,7 @@ pub(crate) fn session_has_content(session: &AppSession) -> bool { || !session.meta.queued_messages.is_empty() || !session.meta.queued_submissions.is_empty() || !session.meta.queued_direct_tools.is_empty() + || session.meta.direct_output.is_some() || session.meta.mode != Some(n00n_storage::sessions::StoredMode::Build) || session.meta.plan_path.is_some() || session.meta.plan_written diff --git a/n00n-ui/src/chat.rs b/n00n-ui/src/chat.rs index c42791436..e4b9fea81 100644 --- a/n00n-ui/src/chat.rs +++ b/n00n-ui/src/chat.rs @@ -247,7 +247,7 @@ impl Chat { "Model stalled after tool calls, nudging...".into(), )); } - AgentEvent::SubagentHistory { .. } => {} + AgentEvent::QueueDrained { .. } | AgentEvent::SubagentHistory { .. } => {} AgentEvent::LiveToolBuf { id, body } => { self.messages_panel.register_live_buf(id, body); } diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index bc2c4decf..0e5d898d7 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -163,6 +163,19 @@ fn has_restorable_work(session: &AppSession) -> bool { || !session.meta.queued_direct_tools.is_empty() } +fn cancel_stored_session(session: &mut AppSession) -> bool { + let had_work = session.meta.lifecycle.is_active() + || !session.meta.queued_messages.is_empty() + || !session.meta.queued_submissions.is_empty() + || !session.meta.queued_direct_tools.is_empty(); + session.meta.lifecycle = StoredSessionLifecycle::Cancelled; + session.meta.queued_messages.clear(); + session.meta.queued_submissions.clear(); + session.meta.queued_direct_tools.clear(); + session.updated_at = n00n_storage::now_epoch(); + had_work +} + fn resolved_root( start: n00nId, parents: &HashMap>, @@ -202,6 +215,40 @@ fn session_identity(session: &AppSession) -> std::result::Result, +) -> u64 { + let Some(snapshot) = snapshot else { + return 0; + }; + let Some(revision) = snapshot.state_revision() else { + return 0; + }; + revision +} + +fn capture_session_plugin_state( + handle: &EventHandle, + session: &mut AppSession, +) -> std::result::Result<(), String> { + let identity = session_identity(session).map_err(|error| error.to_string())?; + let persisted_revision = state_revision_or_initial(session.meta.state_snapshot.as_ref()); + let revision = session.meta.revision.max( + persisted_revision + .checked_add(1) + .ok_or_else(|| "plugin state revision exhausted".to_owned())?, + ); + let snapshot = handle + .capture_state(&identity, revision) + .map_err(|error| error.to_string())?; + let captured_revision = snapshot + .state_revision() + .ok_or_else(|| "captured plugin state has no revision".to_owned())?; + session.meta.revision = session.meta.revision.max(captured_revision); + session.meta.state_snapshot = Some(snapshot); + Ok(()) +} + fn warn_lineage_cleanup( result: std::result::Result, session_id: n00nId, @@ -267,8 +314,6 @@ struct SessionRuntime { shell_rx: flume::Receiver, last_status: SessionStatus, direct_bootstrap_active: bool, - direct_output: Option, - direct_output_is_error: bool, } impl SessionRuntime { @@ -382,8 +427,6 @@ impl SpawnCtx { shell_rx, last_status, direct_bootstrap_active, - direct_output: None, - direct_output_is_error: false, }) } } @@ -414,6 +457,7 @@ pub(crate) struct EventLoop<'t> { struct SubmissionPersistence { session_id: n00nId, dispatch: SubmissionDispatch, + execution_started: bool, result: Result<(), SessionError>, } @@ -931,6 +975,16 @@ impl<'t> EventLoop<'t> { } fn handle_agent(&mut self, idx: usize, envelope: Box) { + if let n00n_agent::AgentEvent::QueueDrained { generation } = &envelope.event { + if self.sessions[idx].handles.queue.is_drained(*generation) { + let id = self.sessions[idx].id(); + if let Err(error) = self.lineage.set_execution_active(id, false) { + warn!(session_id = %id, error = %error, "failed to release drained session activity"); + } + self.sessions[idx].app.save_session(); + } + return; + } if envelope.run_id != self.sessions[idx].app.run_id { let actions = self.sessions[idx].app.update(Msg::Agent(envelope)); self.dispatch(idx, actions); @@ -939,14 +993,31 @@ impl<'t> EventLoop<'t> { if self.sessions[idx].direct_bootstrap_active { match &envelope.event { n00n_agent::AgentEvent::ToolDone(done) => { - self.sessions[idx].direct_output = Some(done.output.as_text()); - self.sessions[idx].direct_output_is_error = done.is_error; + self.sessions[idx].app.state.session.meta.direct_output = + Some(done.output.as_text()); + self.sessions[idx] + .app + .state + .session + .meta + .direct_output_is_error = done.is_error; } n00n_agent::AgentEvent::Error { message } - if self.sessions[idx].direct_output.is_none() => + if self.sessions[idx] + .app + .state + .session + .meta + .direct_output + .is_none() => { - self.sessions[idx].direct_output = Some(message.clone()); - self.sessions[idx].direct_output_is_error = true; + self.sessions[idx].app.state.session.meta.direct_output = Some(message.clone()); + self.sessions[idx] + .app + .state + .session + .meta + .direct_output_is_error = true; } _ => {} } @@ -961,7 +1032,10 @@ impl<'t> EventLoop<'t> { } n00n_agent::AgentEvent::ToolStart(_) | n00n_agent::AgentEvent::TextDelta { .. } - | n00n_agent::AgentEvent::ThinkingDelta { .. } => Some(StoredSessionLifecycle::Running), + | n00n_agent::AgentEvent::ThinkingDelta { .. } + | n00n_agent::AgentEvent::QueueItemConsumed { .. } => { + Some(StoredSessionLifecycle::Running) + } _ => None, }; let capture = matches!( @@ -989,10 +1063,6 @@ impl<'t> EventLoop<'t> { .meta .queued_direct_tools .clear(); - let id = self.sessions[idx].id(); - if let Err(error) = self.lineage.set_execution_active(id, false) { - warn!(session_id = %id, error = %error, "failed to release session lineage activity"); - } self.sessions[idx].app.save_session(); } } @@ -1000,26 +1070,52 @@ impl<'t> EventLoop<'t> { } fn capture_plugin_state(&mut self, idx: usize) -> std::result::Result<(), String> { - let Some(handle) = &self.ctx.lua_event_handle else { + let Some(handle) = self.ctx.lua_event_handle.clone() else { return Ok(()); }; - let session = &self.sessions[idx].app.state.session; - let identity = session_identity(session).map_err(|error| error.to_string())?; - let revision = match session + let root_id = self.sessions[idx] + .app + .state + .session .meta - .state_snapshot - .as_ref() - .and_then(n00n_storage::sessions::StoredSessionStateSnapshot::state_revision) + .root_session_id + .map_or(self.sessions[idx].id(), |root| root); + capture_session_plugin_state(&handle, &mut self.sessions[idx].app.state.session)?; + if root_id == self.sessions[idx].id() { + return Ok(()); + } + + let mut root = if let Some(root_idx) = self.position(root_id) { + self.sessions[root_idx].app.session_snapshot() + } else if let Some(root) = self + .ctx + .storage_writer + .latest_snapshot(root_id) + .map_err(|error| error.to_string())? { - Some(revision) => revision - .checked_add(1) - .ok_or_else(|| "plugin state revision exhausted".to_owned())?, - None => 1, + root + } else { + AppSession::load(root_id, &self.ctx.storage).map_err(|error| error.to_string())? }; - let snapshot = handle - .capture_state(&identity, revision) - .map_err(|error| error.to_string())?; - self.sessions[idx].app.state.session.meta.state_snapshot = Some(snapshot); + root.meta.revision = root + .meta + .revision + .checked_add(1) + .ok_or_else(|| "root session revision exhausted".to_owned())?; + root.updated_at = n00n_storage::now_epoch(); + capture_session_plugin_state(&handle, &mut root)?; + if let Some(root_idx) = self.position(root_id) { + self.sessions[root_idx] + .app + .state + .session + .meta + .state_snapshot + .clone_from(&root.meta.state_snapshot); + self.sessions[root_idx].app.state.session.meta.revision = root.meta.revision; + self.sessions[root_idx].app.state.session.updated_at = root.updated_at; + } + self.ctx.storage_writer.send(Box::new(root)); Ok(()) } @@ -1240,11 +1336,12 @@ impl<'t> EventLoop<'t> { .then(|| message.first_text_content()) .flatten() }); - let output = assistant_output.or(rt.direct_output.as_deref()); + let direct_output = rt.app.state.session.meta.direct_output.as_deref(); + let output = assistant_output.or(direct_output); let direct_error = assistant_output .is_none() - .then_some(rt.direct_output_is_error) - .filter(|_| rt.direct_output.is_some()); + .then_some(rt.app.state.session.meta.direct_output_is_error) + .filter(|_| direct_output.is_some()); let paused_team = paused_team_run(&history); Ok(json!({ "id": rt.id(), @@ -1273,9 +1370,10 @@ impl<'t> EventLoop<'t> { let reply = (|| { let caller = caller_session_id(caller_id)?; let explicit_parent = parent_id.as_deref().map(parse_session_id).transpose()?; + let execution_active = prompt.is_some() || bootstrap.is_some(); let reservation = self .lineage - .reserve_new(caller, explicit_parent) + .reserve_new(caller, explicit_parent, execution_active) .map_err(|error| error.to_string())?; let caller_lineage = match self.lineage.lineage(caller) { Ok(lineage) => lineage, @@ -1343,6 +1441,9 @@ impl<'t> EventLoop<'t> { }); Ok(json!("started")) } else if let Some(prompt) = prompt { + self.lineage + .set_execution_active(id, false) + .map_err(|error| error.to_string())?; self.submit_text(idx, prompt, false, false) } else { self.sessions[idx].app.state.session.meta.lifecycle = @@ -1441,12 +1542,17 @@ impl<'t> EventLoop<'t> { let mut cancelled = false; for session_id in targets { let Some(idx) = self.position(session_id) else { - let mut session = AppSession::load(session_id, &self.ctx.storage) - .map_err(|error| error.to_string())?; - cancelled |= session.meta.lifecycle.is_active(); - session.meta.lifecycle = StoredSessionLifecycle::Cancelled; - session.meta.queued_direct_tools.clear(); - session.updated_at = n00n_storage::now_epoch(); + let mut session = match self + .ctx + .storage_writer + .latest_snapshot(session_id) + .map_err(|error| error.to_string())? + { + Some(session) => session, + None => AppSession::load(session_id, &self.ctx.storage) + .map_err(|error| error.to_string())?, + }; + cancelled |= cancel_stored_session(&mut session); self.ctx.storage_writer.send(Box::new(session)); warn_lineage_cleanup( self.lineage.set_execution_active(session_id, false), @@ -1463,6 +1569,8 @@ impl<'t> EventLoop<'t> { .meta .lifecycle .is_active() + || !self.sessions[idx].app.queue.is_empty() + || has_restorable_work(&self.sessions[idx].app.state.session) { let actions = self.sessions[idx].app.cancel_current_run(); self.dispatch(idx, actions); @@ -1758,12 +1866,28 @@ impl<'t> EventLoop<'t> { if completion.result.is_err() { rt.app .handle_submission_persistence_failure(&completion.dispatch); + if completion.execution_started && rt.app.queue.is_empty() { + warn_lineage_cleanup( + self.lineage + .set_execution_active(completion.session_id, false), + completion.session_id, + "release failed submission activity", + ); + } return; } if !rt.app.accepts_submission_persistence(&completion.dispatch) { rt.app .queue .remove_submission(completion.dispatch.submission_id); + if completion.execution_started && rt.app.queue.is_empty() { + warn_lineage_cleanup( + self.lineage + .set_execution_active(completion.session_id, false), + completion.session_id, + "release superseded submission activity", + ); + } return; } let submission_id = completion.dispatch.submission_id; @@ -1773,15 +1897,40 @@ impl<'t> EventLoop<'t> { .mark_submission_ready(submission_id, completion.dispatch.input) { rt.app.queue.remove_submission(submission_id); + if completion.execution_started && rt.app.queue.is_empty() { + warn_lineage_cleanup( + self.lineage + .set_execution_active(completion.session_id, false), + completion.session_id, + "release removed submission activity", + ); + } } } fn handle_action(&mut self, idx: usize, action: Action) { match action { Action::SendMessage(mut dispatch) => { + let session_id = self.sessions[idx].id(); + let execution_started = match self.lineage.begin_execution(session_id) { + Ok(started) => started, + Err(error) => { + self.sessions[idx] + .app + .handle_submission_failure(&dispatch, &error.to_string()); + return; + } + }; let rt = &mut self.sessions[idx]; if !rt.app.stage_submission_preamble(&mut dispatch) { rt.app.queue.remove_submission(dispatch.submission_id); + if execution_started { + warn_lineage_cleanup( + self.lineage.set_execution_active(session_id, false), + session_id, + "release rejected submission activity", + ); + } return; } let session_id = rt.app.state.session.id; @@ -1793,38 +1942,35 @@ impl<'t> EventLoop<'t> { let _ = completion_tx.send(SubmissionPersistence { session_id, dispatch: *dispatch, + execution_started, result, }); }); } Action::CancelAgent { run_id } => { let id = self.sessions[idx].id(); - match self.sessions[idx] + if let Err(error) = self.sessions[idx] .handles .cmd_tx .try_send(AgentCommand::Cancel { run_id }) { - Ok(()) => { - self.sessions[idx].app.state.session.meta.lifecycle = - StoredSessionLifecycle::Cancelled; - self.sessions[idx] - .app - .state - .session - .meta - .queued_direct_tools - .clear(); - warn_lineage_cleanup( - self.lineage.set_execution_active(id, false), - id, - "clear keyboard-cancelled activity", - ); - self.sessions[idx].app.save_session(); - } - Err(error) => { - warn!(session_id = %id, %error, "failed to send agent cancellation"); - } + warn!(session_id = %id, %error, "failed to send agent cancellation"); } + self.sessions[idx].app.state.session.meta.lifecycle = + StoredSessionLifecycle::Cancelled; + self.sessions[idx] + .app + .state + .session + .meta + .queued_direct_tools + .clear(); + warn_lineage_cleanup( + self.lineage.set_execution_active(id, false), + id, + "clear keyboard-cancelled activity", + ); + self.sessions[idx].app.save_session(); } Action::CancelSubagent { tool_use_id } => { let _ = self.sessions[idx] @@ -2130,12 +2276,16 @@ fn scroll_delta(kind: MouseEventKind, lines: u32) -> i32 { #[cfg(test)] mod tests { use super::{ - DRAIN_BUDGET, DrainScheduler, TEAM_TOOL_NAME, draw_then_post_terminal, paused_team_run, - should_save_periodically, take_painted_submissions, + DRAIN_BUDGET, DrainScheduler, TEAM_TOOL_NAME, cancel_stored_session, + draw_then_post_terminal, paused_team_run, should_save_periodically, + take_painted_submissions, }; - use crate::components::Status; + use crate::{AppSession, components::Status}; use n00n_providers::{ContentBlock, Message, Role}; - use n00n_storage::id::n00nId; + use n00n_storage::{ + id::n00nId, + sessions::{StoredDelivery, StoredDirectTool, StoredQueuedMessage, StoredSessionLifecycle}, + }; use ratatui::{ Terminal, backend::{Backend, ClearType, TestBackend, WindowSize}, @@ -2194,6 +2344,41 @@ mod tests { assert!(paused_team_run(&[tool_call, tool_result]).is_none()); } + #[test] + fn cancel_stored_session_clears_all_persisted_work() { + let mut session = AppSession::new("model", "/project"); + session.meta.lifecycle = StoredSessionLifecycle::Running; + session.meta.queued_messages = vec!["legacy".into()]; + session.meta.queued_submissions = vec![StoredQueuedMessage { + text: "queued".into(), + images: Vec::new(), + mode: None, + plan_path: None, + thinking: None, + fast: false, + workflow: false, + control: false, + delivery: StoredDelivery::TurnEnd, + prompt: None, + }]; + session.meta.queued_direct_tools = vec![StoredDirectTool { + tool: "task".into(), + input: serde_json::json!({}), + }]; + + assert!(cancel_stored_session(&mut session)); + assert_eq!(session.meta.lifecycle, StoredSessionLifecycle::Cancelled); + assert!(session.meta.queued_messages.is_empty()); + assert!(session.meta.queued_submissions.is_empty()); + assert!(session.meta.queued_direct_tools.is_empty()); + + let mut inactive = AppSession::new("model", "/project"); + inactive.meta.lifecycle = StoredSessionLifecycle::Succeeded; + inactive.meta.queued_messages = vec!["pending".into()]; + assert!(cancel_stored_session(&mut inactive)); + assert!(inactive.meta.queued_messages.is_empty()); + } + struct FailingBackend(TestBackend); fn infallible(result: Result) -> T { diff --git a/n00n-ui/src/session_lineage.rs b/n00n-ui/src/session_lineage.rs index 66440dbba..896fad386 100644 --- a/n00n-ui/src/session_lineage.rs +++ b/n00n-ui/src/session_lineage.rs @@ -95,6 +95,7 @@ struct PendingReservation { parent: n00nId, root: n00nId, depth: usize, + execution_active: bool, } #[derive(Debug, Clone, Copy)] @@ -274,6 +275,7 @@ impl SessionLineageGuard { &mut self, caller: n00nId, explicit_parent: Option, + execution_active: bool, ) -> Result { let caller_lineage = self.lineage(caller)?; let parent = match explicit_parent { @@ -305,11 +307,20 @@ impl SessionLineageGuard { limit: self.limits.max_total_descendants, }); } - if limit_reached( - counts.active, - counts.reserved, - self.limits.max_active_descendants, - ) { + let active_reservations = self + .reservations + .values() + .filter(|reservation| { + reservation.root == caller_lineage.root && reservation.execution_active + }) + .count(); + if execution_active + && limit_reached( + counts.active, + active_reservations, + self.limits.max_active_descendants, + ) + { return Err(LineageError::ActiveDescendantsExceeded { limit: self.limits.max_active_descendants, }); @@ -327,6 +338,7 @@ impl SessionLineageGuard { parent, root: caller_lineage.root, depth, + execution_active, }, ); Ok(NewReservation { id }) @@ -354,7 +366,7 @@ impl SessionLineageGuard { root_session_id: pending.root, parent_id: Some(pending.parent), runtime_present: true, - execution_active: true, + execution_active: pending.execution_active, deleted: false, }, ); @@ -400,7 +412,9 @@ impl SessionLineageGuard { .collect::>(); let mut descendants = Vec::new(); while let Some(id) = pending.pop() { - descendants.push(id); + if self.sessions.get(&id).is_some_and(|node| !node.deleted) { + descendants.push(id); + } if let Some(children) = self.children.get(&id) { pending.extend(children.iter().copied()); } @@ -666,14 +680,18 @@ mod tests { ) .expect("valid roots"); - let reservation = guard.reserve_new(root_a, None).expect("root A capacity"); + let reservation = guard + .reserve_new(root_a, None, true) + .expect("root A capacity"); guard.commit_new(reservation, child_a).expect("child A"); assert!(matches!( - guard.reserve_new(root_a, None), + guard.reserve_new(root_a, None, true), Err(LineageError::TotalDescendantsExceeded { .. }) )); - let reservation = guard.reserve_new(root_b, None).expect("root B capacity"); + let reservation = guard + .reserve_new(root_b, None, true) + .expect("root B capacity"); guard.commit_new(reservation, child_b).expect("child B"); assert_eq!(guard.descendant_counts(root_a).expect("counts").total, 1); assert_eq!(guard.descendant_counts(root_b).expect("counts").total, 1); @@ -695,11 +713,11 @@ mod tests { .expect("valid graph"); assert!(matches!( - guard.reserve_new(root, Some(foreign)), + guard.reserve_new(root, Some(foreign), true), Err(LineageError::ParentMismatch) )); assert!(matches!( - guard.reserve_new(id(99), None), + guard.reserve_new(id(99), None, true), Err(LineageError::CallerNotLive(_)) )); assert_eq!( @@ -741,12 +759,12 @@ mod tests { limits(2, 2, 2), ) .expect("valid graph"); - let reservation = guard.reserve_new(child, None).expect("depth one"); + let reservation = guard.reserve_new(child, None, true).expect("depth one"); guard .commit_new(reservation, grandchild) .expect("grandchild"); assert!(matches!( - guard.reserve_new(grandchild, None), + guard.reserve_new(grandchild, None, true), Err(LineageError::DepthExceeded { limit: 2 }) )); @@ -756,11 +774,48 @@ mod tests { ) .expect("valid graph"); assert!(matches!( - active_limited.reserve_new(root, None), + active_limited.reserve_new(root, None, true), Err(LineageError::ActiveDescendantsExceeded { limit: 1 }) )); } + #[test] + fn idle_reservation_does_not_consume_active_capacity() { + let root = id(1); + let active_child = id(2); + let idle_child = id(3); + let mut guard = SessionLineageGuard::from_live( + [session(root, None), session(active_child, Some(root))], + limits(4, 3, 1), + ) + .expect("valid graph"); + + assert!(matches!( + guard.reserve_new(root, None, true), + Err(LineageError::ActiveDescendantsExceeded { limit: 1 }) + )); + let reservation = guard.reserve_new(root, None, false).expect("idle capacity"); + guard + .commit_new(reservation, idle_child) + .expect("idle child"); + assert_eq!( + guard.descendant_counts(root).expect("counts"), + DescendantCounts { + total: 2, + active: 1, + reserved: 0, + } + ); + assert!(matches!( + guard.begin_execution(idle_child), + Err(LineageError::ActiveDescendantsExceeded { limit: 1 }) + )); + guard + .set_execution_active(active_child, false) + .expect("release active child"); + assert!(guard.begin_execution(idle_child).expect("start idle child")); + } + #[test] fn restored_active_descendants_must_fit_limit() { let root = id(1); @@ -786,7 +841,7 @@ mod tests { let child = id(2); let mut guard = SessionLineageGuard::from_live([session(root, None)], limits(4, 1, 1)) .expect("valid root"); - let reservation = guard.reserve_new(root, None).expect("reserve"); + let reservation = guard.reserve_new(root, None, true).expect("reserve"); assert_eq!( guard.descendant_counts(root).expect("counts"), DescendantCounts { @@ -805,7 +860,7 @@ mod tests { } ); - let reservation = guard.reserve_new(root, None).expect("reserve again"); + let reservation = guard.reserve_new(root, None, true).expect("reserve again"); guard.commit_new(reservation, child).expect("commit"); guard.remove_runtime(child).expect("remove"); assert_eq!( @@ -817,7 +872,7 @@ mod tests { } ); assert!(matches!( - guard.reserve_new(root, None), + guard.reserve_new(root, None, true), Err(LineageError::TotalDescendantsExceeded { limit: 1 }) )); } @@ -827,11 +882,42 @@ mod tests { let root = id(1); let mut guard = SessionLineageGuard::from_live([session(root, None)], limits(4, 1, 1)) .expect("valid root"); - let reservation = guard.reserve_new(root, None).expect("reserve"); + let reservation = guard.reserve_new(root, None, true).expect("reserve"); assert!(matches!( guard.commit_new(reservation, root), Err(LineageError::DuplicateSession(_)) )); assert_eq!(guard.descendant_counts(root).expect("counts").reserved, 0); } + + #[test] + fn descendants_of_omits_tombstoned_sessions() { + let root = id(1); + let child = id(2); + let grandchild = id(3); + let sibling = id(4); + let grandchild_session = LiveSession { + id: grandchild, + root_session_id: root, + parent_id: Some(child), + runtime_present: true, + execution_active: true, + }; + let mut guard = SessionLineageGuard::from_live( + [ + session(root, None), + session(child, Some(root)), + grandchild_session, + session(sibling, Some(root)), + ], + limits(4, 4, 4), + ) + .expect("valid graph"); + + guard.remove_sessions(&[child, grandchild]); + assert_eq!( + guard.descendants_of(root).expect("descendants"), + vec![sibling] + ); + } } diff --git a/n00n-ui/src/storage_writer.rs b/n00n-ui/src/storage_writer.rs index ea74dedd7..814ab5789 100644 --- a/n00n-ui/src/storage_writer.rs +++ b/n00n-ui/src/storage_writer.rs @@ -114,6 +114,10 @@ enum Op { generation: u64, done: DeleteCallback, }, + Latest { + id: n00nId, + done: flume::Sender>, + }, #[cfg(test)] Pause { entered: flume::Sender<()>, @@ -197,6 +201,14 @@ impl StorageWriter { state.collect_barriers(id, &writer_tracker); done(result); } + Op::Latest { id, done } => { + state.stage_snapshots(take_snapshots(&writer_inbox), &writer_tracker); + let session = state + .latest_snapshots + .get(&id) + .map(|snapshot| (*snapshot.session).clone()); + let _ = done.send(session); + } #[cfg(test)] Op::Pause { entered, release } => { let _ = entered.send(()); @@ -242,6 +254,14 @@ impl StorageWriter { self.enqueue_reserved_snapshot(generation, session); } + pub(crate) fn latest_snapshot(&self, id: n00nId) -> Result, SessionError> { + let (done_tx, done_rx) = flume::bounded(1); + self.ops + .send(Op::Latest { id, done: done_tx }) + .map_err(|_| writer_gone())?; + done_rx.recv().map_err(|_| writer_gone()) + } + /// Persists this snapshot before invoking `done` on the writer thread. pub fn persist( &self, @@ -879,6 +899,31 @@ mod tests { assert_eq!(AppSession::load(b_id, &dir).unwrap().title, "renamed"); } + #[test] + fn latest_snapshot_includes_queued_writes() { + let (_tmp, dir) = state_dir(); + let writer = StorageWriter::new(dir.clone()).unwrap(); + let release = pause_writer(&writer); + let mut session = AppSession::new("test-model", "/tmp/latest"); + let id = session.id; + session.title = "first".into(); + writer.send(Box::new(session.clone())); + session.title = "latest".into(); + writer.send(Box::new(session)); + + let latest = std::thread::scope(|scope| { + let query = scope.spawn(|| writer.latest_snapshot(id).unwrap()); + release.send(()).unwrap(); + query.join().unwrap() + }); + let latest = match latest { + Some(session) => session, + None => AppSession::load(id, &dir).unwrap(), + }; + assert_eq!(latest.title, "latest"); + writer.shutdown(DRAIN_TIMEOUT).unwrap(); + } + #[test] fn blocked_writer_coalesces_same_session_snapshots_and_persists_latest() { let (_tmp, dir) = state_dir(); diff --git a/src/print.rs b/src/print.rs index 28a5895d1..2e21ce08b 100644 --- a/src/print.rs +++ b/src/print.rs @@ -367,6 +367,7 @@ fn handle_print_event( | AgentEvent::ToolOutput { .. } | AgentEvent::ToolDone(_) | AgentEvent::QueueItemConsumed { .. } + | AgentEvent::QueueDrained { .. } | AgentEvent::AutoCompacting | AgentEvent::CompactionDone | AgentEvent::FusionPhase { .. } diff --git a/src/sdk_mode.rs b/src/sdk_mode.rs index 2d6262c4f..5daff8644 100644 --- a/src/sdk_mode.rs +++ b/src/sdk_mode.rs @@ -1055,6 +1055,7 @@ impl EventPump { | AgentEvent::ToolOutput { .. } | AgentEvent::ToolDone(_) | AgentEvent::QueueItemConsumed { .. } + | AgentEvent::QueueDrained { .. } | AgentEvent::AutoCompacting | AgentEvent::CompactionDone | AgentEvent::FusionPhase { .. } From cc88cf564e7f6053cf99edb41f95f74e8df301cf Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 10 Aug 2026 03:45:01 -0400 Subject: [PATCH 19/24] test(lua): widen deadline cleanup timing margin --- n00n-lua/tests/plugin_host.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/n00n-lua/tests/plugin_host.rs b/n00n-lua/tests/plugin_host.rs index 61a6ae31e..5968b76c1 100644 --- a/n00n-lua/tests/plugin_host.rs +++ b/n00n-lua/tests/plugin_host.rs @@ -3057,7 +3057,7 @@ fn ctx_set_deadline_normalizes_watchdog_error() { fn caught_deadline_interrupt_allows_cleanup_before_timeout_reply() { let reg = fresh_registry(); let host = PluginHost::new(Arc::clone(®)).unwrap(); - let cleanup_secs = CANCEL_INTERRUPT_GRACE.saturating_mul(2).as_secs_f64(); + let cleanup_secs = CANCEL_INTERRUPT_GRACE.as_secs_f64() / 4.0; let src = format!( r#"local cleanup_finished = false n00n.api.register_tool({{ From 699c9ebde72f5bf6bbd0c79d90faabc312dfa9fb Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 10 Aug 2026 10:04:38 -0400 Subject: [PATCH 20/24] fix(orchestration): address lifecycle review findings --- n00n-ui/src/app/session.rs | 4 +- n00n-ui/src/app/tests.rs | 7 +- n00n-ui/src/event_loop.rs | 135 +++++++++++++++++++--------- n00n-ui/src/session_lineage.rs | 65 +++++++++++++- n00n-ui/src/storage_writer.rs | 54 +++++++++-- plugins/lib/n00n/subagent.lua | 4 + site/docs/content/lua-api/_index.md | 5 ++ 7 files changed, 221 insertions(+), 53 deletions(-) diff --git a/n00n-ui/src/app/session.rs b/n00n-ui/src/app/session.rs index 6217ccfda..8ca323786 100644 --- a/n00n-ui/src/app/session.rs +++ b/n00n-ui/src/app/session.rs @@ -12,7 +12,8 @@ use n00n_providers::{Model, TokenUsage}; use n00n_storage::id::{SessionRef, n00nId}; use n00n_storage::sessions::{ StoredDelivery, StoredDirectTool, StoredImageMediaType, StoredImageSource, StoredMcpPrompt, - StoredMode, StoredQueuedMessage, StoredSessionStateSnapshot, StoredSubagent, StoredThinking, + StoredMode, StoredQueuedMessage, StoredSessionLifecycle, StoredSessionStateSnapshot, + StoredSubagent, StoredThinking, }; use crate::AppSession; @@ -57,6 +58,7 @@ pub(crate) fn session_has_content(session: &AppSession) -> bool { || !session.meta.queued_submissions.is_empty() || !session.meta.queued_direct_tools.is_empty() || session.meta.direct_output.is_some() + || session.meta.lifecycle == StoredSessionLifecycle::Cancelled || session.meta.mode != Some(n00n_storage::sessions::StoredMode::Build) || session.meta.plan_path.is_some() || session.meta.plan_written diff --git a/n00n-ui/src/app/tests.rs b/n00n-ui/src/app/tests.rs index 8215374a1..497a29755 100644 --- a/n00n-ui/src/app/tests.rs +++ b/n00n-ui/src/app/tests.rs @@ -19,7 +19,8 @@ use n00n_lua::{HintReader, KeymapReader, LuaCommandReader, PluginHost}; use n00n_providers::{ContentBlock, Effort, Role, TokenUsage}; use n00n_storage::id::SessionRef; use n00n_storage::sessions::{ - StoredMode, StoredSessionStateSnapshot, StoredStateScope, StoredThinking, TranscriptEntry, + StoredMode, StoredSessionLifecycle, StoredSessionStateSnapshot, StoredStateScope, + StoredThinking, TranscriptEntry, }; use ratatui::{Terminal, backend::TestBackend, layout::Rect}; use ratatui_image::picker::Picker; @@ -2465,6 +2466,10 @@ fn session_has_content_covers_each_branch() { assert!(session_has_content(&session)); session.meta.queued_messages.clear(); + session.meta.lifecycle = StoredSessionLifecycle::Cancelled; + assert!(session_has_content(&session)); + session.meta.lifecycle = StoredSessionLifecycle::Idle; + session.meta.mode = Some(StoredMode::Plan); assert!(session_has_content(&session)); session.meta.mode = Some(StoredMode::Build); diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index 0e5d898d7..3a38ca5d6 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -23,7 +23,8 @@ use crossterm::event::{ use n00n_agent::command::CustomCommand; use n00n_agent::permissions::PermissionManager; use n00n_agent::{ - AgentConfig, CancelToken, McpCommand, McpConfigErrors, McpHandle, mcp, tools::SessionIdentity, + AgentConfig, CancelToken, McpCommand, McpConfigErrors, McpHandle, mcp, + tools::{SessionIdentity, truncate_output}, }; use n00n_config::UiConfig; use n00n_lua::{ @@ -71,6 +72,7 @@ const AGENT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(3); const STORAGE_WRITER_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); const STORAGE_WRITER_REFS_ERR: &str = "storage writer has outstanding references, skipping graceful shutdown"; +const DIRECT_OUTPUT_MAX_BYTES: usize = 1024 * 1024; const DELETE_FOCUSED_ERR: &str = "cannot delete the focused session"; const NOT_LIVE_ERR: &str = "session not live"; const TEAM_TOOL_NAME: &str = "team"; @@ -168,12 +170,43 @@ fn cancel_stored_session(session: &mut AppSession) -> bool { || !session.meta.queued_messages.is_empty() || !session.meta.queued_submissions.is_empty() || !session.meta.queued_direct_tools.is_empty(); + if !had_work { + return false; + } session.meta.lifecycle = StoredSessionLifecycle::Cancelled; session.meta.queued_messages.clear(); session.meta.queued_submissions.clear(); session.meta.queued_direct_tools.clear(); session.updated_at = n00n_storage::now_epoch(); - had_work + true +} + +fn bounded_direct_output(text: &str, config: &AgentConfig) -> String { + truncate_output( + text, + config.max_output_lines, + config.max_output_bytes.min(DIRECT_OUTPUT_MAX_BYTES), + ) +} + +fn delete_sessions_sequentially( + writer: &Arc, + mut targets: Vec, + reply_tx: flume::Sender, +) { + let Some(target) = targets.pop() else { + let _ = reply_tx.send(Ok(json!(true))); + return; + }; + let next_writer = Arc::clone(writer); + writer.delete(target, move |result| match result { + Ok(()) | Err(SessionError::Storage(StorageError::NotFound(_))) => { + delete_sessions_sequentially(&next_writer, targets, reply_tx); + } + Err(error) => { + let _ = reply_tx.send(Err(error.to_string())); + } + }); } fn resolved_root( @@ -993,8 +1026,9 @@ impl<'t> EventLoop<'t> { if self.sessions[idx].direct_bootstrap_active { match &envelope.event { n00n_agent::AgentEvent::ToolDone(done) => { - self.sessions[idx].app.state.session.meta.direct_output = - Some(done.output.as_text()); + self.sessions[idx].app.state.session.meta.direct_output = Some( + bounded_direct_output(&done.output.as_text(), &self.ctx.config), + ); self.sessions[idx] .app .state @@ -1011,7 +1045,8 @@ impl<'t> EventLoop<'t> { .direct_output .is_none() => { - self.sessions[idx].app.state.session.meta.direct_output = Some(message.clone()); + self.sessions[idx].app.state.session.meta.direct_output = + Some(bounded_direct_output(message, &self.ctx.config)); self.sessions[idx] .app .state @@ -1093,7 +1128,7 @@ impl<'t> EventLoop<'t> { .latest_snapshot(root_id) .map_err(|error| error.to_string())? { - root + Arc::unwrap_or_clone(root) } else { AppSession::load(root_id, &self.ctx.storage).map_err(|error| error.to_string())? }; @@ -1258,7 +1293,7 @@ impl<'t> EventLoop<'t> { return; } }; - let mut targets = match self.lineage.descendants_of(id) { + let mut targets = match self.lineage.descendants_for_delete(id) { Ok(targets) => targets, Err(LineageError::UnknownSession(_)) => Vec::new(), Err(error) => { @@ -1284,27 +1319,8 @@ impl<'t> EventLoop<'t> { rt.handles.cancel(); } self.lineage.remove_sessions(&targets); - let (done_tx, done_rx) = flume::unbounded(); - let count = targets.len(); - for target in targets { - let done_tx = done_tx.clone(); - self.ctx.storage_writer.delete(target, move |result| { - let _ = done_tx.send(result); - }); - } - smol::spawn(async move { - let mut failure = None; - for _ in 0..count { - match done_rx.recv_async().await { - Ok(Ok(()) | Err(SessionError::Storage(StorageError::NotFound(_)))) => {} - Ok(Err(error)) => failure = Some(error.to_string()), - Err(error) => failure = Some(error.to_string()), - } - } - let reply = failure.map_or_else(|| Ok(json!(true)), Err); - let _ = reply_tx.send(reply); - }) - .detach(); + targets.reverse(); + delete_sessions_sequentially(&self.ctx.storage_writer, targets, reply_tx); } SessionRequest::Live => { let list: Vec<_> = self @@ -1443,8 +1459,8 @@ impl<'t> EventLoop<'t> { } else if let Some(prompt) = prompt { self.lineage .set_execution_active(id, false) - .map_err(|error| error.to_string())?; - self.submit_text(idx, prompt, false, false) + .map_err(|error| error.to_string()) + .and_then(|_| self.submit_text(idx, prompt, false, false)) } else { self.sessions[idx].app.state.session.meta.lifecycle = StoredSessionLifecycle::Idle; @@ -1548,17 +1564,19 @@ impl<'t> EventLoop<'t> { .latest_snapshot(session_id) .map_err(|error| error.to_string())? { - Some(session) => session, + Some(session) => Arc::unwrap_or_clone(session), None => AppSession::load(session_id, &self.ctx.storage) .map_err(|error| error.to_string())?, }; - cancelled |= cancel_stored_session(&mut session); - self.ctx.storage_writer.send(Box::new(session)); - warn_lineage_cleanup( - self.lineage.set_execution_active(session_id, false), - session_id, - "clear cancelled activity", - ); + if cancel_stored_session(&mut session) { + cancelled = true; + self.ctx.storage_writer.send(Box::new(session)); + warn_lineage_cleanup( + self.lineage.set_execution_active(session_id, false), + session_id, + "clear cancelled activity", + ); + } continue; }; if SessionStatus::of(&self.sessions[idx].app) != SessionStatus::Idle @@ -1574,11 +1592,11 @@ impl<'t> EventLoop<'t> { { let actions = self.sessions[idx].app.cancel_current_run(); self.dispatch(idx, actions); + self.sessions[idx].app.state.session.meta.lifecycle = + StoredSessionLifecycle::Cancelled; + self.sessions[idx].app.save_session(); cancelled = true; } - self.sessions[idx].app.state.session.meta.lifecycle = - StoredSessionLifecycle::Cancelled; - self.sessions[idx].app.save_session(); warn_lineage_cleanup( self.lineage.set_execution_active(session_id, false), session_id, @@ -2276,11 +2294,12 @@ fn scroll_delta(kind: MouseEventKind, lines: u32) -> i32 { #[cfg(test)] mod tests { use super::{ - DRAIN_BUDGET, DrainScheduler, TEAM_TOOL_NAME, cancel_stored_session, - draw_then_post_terminal, paused_team_run, should_save_periodically, - take_painted_submissions, + DIRECT_OUTPUT_MAX_BYTES, DRAIN_BUDGET, DrainScheduler, TEAM_TOOL_NAME, + bounded_direct_output, cancel_stored_session, draw_then_post_terminal, paused_team_run, + should_save_periodically, take_painted_submissions, }; use crate::{AppSession, components::Status}; + use n00n_agent::AgentConfig; use n00n_providers::{ContentBlock, Message, Role}; use n00n_storage::{ id::n00nId, @@ -2344,6 +2363,31 @@ mod tests { assert!(paused_team_run(&[tool_call, tool_result]).is_none()); } + #[test] + fn bounded_direct_output_respects_session_record_limits() { + let config = AgentConfig { + max_output_lines: 2, + max_output_bytes: 24, + ..AgentConfig::default() + }; + + let output = bounded_direct_output("αβγδεζηθ\nsecond\nthird", &config); + + assert!(output.len() <= config.max_output_bytes); + assert!(output.lines().count() <= config.max_output_lines); + assert!(output.contains("[truncated]")); + assert!(std::str::from_utf8(output.as_bytes()).is_ok()); + + let unbounded_config = AgentConfig { + max_output_lines: usize::MAX, + max_output_bytes: usize::MAX, + ..AgentConfig::default() + }; + let capped = + bounded_direct_output(&"x".repeat(DIRECT_OUTPUT_MAX_BYTES + 1), &unbounded_config); + assert!(capped.len() <= DIRECT_OUTPUT_MAX_BYTES); + } + #[test] fn cancel_stored_session_clears_all_persisted_work() { let mut session = AppSession::new("model", "/project"); @@ -2374,6 +2418,9 @@ mod tests { let mut inactive = AppSession::new("model", "/project"); inactive.meta.lifecycle = StoredSessionLifecycle::Succeeded; + assert!(!cancel_stored_session(&mut inactive)); + assert_eq!(inactive.meta.lifecycle, StoredSessionLifecycle::Succeeded); + inactive.meta.queued_messages = vec!["pending".into()]; assert!(cancel_stored_session(&mut inactive)); assert!(inactive.meta.queued_messages.is_empty()); diff --git a/n00n-ui/src/session_lineage.rs b/n00n-ui/src/session_lineage.rs index 896fad386..84f11575b 100644 --- a/n00n-ui/src/session_lineage.rs +++ b/n00n-ui/src/session_lineage.rs @@ -2,6 +2,7 @@ use std::collections::{HashMap, HashSet}; use n00n_storage::id::n00nId; use thiserror::Error; +use tracing::warn; #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[allow(clippy::struct_field_names)] @@ -197,7 +198,13 @@ impl SessionLineageGuard { ); if let Err(error) = self.rebuild_topology() { self.sessions.remove(&session.id); - self.rebuild_topology()?; + if let Err(rollback_error) = self.rebuild_topology() { + warn!( + session_id = %session.id, + error = %rollback_error, + "failed to rebuild session lineage topology after activation rollback" + ); + } return Err(error); } Ok(()) @@ -372,7 +379,13 @@ impl SessionLineageGuard { ); if let Err(error) = self.rebuild_topology() { self.sessions.remove(&child_id); - self.rebuild_topology()?; + if let Err(rollback_error) = self.rebuild_topology() { + warn!( + session_id = %child_id, + error = %rollback_error, + "failed to rebuild session lineage topology after reservation rollback" + ); + } return Err(error); } Ok(()) @@ -401,6 +414,38 @@ impl SessionLineageGuard { } pub(crate) fn descendants_of(&self, parent: n00nId) -> Result, LineageError> { + self.descendants(parent) + } + + pub(crate) fn descendants_for_delete( + &self, + parent: n00nId, + ) -> Result, LineageError> { + if !self.sessions.contains_key(&parent) { + return Err(LineageError::UnknownSession(parent)); + } + let mut pending = self + .children + .get(&parent) + .into_iter() + .flat_map(|children| children.iter().copied()) + .map(|id| (id, false)) + .collect::>(); + let mut descendants = Vec::new(); + while let Some((id, visited)) = pending.pop() { + if visited { + descendants.push(id); + continue; + } + pending.push((id, true)); + if let Some(children) = self.children.get(&id) { + pending.extend(children.iter().copied().map(|child| (child, false))); + } + } + Ok(descendants) + } + + fn descendants(&self, parent: n00nId) -> Result, LineageError> { if !self.sessions.contains_key(&parent) { return Err(LineageError::UnknownSession(parent)); } @@ -919,5 +964,21 @@ mod tests { guard.descendants_of(root).expect("descendants"), vec![sibling] ); + let delete_descendants = guard + .descendants_for_delete(root) + .expect("delete descendants"); + assert_eq!( + delete_descendants.iter().copied().collect::>(), + HashSet::from([child, grandchild, sibling]) + ); + let grandchild_index = delete_descendants + .iter() + .position(|id| *id == grandchild) + .expect("grandchild position"); + let child_index = delete_descendants + .iter() + .position(|id| *id == child) + .expect("child position"); + assert!(grandchild_index < child_index); } } diff --git a/n00n-ui/src/storage_writer.rs b/n00n-ui/src/storage_writer.rs index 814ab5789..f73ef79fb 100644 --- a/n00n-ui/src/storage_writer.rs +++ b/n00n-ui/src/storage_writer.rs @@ -22,6 +22,7 @@ use tracing::warn; use crate::AppSession; const RETRY_DELAY: Duration = Duration::from_secs(1); +const LATEST_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(2); const MAX_RETRY_ATTEMPTS: u32 = 5; #[derive(Clone)] @@ -116,7 +117,7 @@ enum Op { }, Latest { id: n00nId, - done: flume::Sender>, + done: flume::Sender>>, }, #[cfg(test)] Pause { @@ -206,7 +207,7 @@ impl StorageWriter { let session = state .latest_snapshots .get(&id) - .map(|snapshot| (*snapshot.session).clone()); + .map(|snapshot| Arc::clone(&snapshot.session)); let _ = done.send(session); } #[cfg(test)] @@ -254,12 +255,27 @@ impl StorageWriter { self.enqueue_reserved_snapshot(generation, session); } - pub(crate) fn latest_snapshot(&self, id: n00nId) -> Result, SessionError> { + pub(crate) fn latest_snapshot( + &self, + id: n00nId, + ) -> Result>, SessionError> { + self.latest_snapshot_with_timeout(id, LATEST_SNAPSHOT_TIMEOUT) + } + + fn latest_snapshot_with_timeout( + &self, + id: n00nId, + timeout: Duration, + ) -> Result>, SessionError> { let (done_tx, done_rx) = flume::bounded(1); self.ops .send(Op::Latest { id, done: done_tx }) .map_err(|_| writer_gone())?; - done_rx.recv().map_err(|_| writer_gone()) + match done_rx.recv_timeout(timeout) { + Ok(session) => Ok(session), + Err(flume::RecvTimeoutError::Timeout) => Err(latest_snapshot_timeout(timeout)), + Err(flume::RecvTimeoutError::Disconnected) => Err(writer_gone()), + } } /// Persists this snapshot before invoking `done` on the writer thread. @@ -670,6 +686,14 @@ fn writer_gone() -> SessionError { StorageError::Io(io::Error::other("storage writer unavailable")).into() } +fn latest_snapshot_timeout(timeout: Duration) -> SessionError { + StorageError::Io(io::Error::new( + io::ErrorKind::TimedOut, + format!("storage writer did not return the latest snapshot within {timeout:?}"), + )) + .into() +} + fn unpersisted_snapshot() -> SessionError { StorageError::Io(io::Error::other( "newer session snapshot remains unpersisted", @@ -917,13 +941,33 @@ mod tests { query.join().unwrap() }); let latest = match latest { - Some(session) => session, + Some(session) => Arc::unwrap_or_clone(session), None => AppSession::load(id, &dir).unwrap(), }; assert_eq!(latest.title, "latest"); writer.shutdown(DRAIN_TIMEOUT).unwrap(); } + #[test] + fn latest_snapshot_times_out_behind_blocked_writer() { + let (_tmp, dir) = state_dir(); + let writer = StorageWriter::new(dir).unwrap(); + let release = pause_writer(&writer); + let id = AppSession::new("test-model", "/tmp/latest-timeout").id; + + let error = writer + .latest_snapshot_with_timeout(id, BLOCKED_TIMEOUT) + .unwrap_err(); + + assert!(matches!( + error, + SessionError::Storage(StorageError::Io(ref io_error)) + if io_error.kind() == io::ErrorKind::TimedOut + )); + release.send(()).unwrap(); + writer.shutdown(DRAIN_TIMEOUT).unwrap(); + } + #[test] fn blocked_writer_coalesces_same_session_snapshots_and_persists_latest() { let (_tmp, dir) = state_dir(); diff --git a/plugins/lib/n00n/subagent.lua b/plugins/lib/n00n/subagent.lua index 4bf86e102..f7a25f927 100644 --- a/plugins/lib/n00n/subagent.lua +++ b/plugins/lib/n00n/subagent.lua @@ -10,6 +10,10 @@ local structured_output = require("n00n.structured_output") local ORCHESTRATION_TOOLS = { "task", "team", "workflow", "agent_control", "batch" } +-- Return a fresh table containing the orchestration tool names. +-- Use it as a denylist when child agents must not launch more orchestration. +-- Example: local excluded = subagent.orchestration_tools() +-- @return string[] function M.orchestration_tools() local copy = {} for index, name in ipairs(ORCHESTRATION_TOOLS) do diff --git a/site/docs/content/lua-api/_index.md b/site/docs/content/lua-api/_index.md index ed28cc3fa..f44dc1d1e 100644 --- a/site/docs/content/lua-api/_index.md +++ b/site/docs/content/lua-api/_index.md @@ -6005,6 +6005,11 @@ function M.make_local_tool(schema, on_submit) -- Subagent launch helper module. -- Provides a unified interface for launching subagents with model resolution, -- system prompts, tool setup, and optional structured output validation. + +-- Return a fresh table containing the orchestration tool names. +-- Use it as a denylist when child agents must not launch more orchestration. +-- Example: local excluded = subagent.orchestration_tools() +-- @return string[] function M.orchestration_tools() -- Launch a subagent with the given options. From b13fd78bf8b0fe419ef783ef9bd58a7db3de203b Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 10 Aug 2026 11:22:48 -0400 Subject: [PATCH 21/24] fix(orchestration): persist resumable team state --- n00n-storage/src/sessions.rs | 15 +++ n00n-ui/src/event_loop.rs | 175 +++++++++++++++++++++++++++-------- 2 files changed, 152 insertions(+), 38 deletions(-) diff --git a/n00n-storage/src/sessions.rs b/n00n-storage/src/sessions.rs index 0b12c2cbd..666773d68 100644 --- a/n00n-storage/src/sessions.rs +++ b/n00n-storage/src/sessions.rs @@ -938,6 +938,8 @@ pub struct SessionMeta { pub direct_output: Option, #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub direct_output_is_error: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub direct_paused_team: Option, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub subagents: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -4187,6 +4189,11 @@ mod tests { session.meta.queued_messages = vec!["queued".into()]; session.meta.direct_output = Some("bootstrap output".into()); session.meta.direct_output_is_error = true; + session.meta.direct_paused_team = Some(serde_json::json!({ + "paused": true, + "run_id": "run-1", + "mode": "swarm", + })); session.title = "updated title".into(); session.updated_at = now_epoch() + 1; log.append(&session).unwrap(); @@ -4199,6 +4206,14 @@ mod tests { Some("bootstrap output") ); assert!(loaded.meta.direct_output_is_error); + assert_eq!( + loaded.meta.direct_paused_team, + Some(serde_json::json!({ + "paused": true, + "run_id": "run-1", + "mode": "swarm", + })) + ); assert_eq!(loaded.title, "updated title"); } diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index 3a38ca5d6..32ee6c4fe 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -76,6 +76,7 @@ const DIRECT_OUTPUT_MAX_BYTES: usize = 1024 * 1024; const DELETE_FOCUSED_ERR: &str = "cannot delete the focused session"; const NOT_LIVE_ERR: &str = "session not live"; const TEAM_TOOL_NAME: &str = "team"; +const PAUSED_TEAM_RUN_ID_MAX_BYTES: usize = 256; /// Tabs carry their in-memory sessions so `/reload` reopens them without a /// disk round-trip; `session_has_content` tells which ones were saved. @@ -177,6 +178,7 @@ fn cancel_stored_session(session: &mut AppSession) -> bool { session.meta.queued_messages.clear(); session.meta.queued_submissions.clear(); session.meta.queued_direct_tools.clear(); + session.meta.direct_paused_team = None; session.updated_at = n00n_storage::now_epoch(); true } @@ -292,6 +294,50 @@ fn warn_lineage_cleanup( } } +fn validated_paused_team_payload(payload: &Value) -> Option { + if payload.get("paused").and_then(Value::as_bool) != Some(true) { + return None; + } + let run_id = payload.get("run_id")?.as_str()?; + if run_id.is_empty() || run_id.len() > PAUSED_TEAM_RUN_ID_MAX_BYTES { + return None; + } + let mode = match payload.get("mode") { + Some(mode) => Some(mode.as_str()?), + None => None, + }; + if mode.is_some_and(|mode| !matches!(mode, "supervised" | "autonomous" | "swarm")) { + return None; + } + let mut validated = json!({ "paused": true, "run_id": run_id }); + if let Some(mode) = mode { + validated["mode"] = Value::String(mode.to_owned()); + } + Some(validated) +} + +fn paused_team_payload(content: &str) -> Option { + if !content.trim_start().starts_with('{') { + return None; + } + let payload: Value = match serde_json::from_str(content) { + Ok(payload) => payload, + Err(error) => { + warn!(%error, "invalid paused team result; ignoring"); + return None; + } + }; + validated_paused_team_payload(&payload) +} + +fn direct_paused_team_payload(tool: &str, content: &str) -> Option { + if tool == TEAM_TOOL_NAME { + paused_team_payload(content) + } else { + None + } +} + fn paused_team_run(history: &[Message]) -> Option { let (user_index, last_user) = history .iter() @@ -317,22 +363,7 @@ fn paused_team_run(history: &[Message]) -> Option { continue; } - if !content.trim_start().starts_with('{') { - continue; - } - let payload: Value = match serde_json::from_str(content) { - Ok(payload) => payload, - Err(error) => { - warn!(%tool_use_id, %error, "invalid paused team result; ignoring"); - continue; - } - }; - let paused = payload.get("paused").and_then(Value::as_bool) == Some(true); - let has_run_id = payload - .get("run_id") - .and_then(Value::as_str) - .is_some_and(|run_id| !run_id.is_empty()); - if paused && has_run_id { + if let Some(payload) = paused_team_payload(content) { return Some(payload); } } @@ -381,7 +412,12 @@ struct SpawnCtx { } impl SpawnCtx { - fn spawn_runtime(&self, session: AppSession) -> Result { + fn spawn_runtime(&self, mut session: AppSession) -> Result { + session.meta.direct_paused_team = session + .meta + .direct_paused_team + .as_ref() + .and_then(validated_paused_team_payload); let resumed = crate::app::session_has_content(&session); let direct_bootstrap_active = !session.meta.queued_direct_tools.is_empty(); let identity = session_identity(&session) @@ -1026,15 +1062,11 @@ impl<'t> EventLoop<'t> { if self.sessions[idx].direct_bootstrap_active { match &envelope.event { n00n_agent::AgentEvent::ToolDone(done) => { - self.sessions[idx].app.state.session.meta.direct_output = Some( - bounded_direct_output(&done.output.as_text(), &self.ctx.config), - ); - self.sessions[idx] - .app - .state - .session - .meta - .direct_output_is_error = done.is_error; + let output = done.output.as_text(); + let meta = &mut self.sessions[idx].app.state.session.meta; + meta.direct_paused_team = direct_paused_team_payload(&done.tool, &output); + meta.direct_output = Some(bounded_direct_output(&output, &self.ctx.config)); + meta.direct_output_is_error = done.is_error; } n00n_agent::AgentEvent::Error { message } if self.sessions[idx] @@ -1358,7 +1390,15 @@ impl<'t> EventLoop<'t> { .is_none() .then_some(rt.app.state.session.meta.direct_output_is_error) .filter(|_| direct_output.is_some()); - let paused_team = paused_team_run(&history); + let paused_team = paused_team_run(&history).or_else(|| { + rt.app + .state + .session + .meta + .direct_paused_team + .as_ref() + .and_then(validated_paused_team_payload) + }); Ok(json!({ "id": rt.id(), "title": rt.app.state.session.title, @@ -1517,8 +1557,9 @@ impl<'t> EventLoop<'t> { .map_err(|error| error.to_string())?; match self.submit_text(idx, text, steer, control) { Ok(state) => { - self.sessions[idx].app.state.session.meta.lifecycle = - StoredSessionLifecycle::Running; + let meta = &mut self.sessions[idx].app.state.session.meta; + meta.lifecycle = StoredSessionLifecycle::Running; + meta.direct_paused_team = None; Ok(state) } Err(error) => { @@ -1592,8 +1633,9 @@ impl<'t> EventLoop<'t> { { let actions = self.sessions[idx].app.cancel_current_run(); self.dispatch(idx, actions); - self.sessions[idx].app.state.session.meta.lifecycle = - StoredSessionLifecycle::Cancelled; + let meta = &mut self.sessions[idx].app.state.session.meta; + meta.lifecycle = StoredSessionLifecycle::Cancelled; + meta.direct_paused_team = None; self.sessions[idx].app.save_session(); cancelled = true; } @@ -1706,8 +1748,16 @@ impl<'t> EventLoop<'t> { self.set_focus(i); return Ok(()); } - let session = AppSession::load(id, &self.ctx.storage) - .map_err(|e| format!("Failed to load session: {e}"))?; + let session = match self + .ctx + .storage_writer + .latest_snapshot(id) + .map_err(|error| format!("Failed to load pending session state: {error}"))? + { + Some(session) => Arc::unwrap_or_clone(session), + None => AppSession::load(id, &self.ctx.storage) + .map_err(|error| format!("Failed to load session: {error}"))?, + }; let restore_execution = has_restorable_work(&session); let mut live = live_session(&session).map_err(|error| error.to_string())?; live.execution_active = false; @@ -1898,6 +1948,7 @@ impl<'t> EventLoop<'t> { rt.app .queue .remove_submission(completion.dispatch.submission_id); + rt.app.save_session(); if completion.execution_started && rt.app.queue.is_empty() { warn_lineage_cleanup( self.lineage @@ -1909,12 +1960,15 @@ impl<'t> EventLoop<'t> { return; } let submission_id = completion.dispatch.submission_id; - if !rt + if rt .app .queue .mark_submission_ready(submission_id, completion.dispatch.input) { + rt.app.state.session.meta.direct_paused_team = None; + } else { rt.app.queue.remove_submission(submission_id); + rt.app.save_session(); if completion.execution_started && rt.app.queue.is_empty() { warn_lineage_cleanup( self.lineage @@ -1952,7 +2006,8 @@ impl<'t> EventLoop<'t> { return; } let session_id = rt.app.state.session.id; - let snapshot = rt.app.session_snapshot(); + let mut snapshot = rt.app.session_snapshot(); + snapshot.meta.direct_paused_team = None; let completion_tx = self.submission_persist_tx.clone(); self.ctx .storage_writer @@ -1983,6 +2038,7 @@ impl<'t> EventLoop<'t> { .meta .queued_direct_tools .clear(); + self.sessions[idx].app.state.session.meta.direct_paused_team = None; warn_lineage_cleanup( self.lineage.set_execution_active(id, false), id, @@ -2294,9 +2350,10 @@ fn scroll_delta(kind: MouseEventKind, lines: u32) -> i32 { #[cfg(test)] mod tests { use super::{ - DIRECT_OUTPUT_MAX_BYTES, DRAIN_BUDGET, DrainScheduler, TEAM_TOOL_NAME, - bounded_direct_output, cancel_stored_session, draw_then_post_terminal, paused_team_run, - should_save_periodically, take_painted_submissions, + DIRECT_OUTPUT_MAX_BYTES, DRAIN_BUDGET, DrainScheduler, PAUSED_TEAM_RUN_ID_MAX_BYTES, + TEAM_TOOL_NAME, bounded_direct_output, cancel_stored_session, direct_paused_team_payload, + draw_then_post_terminal, paused_team_payload, paused_team_run, should_save_periodically, + take_painted_submissions, validated_paused_team_payload, }; use crate::{AppSession, components::Status}; use n00n_agent::AgentConfig; @@ -2363,6 +2420,43 @@ mod tests { assert!(paused_team_run(&[tool_call, tool_result]).is_none()); } + #[test] + fn paused_team_payload_keeps_only_valid_resume_fields() { + let paused = paused_team_payload( + r#"{"paused":true,"run_id":"run-1","mode":"swarm","output":"large"}"#, + ) + .expect("paused team payload"); + assert_eq!( + paused, + serde_json::json!({"paused": true, "run_id": "run-1", "mode": "swarm"}) + ); + assert!(paused_team_payload(r#"{"paused":false,"run_id":"run-1"}"#).is_none()); + assert!(paused_team_payload(r#"{"paused":true,"run_id":""}"#).is_none()); + assert!( + paused_team_payload(r#"{"paused":true,"run_id":"run-1","mode":"invalid"}"#).is_none() + ); + let oversized = serde_json::json!({ + "paused": true, + "run_id": "x".repeat(PAUSED_TEAM_RUN_ID_MAX_BYTES + 1), + }); + assert!(validated_paused_team_payload(&oversized).is_none()); + assert!(validated_paused_team_payload(&serde_json::json!({"run_id": "run-1"})).is_none()); + } + + #[test] + fn direct_paused_team_payload_requires_team_tool_event() { + let output = format!( + r#"{{"paused":true,"run_id":"run-1","output":"{}"}}"#, + "x".repeat(DIRECT_OUTPUT_MAX_BYTES) + ); + + assert_eq!( + direct_paused_team_payload(TEAM_TOOL_NAME, &output), + Some(serde_json::json!({"paused": true, "run_id": "run-1"})) + ); + assert!(direct_paused_team_payload("task", &output).is_none()); + } + #[test] fn bounded_direct_output_respects_session_record_limits() { let config = AgentConfig { @@ -2409,12 +2503,17 @@ mod tests { tool: "task".into(), input: serde_json::json!({}), }]; + session.meta.direct_paused_team = Some(serde_json::json!({ + "paused": true, + "run_id": "run-1", + })); assert!(cancel_stored_session(&mut session)); assert_eq!(session.meta.lifecycle, StoredSessionLifecycle::Cancelled); assert!(session.meta.queued_messages.is_empty()); assert!(session.meta.queued_submissions.is_empty()); assert!(session.meta.queued_direct_tools.is_empty()); + assert!(session.meta.direct_paused_team.is_none()); let mut inactive = AppSession::new("model", "/project"); inactive.meta.lifecycle = StoredSessionLifecycle::Succeeded; From 81d3af70f31978ce920237e966d30b5d91b0e461 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Mon, 10 Aug 2026 11:36:17 -0400 Subject: [PATCH 22/24] fix(orchestration): close final review gaps --- n00n-lua/src/api/session.rs | 5 +++-- n00n-ui/src/event_loop.rs | 1 - plugins/lib/n00n/subagent.lua | 2 +- site/docs/content/lua-api/_index.md | 7 +++++-- 4 files changed, 9 insertions(+), 6 deletions(-) diff --git a/n00n-lua/src/api/session.rs b/n00n-lua/src/api/session.rs index 41cf7c627..3f746bd3c 100644 --- a/n00n-lua/src/api/session.rs +++ b/n00n-lua/src/api/session.rs @@ -118,8 +118,9 @@ async fn delete( /// /// @param opts table? Optional fields: prompt (string) first user message /// to submit right away; focus (boolean) switch the UI to the new session; -/// parent_id (string?) session that spawned this session; tool (string), -/// input (table), and title (string?) for a direct host-executed bootstrap. +/// parent_id (string?) session that spawned this session; tool (string) for a +/// direct host-executed bootstrap. Tool cannot be combined with prompt. Input +/// (table) and title (string?) require tool. /// @return (string|nil, string|nil) New session id, or nil and an error. /// @example /// local id, err = n00n.session.new({ prompt = "fix the tests", focus = true }) diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index 32ee6c4fe..3ee7d3f53 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -482,7 +482,6 @@ impl SpawnCtx { picker: Arc::clone(&self.picker), }); app.lua_event_handle.clone_from(&self.lua_event_handle); - app.hydrate_plugin_state(); handles.apply_to_app(&mut app); if resumed { restore_session(&mut app, &handles); diff --git a/plugins/lib/n00n/subagent.lua b/plugins/lib/n00n/subagent.lua index f7a25f927..6c7b2b461 100644 --- a/plugins/lib/n00n/subagent.lua +++ b/plugins/lib/n00n/subagent.lua @@ -25,7 +25,7 @@ end local function excluded_tools(opts) local excluded = {} local seen = {} - if not opts.allow_orchestration then + if opts.allow_orchestration ~= true then for _, name in ipairs(ORCHESTRATION_TOOLS) do excluded[#excluded + 1] = name seen[name] = true diff --git a/site/docs/content/lua-api/_index.md b/site/docs/content/lua-api/_index.md index f44dc1d1e..e7b85c149 100644 --- a/site/docs/content/lua-api/_index.md +++ b/site/docs/content/lua-api/_index.md @@ -2806,9 +2806,12 @@ Starts a new session in the current project. to submit right away; focus (boolean) switch the UI to the new session; - - `parent_id` (`string?`) session that spawned this session; tool (string), + - `parent_id` (`string?`) session that spawned this session; tool (string) for a - input (table), and title (string?) for a direct host-executed bootstrap. + direct host-executed bootstrap. Tool cannot be combined with prompt. Input + + + (table) and title (string?) require tool. **Returns:** (`string|nil`, `string|nil`) New session id, or nil and an error. From 30e4dabd4c2b71c0ee2a8de1a09d5ccaa64ddc5b Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Wed, 12 Aug 2026 01:54:32 -0400 Subject: [PATCH 23/24] fix(ui): refresh lineage after session reset --- n00n-ui/src/agent/mod.rs | 2 + n00n-ui/src/app/session.rs | 5 +- n00n-ui/src/app/tests.rs | 15 +++-- n00n-ui/src/components/mod.rs | 6 +- n00n-ui/src/event_loop.rs | 38 ++++++++++- n00n-ui/src/session_lineage.rs | 116 +++++++++++++++++++++++++++++++++ 6 files changed, 171 insertions(+), 11 deletions(-) diff --git a/n00n-ui/src/agent/mod.rs b/n00n-ui/src/agent/mod.rs index 6bc361d5d..29ef2c977 100644 --- a/n00n-ui/src/agent/mod.rs +++ b/n00n-ui/src/agent/mod.rs @@ -135,8 +135,10 @@ impl AgentHandles { tool_output_lines: ToolOutputLines, permissions: &Arc, app: &mut App, + identity: Option, lua_handle: Option, ) { + self.identity = identity; let slot = model_slot.load(); if let Err(e) = smol::block_on(slot.provider.reload_auth()) { warn!(error = %e, "failed to reload auth, continuing with existing credentials"); diff --git a/n00n-ui/src/app/session.rs b/n00n-ui/src/app/session.rs index 8ca323786..309809a7b 100644 --- a/n00n-ui/src/app/session.rs +++ b/n00n-ui/src/app/session.rs @@ -454,7 +454,8 @@ impl App { pub(super) fn reset_session(&mut self) -> Vec { self.save_session(); - self.drop_plugin_state(self.state.session.id); + let previous_id = self.state.session.id; + self.drop_plugin_state(previous_id); self.reset_ui_chrome(); self.state.token_usage = TokenUsage::default(); self.state.context_size = 0; @@ -466,7 +467,7 @@ impl App { self.hydrate_plugin_state(); self.fire_session_autocmd("SessionReset", serde_json::json!({})); self.fire_session_focus_autocmd(); - vec![Action::NewSession] + vec![Action::NewSession { previous_id }] } pub(super) fn open_rewind_picker(&mut self) -> Vec { diff --git a/n00n-ui/src/app/tests.rs b/n00n-ui/src/app/tests.rs index 497a29755..4259535f4 100644 --- a/n00n-ui/src/app/tests.rs +++ b/n00n-ui/src/app/tests.rs @@ -861,7 +861,7 @@ fn enter_executes_new_command() { type_slash(&mut app); app.update(Msg::Key(key(KeyCode::Char('n')))); let actions = app.update(Msg::Key(key(KeyCode::Enter))); - assert!(matches!(&actions[0], Action::NewSession)); + assert!(matches!(&actions[0], Action::NewSession { .. })); assert!(!app.command_palette.is_active()); } @@ -887,8 +887,13 @@ fn reset_session_clears_plan() { app.help_modal.toggle(); let (_tx, rx) = flume::bounded::(1); app.btw_modal.open("q", rx); + let previous_id = app.state.session.id; let actions = app.reset_session(); - assert!(matches!(&actions[0], Action::NewSession)); + assert!(matches!( + &actions[0], + Action::NewSession { previous_id: id } if *id == previous_id + )); + assert_ne!(app.state.session.id, previous_id); assert_eq!(app.status, Status::Idle); assert_eq!(app.state.token_usage.input, 0); assert_eq!(app.chats[0].context_size, 0); @@ -3846,7 +3851,9 @@ fn plan_form_menu_options( assert!(matches!(app.state.plan, PlanState::Ready(_))); } assert_eq!( - actions.iter().any(|a| matches!(a, Action::NewSession)), + actions + .iter() + .any(|a| matches!(a, Action::NewSession { .. })), has_new_session ); let expected_msg = implement_msg(PlanForm::new().parallel()); @@ -3873,7 +3880,7 @@ fn clear_and_implement_defers_submission_until_new_session() { let actions = app.implement_plan(true); - assert!(matches!(&actions[..], [Action::NewSession])); + assert!(matches!(&actions[..], [Action::NewSession { .. }])); assert_ne!(app.state.session.id, old_session_id); let pending = app .pending_plan_submit diff --git a/n00n-ui/src/components/mod.rs b/n00n-ui/src/components/mod.rs index d843446fa..2faa0c755 100644 --- a/n00n-ui/src/components/mod.rs +++ b/n00n-ui/src/components/mod.rs @@ -36,7 +36,7 @@ use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; use n00n_agent::{AgentInput, PreDispatchGate}; use n00n_agent::{BufferSnapshot, ImageSource, ToolInput, ToolOutput}; use n00n_providers::{Message, ModelTier}; -use n00n_storage::sessions::TranscriptEntry; +use n00n_storage::{id::n00nId, sessions::TranscriptEntry}; use ratatui::text::{Line, Span}; pub(crate) const CHEVRON: &str = "❯ "; @@ -205,7 +205,9 @@ pub enum Action { CancelSubagent { tool_use_id: String, }, - NewSession, + NewSession { + previous_id: n00nId, + }, LoadSession(Box), ChangeModel(String), RefreshProvider { diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index 3ee7d3f53..ec34405c3 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -1909,6 +1909,7 @@ impl<'t> EventLoop<'t> { idx: usize, history: Vec, transcript: Vec>, + identity: SessionIdentity, ) { let rt = &mut self.sessions[idx]; let lua_handle = rt.app.lua_event_handle.clone(); @@ -1921,6 +1922,7 @@ impl<'t> EventLoop<'t> { self.ctx.ui_config.tool_output_lines, &permissions, &mut rt.app, + Some(identity), lua_handle, ); } @@ -2051,8 +2053,31 @@ impl<'t> EventLoop<'t> { .cmd_tx .try_send(AgentCommand::CancelSubagent { tool_use_id }); } - Action::NewSession => { - self.respawn_agent(idx, Vec::new(), Vec::new()); + Action::NewSession { previous_id } => { + let replacement = match live_session(&self.sessions[idx].app.state.session) { + Ok(replacement) => replacement, + Err(error) => { + warn!(session_id = %self.sessions[idx].id(), %error, "invalid reset session lineage"); + return; + } + }; + let identity = match session_identity(&self.sessions[idx].app.state.session) { + Ok(identity) => identity, + Err(error) => { + warn!(session_id = %replacement.id, %error, "invalid reset session identity"); + return; + } + }; + if let Err(error) = self.lineage.replace_runtime(previous_id, replacement) { + warn!( + previous_session_id = %previous_id, + replacement_session_id = %replacement.id, + %error, + "failed to replace reset session lineage" + ); + return; + } + self.respawn_agent(idx, Vec::new(), Vec::new(), identity); if let Some(pending) = self.sessions[idx].app.pending_plan_submit.take() { let actions = { let app = &mut self.sessions[idx].app; @@ -2084,7 +2109,14 @@ impl<'t> EventLoop<'t> { provider: Arc::from(new_provider), })); } - self.respawn_agent(idx, loaded.messages, loaded.transcript); + let identity = match session_identity(&self.sessions[idx].app.state.session) { + Ok(identity) => identity, + Err(error) => { + warn!(session_id = %self.sessions[idx].id(), %error, "invalid loaded session identity"); + return; + } + }; + self.respawn_agent(idx, loaded.messages, loaded.transcript, identity); *self.sessions[idx] .handles .tool_outputs diff --git a/n00n-ui/src/session_lineage.rs b/n00n-ui/src/session_lineage.rs index 84f11575b..88e63cad6 100644 --- a/n00n-ui/src/session_lineage.rs +++ b/n00n-ui/src/session_lineage.rs @@ -220,6 +220,56 @@ impl SessionLineageGuard { Ok(()) } + pub(crate) fn replace_runtime( + &mut self, + previous_id: n00nId, + replacement: LiveSession, + ) -> Result<(), LineageError> { + if self.sessions.contains_key(&replacement.id) { + return Err(LineageError::DuplicateSession(replacement.id)); + } + let previous = *self + .sessions + .get(&previous_id) + .ok_or(LineageError::UnknownSession(previous_id))?; + if !previous.runtime_present || previous.deleted { + return Err(LineageError::TargetNotLive(previous_id)); + } + + self.sessions + .get_mut(&previous_id) + .ok_or(LineageError::UnknownSession(previous_id))? + .runtime_present = false; + self.sessions + .get_mut(&previous_id) + .ok_or(LineageError::UnknownSession(previous_id))? + .execution_active = false; + self.sessions.insert( + replacement.id, + SessionNode { + root_session_id: replacement.root_session_id, + parent_id: replacement.parent_id, + runtime_present: replacement.runtime_present, + execution_active: replacement.execution_active, + deleted: false, + }, + ); + if let Err(error) = self.rebuild_topology() { + self.sessions.remove(&replacement.id); + self.sessions.insert(previous_id, previous); + if let Err(rollback_error) = self.rebuild_topology() { + warn!( + previous_session_id = %previous_id, + replacement_session_id = %replacement.id, + error = %rollback_error, + "failed to rebuild session lineage topology after runtime replacement rollback" + ); + } + return Err(error); + } + Ok(()) + } + pub(crate) fn set_execution_active( &mut self, id: n00nId, @@ -880,6 +930,72 @@ mod tests { )); } + #[test] + fn runtime_replacement_atomically_moves_liveness_to_new_root() { + let previous = id(1); + let child = id(2); + let replacement = id(3); + let mut guard = SessionLineageGuard::from_live( + [session(previous, None), session(child, Some(previous))], + limits(4, 4, 4), + ) + .expect("valid graph"); + + guard + .replace_runtime(previous, session(replacement, None)) + .expect("replace root runtime"); + + assert!(matches!( + guard.reserve_new(previous, None, false), + Err(LineageError::CallerNotLive(id)) if id == previous + )); + assert_eq!( + guard.lineage(replacement).expect("replacement lineage"), + SessionLineage { + caller: replacement, + root: replacement, + parent: None, + depth: 0, + } + ); + assert_eq!( + guard.descendants_of(previous).expect("old descendants"), + vec![child] + ); + } + + #[test] + fn failed_runtime_replacement_leaves_previous_runtime_live() { + let previous = id(1); + let replacement = id(2); + let missing_parent = id(3); + let mut guard = SessionLineageGuard::from_live([session(previous, None)], limits(4, 4, 4)) + .expect("valid root"); + let invalid_replacement = LiveSession { + id: replacement, + root_session_id: missing_parent, + parent_id: Some(missing_parent), + runtime_present: true, + execution_active: false, + }; + + assert!(matches!( + guard.replace_runtime(previous, invalid_replacement), + Err(LineageError::MissingParent { .. }) + )); + assert_eq!( + guard.lineage(previous).expect("previous lineage").root, + previous + ); + guard + .reserve_new(previous, None, false) + .expect("previous remains live"); + assert!(matches!( + guard.lineage(replacement), + Err(LineageError::CallerNotLive(id)) if id == replacement + )); + } + #[test] fn reservation_release_is_exact_and_removal_releases_only_active_capacity() { let root = id(1); From 0ee6fb87b6fb71e7ee469bf10c695ad1b9699d44 Mon Sep 17 00:00:00 2001 From: w0wl0lxd Date: Wed, 12 Aug 2026 12:05:46 -0400 Subject: [PATCH 24/24] fix(orchestration): preserve crash recovery hardening --- changelog.d/327.fixed.md | 2 +- n00n-lua/src/api/session.rs | 38 ++++++- n00n-lua/src/api/util/command.rs | 2 + n00n-lua/src/runtime.rs | 41 ++++++- .../src/providers/openai/platform.rs | 103 +++++++++++------- n00n-ui/src/app/session.rs | 16 ++- n00n-ui/src/app/tests.rs | 19 ++++ n00n-ui/src/event_loop.rs | 65 +++++++++-- n00n-ui/src/storage_writer.rs | 21 ++++ 9 files changed, 244 insertions(+), 63 deletions(-) diff --git a/changelog.d/327.fixed.md b/changelog.d/327.fixed.md index 51e0f2b75..b0acb016d 100644 --- a/changelog.d/327.fixed.md +++ b/changelog.d/327.fixed.md @@ -1 +1 @@ -Hardened agent spawn recursion and session lineage handling to prevent runaway subagent loops. \ No newline at end of file +Hardened agent session recovery, OpenAI response continuation, UI-only deletion, and lineage handling to prevent lost progress and runaway subagent loops. \ No newline at end of file diff --git a/n00n-lua/src/api/session.rs b/n00n-lua/src/api/session.rs index 3f746bd3c..c1c7bf6a9 100644 --- a/n00n-lua/src/api/session.rs +++ b/n00n-lua/src/api/session.rs @@ -7,7 +7,7 @@ use n00n_lua_macro::{lua_fn, lua_table}; use crate::api::util::command::{SessionBootstrap, SessionReply, SessionRequest, UiAction}; use crate::api::util::convert::{json_to_lua, lua_to_json}; -use crate::runtime::active_session_identity; +use crate::runtime::{active_session_identity, active_trusted_ui_control}; const NO_UI_ERR: &str = "no interactive UI attached"; @@ -111,7 +111,18 @@ async fn delete( #[ctx] tx: Option>, id: String, ) -> LuaResult { - roundtrip(lua, tx, SessionRequest::Delete { id }).await + let caller_id = active_session_identity(&lua).map(|identity| identity.session_id().clone()); + let trusted_ui_control = active_trusted_ui_control(&lua); + roundtrip( + lua, + tx, + SessionRequest::Delete { + id, + caller_id, + trusted_ui_control, + }, + ) + .await } /// Starts a new session in the current project. @@ -342,16 +353,34 @@ mod tests { }; assert_eq!(actual_caller_id.as_ref(), Some(&expected_caller_id)); reply_tx.send(Ok(json!("queued"))).unwrap(); + let Ok(UiAction::Session { + req: + SessionRequest::Delete { + id, + caller_id: actual_caller_id, + trusted_ui_control, + }, + reply_tx, + }) = rx.recv() + else { + panic!("expected delete request"); + }; + assert_eq!(id, "target"); + assert_eq!(actual_caller_id.as_ref(), Some(&expected_caller_id)); + assert!(!trusted_ui_control); + reply_tx.send(Ok(json!(true))).unwrap(); }); - let (child_id, prompt_status): (String, String) = smol::block_on( + let (child_id, prompt_status, deleted): (String, String, bool) = smol::block_on( lua.load( r#" local child, new_err = session.new({ caller_id = "spoof" }) if new_err then error(new_err) end local status, prompt_err = session.prompt("hello", { caller_id = "spoof" }) if prompt_err then error(prompt_err) end - return child, status + local deleted, delete_err = session.delete("target") + if delete_err then error(delete_err) end + return child, status, deleted "#, ) .eval_async(), @@ -360,6 +389,7 @@ mod tests { checker.join().unwrap(); assert_eq!(child_id, "child"); assert_eq!(prompt_status, "queued"); + assert!(deleted); } #[test] diff --git a/n00n-lua/src/api/util/command.rs b/n00n-lua/src/api/util/command.rs index 32bdc043f..f82430aff 100644 --- a/n00n-lua/src/api/util/command.rs +++ b/n00n-lua/src/api/util/command.rs @@ -444,6 +444,8 @@ pub enum SessionRequest { }, Delete { id: String, + caller_id: Option, + trusted_ui_control: bool, }, SetTitle { id: String, diff --git a/n00n-lua/src/runtime.rs b/n00n-lua/src/runtime.rs index 3caba47f1..53b4be2ef 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -348,6 +348,7 @@ pub(crate) struct TaskCell { /// `n00n.agent.call_tool(on_live_buf/on_annotation)`. pub(crate) live_sink: Option>, pub(crate) identity: Option, + pub(crate) trusted_ui_control: bool, /// When `Some`, `n00n.async.run` tasks queue here instead of the global /// `SpawnQueue` so restore can run them inline before snapshotting. pub(crate) inline_spawn: Option>, @@ -384,6 +385,7 @@ impl TaskCell { root_buf: None, live_sink: None, identity, + trusted_ui_control: false, inline_spawn: None, bufs_claim: Weak::new(), async_tasks: Cell::new(0), @@ -663,6 +665,12 @@ impl TaskScope { Self::new(lua, TaskCell::new(CancelToken::none(), None, None, None)) } + pub(crate) fn trusted_ui(lua: &Lua) -> Self { + let scope = Self::detached(lua); + lock_cell(&scope.handle).trusted_ui_control = true; + scope + } + pub(crate) fn handle(&self) -> &TaskHandle { &self.handle } @@ -689,12 +697,17 @@ impl TaskScope { /// /// [detached]: TaskScope::detached pub(crate) async fn run_detached(lua: &Lua, fut: F) -> F::Output { - run_callback(lua, None, fut).await + run_callback(lua, None, false, fut).await +} + +pub(crate) async fn run_trusted_ui(lua: &Lua, fut: F) -> F::Output { + run_callback(lua, None, true, fut).await } async fn run_callback( lua: &Lua, identity: Option, + trusted_ui_control: bool, fut: F, ) -> F::Output { let scope = TaskScope::new( @@ -702,6 +715,7 @@ async fn run_callback( TaskCell::new(CancelToken::none(), None, None, identity), ); let handle = Arc::clone(scope.handle()); + lock_cell(&handle).trusted_ui_control = trusted_ui_control; let pump = async { let mut event_buf = Vec::new(); loop { @@ -786,6 +800,13 @@ pub(crate) fn active_session_identity(lua: &Lua) -> Option { lock_cell(&handle).identity.clone() } +pub(crate) fn active_trusted_ui_control(lua: &Lua) -> bool { + let Some(handle) = lua.app_data_ref::() else { + return false; + }; + lock_cell(&handle).trusted_ui_control +} + pub(crate) fn with_task_jobs(lua: &Lua, f: impl FnOnce(&mut JobStore) -> R) -> R { f(&mut lock_cell(&active_task(lua)).jobs) } @@ -3066,7 +3087,7 @@ pub fn spawn( let thread = lua.create_thread(func)?; thread.into_async::<()>(args)?.await }; - if let Err(e) = run_callback(&lua, identity, run).await { + if let Err(e) = run_callback(&lua, identity, false, run).await { tracing::warn!(plugin = %plugin, command = %command, error = %e, "command handler failed"); } }) @@ -3296,7 +3317,7 @@ pub fn spawn( } Err(_) => LuaValue::Nil, }; - let scope = TaskScope::detached(&rt.lua); + let scope = TaskScope::trusted_ui(&rt.lua); if let Err(e) = scope.scope_future(func.call_async::<()>(arg)).await { tracing::warn!(error = %e, "window buffer click failed"); } @@ -3325,7 +3346,7 @@ pub fn spawn( if let Some(func) = func { let lua = rt.lua.clone(); ex.spawn(async move { - if let Err(e) = run_detached(&lua, func.call_async::<()>(())).await { + if let Err(e) = run_trusted_ui(&lua, func.call_async::<()>(())).await { tracing::warn!(keybind_id = id, error = %e, "keybind callback failed"); } }).detach(); @@ -4016,6 +4037,18 @@ mod tests { assert_eq!(interrupt_reason(&lua), None); } + #[test] + fn trusted_ui_scope_marks_only_its_callback() { + let (lua, _watchdog) = watchdog_lua(false); + assert!(!active_trusted_ui_control(&lua)); + + let scope = TaskScope::trusted_ui(&lua); + assert!(active_trusted_ui_control(&lua)); + drop(scope); + + assert!(!active_trusted_ui_control(&lua)); + } + #[test] fn expired_cleanup_window_interrupts_cancelled_task() { let (lua, _watchdog) = watchdog_lua(false); diff --git a/n00n-providers/src/providers/openai/platform.rs b/n00n-providers/src/providers/openai/platform.rs index 97d4f56a4..59c15e78e 100644 --- a/n00n-providers/src/providers/openai/platform.rs +++ b/n00n-providers/src/providers/openai/platform.rs @@ -1284,7 +1284,7 @@ impl OpenAi { ) -> CodexAttempt { if attempt.previous_response_id.is_some() && (is_missing_previous_response(&attempt) - || should_clear_response_chain(&attempt.result)) + || should_clear_response_chain(&attempt.result, response_chain_lock.is_some())) { self.clear_response_chain(session_id, response_chain_lock) .await; @@ -1347,6 +1347,7 @@ impl OpenAi { } else { None }; + let persist_response_chain = response_chain_lock.is_some(); let stream_timeout = self.compat.stream_timeout(); let connection_reusable = self .response_connection_is_reusable( @@ -1356,7 +1357,7 @@ impl OpenAi { attempt_nonce, ) .await; - if !connection_reusable { + if !connection_reusable && !persist_response_chain { debug!( chain_reset = true, chain_reset_reason = "socket_not_reusable", @@ -1412,12 +1413,13 @@ impl OpenAi { tools, previous_response_id.as_deref(), Some(&prompt_cache_key), - false, + persist_response_chain, &opts, true, ); let mut full_history_body = None; let full_history_fallback_available = previous_response_id.is_some() + && !persist_response_chain && (!opts.protect_history_replay || opts.allow_history_replay); log_responses_request( "websocket", @@ -1447,7 +1449,7 @@ impl OpenAi { tools, None, Some(&prompt_cache_key), - false, + persist_response_chain, &opts, true, ) @@ -1486,26 +1488,34 @@ impl OpenAi { .await; } warn!("OpenAI Responses WebSocket unavailable; falling back to HTTP"); - let fallback_body = full_history_body.get_or_insert_with(|| { - super::websocket::build_request_body( - model, - messages, - system, - tools, - None, - Some(&prompt_cache_key), - false, - &opts, - true, - ) - }); + let fallback_body = if persist_response_chain { + &body + } else { + full_history_body.get_or_insert_with(|| { + super::websocket::build_request_body( + model, + messages, + system, + tools, + None, + Some(&prompt_cache_key), + false, + &opts, + true, + ) + }) + }; log_responses_request( "http_sse", fallback_body, messages.len(), - messages.len(), - false, - true, + if persist_response_chain { + incremental_messages.len() + } else { + messages.len() + }, + persist_response_chain && previous_response_id.is_some(), + !persist_response_chain, ); let fallback_auth = loop { let preflight = match self.pre_send_auth(attempt_nonce).await { @@ -1585,7 +1595,9 @@ impl OpenAi { ) .await { - Ok((response_id, response)) => (response_id, response, false), + Ok((response_id, response)) => { + (response_id, response, persist_response_chain) + } Err(error) => { return self .finish_codex_attempt( @@ -1622,7 +1634,7 @@ impl OpenAi { tools_hash, &state_scope_hash, messages, - false, + persist_response_chain, response_chain_lock.as_ref(), ) .await; @@ -2320,8 +2332,8 @@ fn is_missing_previous_response(attempt: &CodexAttempt) -> bool { && normalized == format!("not found: {}", previous_response_id.to_ascii_lowercase()) } -fn should_clear_response_chain(result: &Result) -> bool { - result.is_err() +fn should_clear_response_chain(result: &Result, durable_chain: bool) -> bool { + result.is_err() && !durable_chain } fn is_definitive_responses_rejection(error: &AgentError) -> bool { @@ -2821,7 +2833,7 @@ mod tests { #[test] #[allow(clippy::too_many_lines)] - fn approved_ephemeral_preflight_failure_rebuilds_second_turn_with_full_history() { + fn durable_preflight_failure_continues_second_turn_from_persisted_response() { smol::block_on(async { let temp_dir = TempDir::new().unwrap(); let listener = smol::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -2915,8 +2927,13 @@ mod tests { .unwrap(); let storage = StateDir::from_path(temp_dir.path().to_path_buf()); provider.storage = Some(storage.clone()); - provider.response_state_storage = Some(storage); + provider.response_state_storage = Some(storage.clone()); let session_id = SessionRef::generate(); + let mut session = n00n_storage::sessions::Session::::new( + "model", "/project", + ); + session.id = session_id.id(); + session.save(&storage).unwrap(); let model = Model::from_spec("codex/gpt-5.3-codex").unwrap(); let tools = serde_json::json!([]); let (event_tx, _event_rx) = flume::unbounded(); @@ -2960,20 +2977,20 @@ mod tests { let first_body = body_rx.recv_async().await.unwrap(); let second_body = body_rx.recv_async().await.unwrap(); assert!(first_body.get("previous_response_id").is_none()); - assert_eq!(first_body["store"], false); - assert!(second_body.get("previous_response_id").is_none()); - assert_eq!(second_body["store"], false); - assert_eq!(second_body["input"].as_array().unwrap().len(), 3); + assert_eq!(first_body["store"], true); + assert_eq!(second_body["previous_response_id"], "resp_first"); + assert_eq!(second_body["store"], true); + assert_eq!(second_body["input"].as_array().unwrap().len(), 1); - let sessions_dir = temp_dir.path().join(n00n_storage::sessions::SESSIONS_DIR); - let session_prefix = session_id.id().to_string(); - assert!(std::fs::read_dir(sessions_dir).unwrap().all(|entry| { - !entry - .unwrap() - .file_name() - .to_string_lossy() - .starts_with(&session_prefix) - })); + let lock = provider + .lock_response_chain(Some(&session_id)) + .await + .unwrap() + .expect("durable response-chain lock"); + let chain = load_openai_response_chain(&storage, session_id.id(), &lock) + .unwrap() + .expect("persisted response chain"); + assert_eq!(chain.response_id, "resp_second"); }); } @@ -4540,17 +4557,19 @@ mod tests { #[test] fn successful_socket_local_continuation_keeps_response_chain() { let success: Result<(), AgentError> = Ok(()); - assert!(!should_clear_response_chain(&success)); + assert!(!should_clear_response_chain(&success, false)); let transport_error: Result<(), AgentError> = Err(std::io::Error::new(std::io::ErrorKind::ConnectionAborted, "closed").into()); - assert!(should_clear_response_chain(&transport_error)); + assert!(should_clear_response_chain(&transport_error, false)); + assert!(!should_clear_response_chain(&transport_error, true)); let api_error: Result<(), AgentError> = Err(AgentError::Api { status: 500, message: "temporary".into(), }); - assert!(should_clear_response_chain(&api_error)); + assert!(should_clear_response_chain(&api_error, false)); + assert!(!should_clear_response_chain(&api_error, true)); } #[test] diff --git a/n00n-ui/src/app/session.rs b/n00n-ui/src/app/session.rs index 309809a7b..252d531de 100644 --- a/n00n-ui/src/app/session.rs +++ b/n00n-ui/src/app/session.rs @@ -1,6 +1,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicBool; +use std::time::Duration; use crate::chat::{Chat, DONE_TEXT, RESTORE_BATCH_SIZE, history_to_display, transcript_to_display}; use crate::components::DisplayRole; @@ -11,9 +12,9 @@ use n00n_agent::{AgentInput, AgentMode, McpPromptRef}; use n00n_providers::{Model, TokenUsage}; use n00n_storage::id::{SessionRef, n00nId}; use n00n_storage::sessions::{ - StoredDelivery, StoredDirectTool, StoredImageMediaType, StoredImageSource, StoredMcpPrompt, - StoredMode, StoredQueuedMessage, StoredSessionLifecycle, StoredSessionStateSnapshot, - StoredSubagent, StoredThinking, + SessionError, StoredDelivery, StoredDirectTool, StoredImageMediaType, StoredImageSource, + StoredMcpPrompt, StoredMode, StoredQueuedMessage, StoredSessionLifecycle, + StoredSessionStateSnapshot, StoredSubagent, StoredThinking, }; use crate::AppSession; @@ -186,6 +187,15 @@ impl App { self.storage_writer.send(Box::new(snapshot)); } + pub(crate) fn checkpoint_session(&mut self, timeout: Duration) -> Result<(), SessionError> { + let snapshot = self.session_snapshot_with_plugin_state(); + if !session_has_content(&snapshot) { + return Ok(()); + } + self.storage_writer + .persist_and_wait(Box::new(snapshot), timeout) + } + pub(crate) fn session_snapshot(&mut self) -> AppSession { self.state.sync_session( self.shared_history.as_ref(), diff --git a/n00n-ui/src/app/tests.rs b/n00n-ui/src/app/tests.rs index 4259535f4..e0420ef9c 100644 --- a/n00n-ui/src/app/tests.rs +++ b/n00n-ui/src/app/tests.rs @@ -2595,6 +2595,25 @@ fn drain_writer(app: App, writer: Arc) { .unwrap(); } +#[test] +fn checkpoint_session_is_durable_before_returning() { + let (_tmp, dir, writer, mut app) = tempdir_app(); + let session_id = app.state.session.id; + app.state + .session + .messages + .push(Message::user("checkpoint".into())); + + app.checkpoint_session(WRITER_DRAIN_TIMEOUT).unwrap(); + + let loaded = AppSession::load(session_id, &dir).unwrap(); + assert_eq!( + serde_json::to_value(&loaded.messages).unwrap(), + serde_json::to_value(&app.state.session.messages).unwrap() + ); + drain_writer(app, writer); +} + #[test] fn save_session_captures_plugin_state_snapshot() { let (_tmp, dir, writer, mut app) = tempdir_app(); diff --git a/n00n-ui/src/event_loop.rs b/n00n-ui/src/event_loop.rs index 2ecc561ad..e320c0a91 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -69,10 +69,12 @@ const PERIODIC_SAVE_INTERVAL: Duration = Duration::from_secs(1); const DRAIN_BUDGET: usize = 256; const AGENT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(3); const STORAGE_WRITER_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +const TERMINAL_CHECKPOINT_TIMEOUT: Duration = Duration::from_secs(5); const STORAGE_WRITER_REFS_ERR: &str = "storage writer has outstanding references, skipping graceful shutdown"; const DIRECT_OUTPUT_MAX_BYTES: usize = 1024 * 1024; const DELETE_FOCUSED_ERR: &str = "cannot delete the focused session"; +const DELETE_UI_ONLY_ERR: &str = "session deletion is available only from trusted UI controls"; const NOT_LIVE_ERR: &str = "session not live"; const TEAM_TOOL_NAME: &str = "team"; const PAUSED_TEAM_RUN_ID_MAX_BYTES: usize = 256; @@ -141,7 +143,17 @@ fn parse_session_id(id: &str) -> Result { fn caller_session_id(caller: Option) -> Result { caller .map(|session| session.id()) - .ok_or_else(|| "authoritative caller session identity is unavailable".to_owned()) + .ok_or_else(|| "session caller identity is required".to_owned()) +} + +fn authorize_ui_delete( + caller: Option<&SessionRef>, + trusted_ui_control: bool, +) -> Result<(), String> { + if caller.is_some() || !trusted_ui_control { + return Err(DELETE_UI_ONLY_ERR.to_owned()); + } + Ok(()) } fn live_session(session: &AppSession) -> std::result::Result { @@ -1173,7 +1185,12 @@ impl<'t> EventLoop<'t> { .meta .queued_direct_tools .clear(); - self.sessions[idx].app.save_session(); + if let Err(error) = self.sessions[idx] + .app + .checkpoint_session(TERMINAL_CHECKPOINT_TIMEOUT) + { + warn!(session_id = %self.sessions[idx].id(), %error, "failed to persist terminal session checkpoint"); + } } } self.dispatch(idx, actions); @@ -1360,7 +1377,15 @@ impl<'t> EventLoop<'t> { // Deletes run on the storage writer thread after any queued // flushes, so the loop never blocks on disk and a queued save // cannot resurrect the files. - SessionRequest::Delete { id } => { + SessionRequest::Delete { + id, + caller_id, + trusted_ui_control, + } => { + if let Err(error) = authorize_ui_delete(caller_id.as_ref(), trusted_ui_control) { + let _ = reply_tx.send(Err(error)); + return; + } let id = match parse_session_id(&id) { Ok(id) => id, Err(e) => { @@ -2102,6 +2127,7 @@ impl<'t> EventLoop<'t> { Ok(replacement) => replacement, Err(error) => { warn!(session_id = %self.sessions[idx].id(), %error, "invalid reset session lineage"); + self.sessions[idx].app.status = Status::error(error.to_string()); return; } }; @@ -2109,6 +2135,7 @@ impl<'t> EventLoop<'t> { Ok(identity) => identity, Err(error) => { warn!(session_id = %replacement.id, %error, "invalid reset session identity"); + self.sessions[idx].app.status = Status::error(error.to_string()); return; } }; @@ -2119,6 +2146,7 @@ impl<'t> EventLoop<'t> { %error, "failed to replace reset session lineage" ); + self.sessions[idx].app.status = Status::error(error.to_string()); return; } self.respawn_agent(idx, Vec::new(), Vec::new(), identity); @@ -2442,11 +2470,12 @@ fn scroll_delta(kind: MouseEventKind, lines: u32) -> i32 { #[cfg(test)] mod tests { use super::{ - DIRECT_OUTPUT_MAX_BYTES, DRAIN_BUDGET, DrainScheduler, PAUSED_TEAM_RUN_ID_MAX_BYTES, - TEAM_TOOL_NAME, bounded_direct_output, cancel_stored_session, complete_model_fetch_with, - direct_paused_team_payload, draw_then_post_terminal, paused_team_payload, paused_team_run, - should_save_periodically, startup_login_completed, startup_provider_with, - take_painted_submissions, validated_paused_team_payload, + DELETE_UI_ONLY_ERR, DIRECT_OUTPUT_MAX_BYTES, DRAIN_BUDGET, DrainScheduler, + PAUSED_TEAM_RUN_ID_MAX_BYTES, TEAM_TOOL_NAME, authorize_ui_delete, bounded_direct_output, + cancel_stored_session, complete_model_fetch_with, direct_paused_team_payload, + draw_then_post_terminal, paused_team_payload, paused_team_run, should_save_periodically, + startup_login_completed, startup_provider_with, take_painted_submissions, + validated_paused_team_payload, }; use crate::{AppSession, agent::ModelSlot, components::Status}; use arc_swap::ArcSwap; @@ -2455,7 +2484,7 @@ mod tests { AgentError, ContentBlock, Message, Model, Role, provider::unconfigured_provider, }; use n00n_storage::{ - id::n00nId, + id::{SessionRef, n00nId}, sessions::{StoredDelivery, StoredDirectTool, StoredQueuedMessage, StoredSessionLifecycle}, }; use ratatui::{ @@ -2560,6 +2589,24 @@ mod tests { assert!(Arc::ptr_eq(&model_slot.load_full(), &replacement)); } + #[test] + fn delete_allows_only_ui_callbacks_without_agent_identity() { + assert!(authorize_ui_delete(None, true).is_ok()); + assert_eq!( + authorize_ui_delete(None, false) + .as_ref() + .map_err(String::as_str), + Err(DELETE_UI_ONLY_ERR) + ); + let caller = SessionRef::generate(); + assert_eq!( + authorize_ui_delete(Some(&caller), true) + .as_ref() + .map_err(String::as_str), + Err(DELETE_UI_ONLY_ERR) + ); + } + #[test] fn paused_team_run_requires_matching_team_tool_call() { let tool_result = Message { diff --git a/n00n-ui/src/storage_writer.rs b/n00n-ui/src/storage_writer.rs index f73ef79fb..5b832a285 100644 --- a/n00n-ui/src/storage_writer.rs +++ b/n00n-ui/src/storage_writer.rs @@ -298,6 +298,27 @@ impl StorageWriter { } /// Deletes a session on the writer thread after superseded commands have + pub(crate) fn persist_and_wait( + &self, + session: Box, + timeout: Duration, + ) -> Result<(), SessionError> { + let (done_tx, done_rx) = flume::bounded(1); + self.persist(session, move |result| { + let _ = done_tx.send(result); + }); + match done_rx.recv_timeout(timeout) { + Ok(result) => result, + Err(flume::RecvTimeoutError::Timeout) => { + Err(SessionError::Storage(StorageError::Io(io::Error::new( + io::ErrorKind::TimedOut, + format!("session checkpoint did not complete within {timeout:?}"), + )))) + } + Err(flume::RecvTimeoutError::Disconnected) => Err(writer_gone()), + } + } + /// been rejected by generation. The caller never waits for filesystem I/O. pub fn delete(&self, id: n00nId, done: impl FnOnce(Result<(), SessionError>) + Send + 'static) { let generation = reserve_command(&self.tracker, id);