diff --git a/changelog.d/327.fixed.md b/changelog.d/327.fixed.md new file mode 100644 index 000000000..b0acb016d --- /dev/null +++ b/changelog.d/327.fixed.md @@ -0,0 +1 @@ +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-agent/src/agent/run.rs b/n00n-agent/src/agent/run.rs index 0a8b8240a..2e8bb21bc 100644 --- a/n00n-agent/src/agent/run.rs +++ b/n00n-agent/src/agent/run.rs @@ -355,6 +355,32 @@ impl<'h> Agent<'h> { self.total_cost } + /// 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, + self.mcp.as_ref(), + id, + name, + input, + &ctx, + tool_dispatch::Emit::Notify, + ) + .await; + self.event_tx + .send(AgentEvent::ToolDone(Box::new(done.clone())))?; + Ok(done) + } + /// Runs the agent loop with the given input. /// /// # Errors @@ -2077,6 +2103,18 @@ mod tests { (agent, event_rx) } + #[test] + 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 identity = SessionIdentity::root(SessionRef::generate()); + agent.identity = Some(identity.clone()); + + let ctx = agent.tool_context(); + + assert_eq!(ctx.identity, Some(identity)); + } + fn make_agent_with_config( provider: MockProvider, history: &mut History, 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-config/src/lib.rs b/n00n-config/src/lib.rs index afcd3a9ec..1f9c336e0 100644 --- a/n00n-config/src/lib.rs +++ b/n00n-config/src/lib.rs @@ -27,6 +27,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; @@ -45,6 +48,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 = @@ -195,6 +201,16 @@ pub enum ConfigError { "invalid config: agent.fusion.sidekick_tier must be weak, medium, or strong, got {tier:?}" )] InvalidFusionSidekickTier { tier: crate::providers::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( @@ -496,6 +512,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, @@ -527,6 +546,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 ); @@ -1111,6 +1133,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, @@ -1226,6 +1257,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), @@ -1239,6 +1277,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)] @@ -1413,6 +1464,7 @@ impl Config { pub fn validate(&self) -> Result<(), ConfigError> { self.ui.validate_all()?; self.agent.validate()?; + self.agent.validate_lineage_limits()?; self.provider.validate()?; self.provider.validate_openai_coding_plan_slots()?; if self.agent.fusion.sidekick_tier == crate::providers::Tier::Compaction { 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/agent.rs b/n00n-lua/src/api/agent.rs index 756b9ab3b..435baa109 100644 --- a/n00n-lua/src/api/agent.rs +++ b/n00n-lua/src/api/agent.rs @@ -596,6 +596,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. @@ -661,11 +662,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")?; @@ -800,6 +796,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(); diff --git a/n00n-lua/src/api/async.rs b/n00n-lua/src/api/async.rs index 14fe989df..84122f193 100644 --- a/n00n-lua/src/api/async.rs +++ b/n00n-lua/src/api/async.rs @@ -593,7 +593,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..c1c7bf6a9 100644 --- a/n00n-lua/src/api/session.rs +++ b/n00n-lua/src/api/session.rs @@ -5,8 +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::api::util::command::{SessionBootstrap, SessionReply, SessionRequest, UiAction}; +use crate::api::util::convert::{json_to_lua, lua_to_json}; +use crate::runtime::{active_session_identity, active_trusted_ui_control}; const NO_UI_ERR: &str = "no interactive UI attached"; @@ -110,14 +111,27 @@ 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. /// /// @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) 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 }) @@ -127,13 +141,47 @@ async fn new( #[ctx] tx: Option>, opts: Option, ) -> LuaResult { - 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), + 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")?, + 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, @@ -142,6 +190,8 @@ async fn new( prompt, focus, parent_id, + caller_id, + bootstrap, }, ) .await @@ -167,6 +217,7 @@ async fn prompt( text: String, opts: Option
, ) -> LuaResult { + let caller_id = active_session_identity(&lua).map(|identity| identity.session_id().clone()); let (id, steer, control) = match opts { Some(opts) => ( opts.get("session")?, @@ -183,6 +234,8 @@ async fn prompt( text, steer, control, + caller_id, + host_control: false, }, ) .await @@ -198,7 +251,17 @@ 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, + host_control: false, + }, + ) + .await } /// Renames a session, live or stored. @@ -234,9 +297,13 @@ lua_table! { #[cfg(test)] mod tests { use super::*; + use n00n_agent::{cancel::CancelToken, tools::SessionIdentity}; + 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 +311,148 @@ 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(SessionIdentity::root(caller_id.clone())), + ), + ); + let expected_caller_id = caller_id; + 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 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, 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 + local deleted, delete_err = session.delete("target") + if delete_err then error(delete_err) end + return child, status, deleted + "#, + ) + .eval_async(), + ) + .unwrap(); + checker.join().unwrap(); + assert_eq!(child_id, "child"); + assert_eq!(prompt_status, "queued"); + assert!(deleted); + } + + #[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; + 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); @@ -303,13 +512,20 @@ 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, + host_control, + }, reply_tx, }) = rx.recv() else { panic!("expected cancel request"); }; assert_eq!(id, "abc"); + assert_eq!(caller_id, None); + assert!(!host_control); reply_tx.send(Ok(json!(true))).unwrap(); }); let (val, err): (bool, Option) = @@ -340,6 +556,8 @@ mod tests { text, steer, control, + caller_id, + host_control, }, reply_tx, }) = rx.recv() @@ -350,6 +568,8 @@ mod tests { assert_eq!(text, "hi"); 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) = @@ -359,6 +579,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/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/command.rs b/n00n-lua/src/api/util/command.rs index c3d55db15..f82430aff 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 { @@ -403,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, @@ -415,21 +423,29 @@ pub enum SessionRequest { prompt: Option, focus: bool, parent_id: Option, + caller_id: Option, + bootstrap: Option, }, Prompt { id: Option, text: String, steer: bool, control: bool, + caller_id: Option, + host_control: bool, }, Cancel { id: String, + caller_id: Option, + host_control: bool, }, Focus { id: String, }, Delete { id: String, + caller_id: Option, + trusted_ui_control: bool, }, SetTitle { id: String, diff --git a/n00n-lua/src/api/util/ctx.rs b/n00n-lua/src/api/util/ctx.rs index f6593db27..6fd185fcd 100644 --- a/n00n-lua/src/api/util/ctx.rs +++ b/n00n-lua/src/api/util/ctx.rs @@ -10,7 +10,7 @@ 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}; @@ -118,6 +118,7 @@ enum Caps { config: Arc, workflow: bool, audience: ToolAudience, + identity: Option, }, Restore { state: Option, @@ -153,6 +154,7 @@ impl LuaCtx { config: Arc::clone(&ctx.config), workflow: ctx.workflow, audience: ctx.audience, + identity: ctx.identity.clone(), }, ) } @@ -179,6 +181,14 @@ impl LuaCtx { } } + pub(crate) fn session_identity(&self) -> Option { + match &self.caps { + Caps::Handler { agent, .. } => agent.identity.clone(), + Caps::Start { identity, .. } => identity.clone(), + Caps::Restore { .. } => None, + } + } + fn config(&self) -> Option<&AgentConfig> { match &self.caps { Caps::Handler { agent, .. } => Some(&agent.config), @@ -613,9 +623,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_identity_and_resets_per_call_state() { + let ctx = populated_ctx(); + 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.identity, expected_identity); assert!(matches!(agent.deadline, Deadline::None)); assert_eq!(agent.tool_output_lines, ToolOutputLines::default()); assert!(agent.local_tools.is_empty()); @@ -636,6 +649,7 @@ mod tests { ); let inner = agent.to_tool_context(); assert_eq!(inner.tool_use_id, None); + 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 87750fc25..aed081fa3 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/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 bdbd88c33..53b4be2ef 100644 --- a/n00n-lua/src/runtime.rs +++ b/n00n-lua/src/runtime.rs @@ -16,7 +16,8 @@ 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; @@ -150,6 +151,7 @@ pub enum Request { input: Value, ctx: Box, deadline: Option, + nested: bool, reply: flume::Sender, live: Option, }, @@ -157,12 +159,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 { @@ -179,6 +183,7 @@ pub enum Request { plugin: Arc, command: Arc, args: String, + identity: Option, }, CollectPromptSlots { reply: flume::Sender, @@ -249,6 +254,7 @@ pub enum Request { input: Value, live: LiveCtx, ctx: Box, + nested: bool, reply: flume::Sender<()>, }, } @@ -341,6 +347,8 @@ 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) 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>, @@ -358,6 +366,7 @@ impl TaskCell { cancel: CancelToken, deadline: Option, live: Option, + identity: Option, ) -> Self { Self { cancel, @@ -375,6 +384,8 @@ impl TaskCell { live, root_buf: None, live_sink: None, + identity, + trusted_ui_control: false, inline_spawn: None, bufs_claim: Weak::new(), async_tasks: Cell::new(0), @@ -651,7 +662,13 @@ 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 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 { @@ -680,8 +697,25 @@ 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, 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( + lua, + 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 { @@ -761,6 +795,18 @@ pub(crate) fn active_task(lua: &Lua) -> TaskHandle { ) } +pub(crate) fn active_session_identity(lua: &Lua) -> Option { + let handle = lua.app_data_ref::()?; + 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) } @@ -781,12 +827,17 @@ pub(crate) fn enqueue_async_task( on_finish: Option, ) -> Result<(), mlua::Error> { let handle = lua.app_data_ref::(); - let (cancel, live_ctx, parent_deadline) = 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.cancel.clone(), + cell.live.clone(), + cell.deadline.get(), + cell.identity.clone(), + ) } - None => (CancelToken::none(), None, None), + None => (CancelToken::none(), None, None, None), }; let deadline = @@ -798,6 +849,7 @@ pub(crate) fn enqueue_async_task( cancel, deadline, live_ctx, + identity, owner: None, parent: None, }; @@ -1002,6 +1054,7 @@ pub(crate) struct PendingAsyncTask { pub cancel: CancelToken, pub deadline: Option, pub live_ctx: 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. @@ -1164,7 +1217,12 @@ fn spawn_async_task( } }; - let mut cell = TaskCell::new(task.cancel.clone(), task.deadline, task.live_ctx.clone()); + let mut cell = TaskCell::new( + task.cancel.clone(), + task.deadline, + task.live_ctx.clone(), + task.identity.clone(), + ); cell.async_cleanup_cutoff.set( task.deadline .map(|deadline| checked_deadline_after(deadline, DEADLINE_CLEANUP_GRACE)), @@ -1234,7 +1292,19 @@ fn spawn_runtime_request( 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, @@ -1242,6 +1312,7 @@ fn spawn_runtime_request( input, mut ctx, deadline, + nested: _, reply, live, } => { @@ -1268,6 +1339,7 @@ fn spawn_runtime_request( plugin, tool, input, + nested: _, reply, } => { let lua = rt.lua.clone(); @@ -1285,6 +1357,7 @@ fn spawn_runtime_request( plugin, tool, input, + nested: _, reply, } => { let lua = rt.lua.clone(); @@ -1306,6 +1379,7 @@ fn spawn_runtime_request( input, live, ctx, + nested: _, reply, } => { let func = { @@ -1387,7 +1461,9 @@ async fn drain_runtime( while !request_closed { match request_rx.try_recv() { Ok(request) => { - if let Some(request) = spawn_runtime_request(rt, ex, gate, lifecycle, request) { + if let Some(request) = + spawn_runtime_request(rt, ex, gate, lifecycle, request, false) + { deferred.push_back(request); } } @@ -1445,7 +1521,9 @@ async fn drain_runtime( 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) { + if let Some(request) = + spawn_runtime_request(rt, ex, gate, lifecycle, *request, false) + { deferred.push_back(request); } } @@ -2285,6 +2363,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 = LuaCtx::restore(item.tool_output_lines, item.state); @@ -2592,7 +2671,11 @@ async fn run_tool_start( ctx: Box, ) { let _context_liveness = ContextLivenessGuard(ctx.context_liveness()); - let scope = TaskScope::new(lua, TaskCell::new(ctx.cancel.clone(), None, Some(live))); + let identity = ctx.session_identity(); + let scope = TaskScope::new( + lua, + TaskCell::new(ctx.cancel.clone(), None, Some(live), identity), + ); let run = async { let input_lua = json_to_lua(lua, &input)?; let ctx_ud = lua.create_userdata(*ctx)?; @@ -2698,6 +2781,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 identity = ctx.session_identity(); let input_lua = match json_to_lua(&lua, &input) { Ok(v) => v, @@ -2714,7 +2798,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, identity.clone()); cell.live_sink = live_sink; let scope = TaskScope::new(&lua, cell); let handle = Arc::clone(scope.handle()); @@ -2805,7 +2889,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, identity); cell.root_buf = Some(root); let mut warm = warm_tools.borrow_mut(); warm.push_back(WarmTool { @@ -2964,7 +3048,7 @@ pub fn spawn( } request @ (Request::CallTool { .. } | Request::StartTool { .. }) => { let deferred_request = - spawn_runtime_request(&rt, &ex, &gate, &lifecycle, request); + spawn_runtime_request(&rt, &ex, &gate, &lifecycle, request, false); debug_assert!(deferred_request.is_none()); } Request::ClearPlugin { plugin, reply } => { @@ -2989,6 +3073,7 @@ pub fn spawn( plugin, command, args, + identity, } => { let handler_fn = rt.lua.app_data_ref::().and_then(|m| { @@ -3002,7 +3087,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, false, run).await { tracing::warn!(plugin = %plugin, command = %command, error = %e, "command handler failed"); } }) @@ -3013,6 +3098,7 @@ pub fn spawn( plugin, tool, input, + nested: _, reply, } => { let res = @@ -3023,6 +3109,7 @@ pub fn spawn( plugin, tool, input, + nested: _, reply, } => { let res = LuaRuntime::compute_permission_scopes( @@ -3230,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"); } @@ -3259,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(); @@ -3392,7 +3479,7 @@ mod tests { } fn task_cell(live: Option) -> TaskCell { - TaskCell::new(CancelToken::none(), None, live) + TaskCell::new(CancelToken::none(), None, live, None) } #[test] @@ -3555,6 +3642,7 @@ mod tests { CancelToken::none(), Some(Instant::now() + Duration::from_millis(10)), None, + None, ), ); @@ -3576,6 +3664,7 @@ mod tests { CancelToken::none(), Some(Instant::now() + DISPATCH_POLL_INTERVAL.saturating_mul(100)), None, + None, ))); let (ready_tx, ready_rx) = std::sync::mpsc::sync_channel(0); let waiter = smol::spawn({ @@ -3620,6 +3709,7 @@ mod tests { CancelToken::none(), Some(Instant::now() + Duration::from_secs(1)), None, + None, ))); let waiter = smol::spawn({ let handle = Arc::clone(&handle); @@ -3645,7 +3735,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(); @@ -3686,7 +3776,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), None).unwrap(); @@ -3714,7 +3804,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), None).unwrap(); let queue = lua.app_data_ref::().unwrap(); @@ -3727,13 +3817,28 @@ mod tests { ); } + #[test] + fn enqueue_async_task_inherits_session_identity() { + let lua = enqueue_test_lua(); + let identity = SessionIdentity::root(n00n_storage::id::SessionRef::generate()); + let _h = set_active( + &lua, + TaskCell::new(CancelToken::none(), None, None, Some(identity.clone())), + ); + enqueue_async_task(&lua, enqueue_dummy(&lua), None).unwrap(); + + let queue = lua.app_data_ref::().unwrap(); + let queued = queue.rx.try_recv().unwrap(); + assert_eq!(queued.identity, Some(identity)); + } + #[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(); @@ -3753,7 +3858,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), None).unwrap(); @@ -3765,7 +3870,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), None).unwrap(); @@ -3778,7 +3883,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()); @@ -3820,6 +3925,7 @@ mod tests { cancel, deadline, live_ctx: None, + identity: None, owner: None, parent: None, } @@ -3856,6 +3962,7 @@ mod tests { cancel: token, deadline: None, live_ctx: None, + identity: None, owner: None, parent: None, }; @@ -3919,7 +4026,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] @@ -3930,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); @@ -3979,7 +4098,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); @@ -3995,6 +4114,7 @@ mod tests { CancelToken::none(), Some(deadline), None, + None, ))); lock_cell(&handle).deadline_secs.set(Some(1)); @@ -4009,6 +4129,7 @@ mod tests { CancelToken::none(), Some(Instant::now()), None, + None, ))); let reply = timeout_reply(&handle, "test", "tool"); diff --git a/n00n-lua/src/state.rs b/n00n-lua/src/state.rs index a0eee9c8a..bb780366a 100644 --- a/n00n-lua/src/state.rs +++ b/n00n-lua/src/state.rs @@ -240,6 +240,22 @@ impl PluginStateStore { } } +fn validate_replacement( + inner: &StateInner, + identity: &PluginStateIdentity, + replacement_key: &StateKey, + replacement_value: &Value, +) -> Result<(), PluginStateError> { + let mut candidate = candidate_for(inner, identity, Some(replacement_key))?; + candidate.set_plugin_state( + &replacement_key.plugin, + PLUGIN_STATE_SCHEMA_VERSION, + replacement_key.scope.stored(), + replacement_value.clone(), + )?; + Ok(()) +} + fn candidate_for( inner: &StateInner, identity: &PluginStateIdentity, @@ -272,22 +288,6 @@ fn candidate_for( Ok(candidate) } -fn validate_replacement( - inner: &StateInner, - identity: &PluginStateIdentity, - replacement_key: &StateKey, - replacement_value: &Value, -) -> Result<(), PluginStateError> { - let mut candidate = candidate_for(inner, identity, Some(replacement_key))?; - 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 { diff --git a/n00n-lua/tests/plugin_host.rs b/n00n-lua/tests/plugin_host.rs index 9666a6870..f2f81e1b6 100644 --- a/n00n-lua/tests/plugin_host.rs +++ b/n00n-lua/tests/plugin_host.rs @@ -3207,7 +3207,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({{ @@ -6368,6 +6368,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 @@ -6452,7 +6453,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)) @@ -6581,7 +6582,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)) @@ -6612,7 +6613,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-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-storage/src/sessions.rs b/n00n-storage/src/sessions.rs index 9366b05fe..1fc8142d7 100644 --- a/n00n-storage/src/sessions.rs +++ b/n00n-storage/src/sessions.rs @@ -155,6 +155,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, @@ -178,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 @@ -876,6 +911,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)] @@ -894,6 +933,14 @@ 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 = "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 = "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")] pub thinking: Option, @@ -5096,6 +5143,13 @@ 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.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(); @@ -5103,6 +5157,19 @@ 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.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/agent/agent_loop.rs b/n00n-ui/src/agent/agent_loop.rs index 7f1103529..511ec996a 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, info, 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, @@ -161,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 }); + } } } @@ -195,6 +202,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 +235,66 @@ 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); + while self.answer_rx.lock().await.try_recv().is_ok() {} + + 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_user_response_rx(Arc::clone(&self.answer_rx)) + .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( @@ -537,9 +608,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/agent/mod.rs b/n00n-ui/src/agent/mod.rs index 6f6d8cc3c..29ef2c977 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, @@ -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"); @@ -152,7 +154,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 +218,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 +272,7 @@ fn spawn_agent_internal( queue: queue_rx, cancel_map, init_cancel, - session_id, + identity: identity.clone(), timeouts, openai_options, lua_handle, @@ -293,6 +295,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..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::{ @@ -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, } } @@ -138,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(()); } @@ -168,7 +192,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) @@ -176,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(()); } @@ -199,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 { @@ -215,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(); } @@ -268,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 @@ -279,7 +311,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 +323,16 @@ impl QueueSender { QueueItem::Message { input, delivery, .. } => Some((input.clone(), *delivery)), - QueueItem::Compact { .. } => None, + QueueItem::Compact { .. } | QueueItem::DirectTool { .. } => None, + }) + .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() } @@ -324,6 +365,7 @@ impl QueueReceiver { delivery: Delivery::TurnEnd, .. } | QueueItem::Compact { .. } + | QueueItem::DirectTool { .. } ) { None @@ -334,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 } @@ -356,9 +405,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) } } @@ -465,6 +517,37 @@ 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(); + 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/mod.rs b/n00n-ui/src/app/mod.rs index 933df0728..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) { @@ -1961,7 +1964,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/queue.rs b/n00n-ui/src/app/queue.rs index b5057e2fa..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() @@ -199,6 +198,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 3839fe52e..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,18 +12,30 @@ 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, + SessionError, StoredDelivery, StoredDirectTool, StoredImageMediaType, StoredImageSource, + StoredMcpPrompt, StoredMode, StoredQueuedMessage, StoredSessionLifecycle, + 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; +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 { + 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 +57,9 @@ 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.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 @@ -158,6 +174,7 @@ fn restored_submission( } impl App { + #[allow(dead_code)] pub(crate) fn has_content(&self) -> bool { session_has_content(&self.state.session) } @@ -170,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(), @@ -210,23 +236,12 @@ 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; }; 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()) { @@ -239,7 +254,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 @@ -275,6 +290,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 @@ -376,6 +402,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); @@ -429,7 +464,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; @@ -441,7 +477,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 { @@ -485,6 +521,7 @@ impl App { ))] } + #[allow(dead_code)] pub(crate) fn apply_loaded_session( &mut self, session: AppSession, @@ -512,6 +549,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/app/tests.rs b/n00n-ui/src/app/tests.rs index 8215374a1..e0420ef9c 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; @@ -860,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()); } @@ -886,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); @@ -2465,6 +2471,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); @@ -2585,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(); @@ -3841,7 +3870,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()); @@ -3868,7 +3899,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/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/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 caa2f268a..e320c0a91 100644 --- a/n00n-ui/src/event_loop.rs +++ b/n00n-ui/src/event_loop.rs @@ -7,6 +7,8 @@ //! 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; use std::time::{Duration, Instant}; @@ -20,7 +22,10 @@ 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, truncate_output}, +}; use n00n_config::UiConfig; use n00n_lua::{ EventHandle, HintReader, KeymapReader, LuaCommandReader, SessionReply, SessionRequest, UiAction, @@ -33,7 +38,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, StoredDirectTool, StoredSessionLifecycle, TranscriptEntry, normalize_title, +}; use serde_json::{Value, json}; use tracing::warn; @@ -47,6 +54,7 @@ use crate::components::{ Action, DisplayMessage, DisplayRole, ExitRequest, Status, SubmissionDispatch, }; use crate::input::InputReader; +use crate::session_lineage::{LineageError, LineageLimits, LiveSession, SessionLineageGuard}; use crate::color_compat; use crate::storage_writer::StorageWriter; @@ -61,11 +69,15 @@ 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; /// Tabs carry their in-memory sessions so `/reload` reopens them without a /// disk round-trip; `session_has_content` tells which ones were saved. @@ -128,6 +140,215 @@ 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(|| "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 { + 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 has_restorable_work(session: &AppSession) -> bool { + !session.meta.queued_messages.is_empty() + || !session.meta.queued_submissions.is_empty() + || !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(); + 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.meta.direct_paused_team = None; + session.updated_at = n00n_storage::now_epoch(); + 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( + start: n00nId, + parents: &HashMap>, +) -> 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); + 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 state_revision_or_initial( + snapshot: Option<&n00n_storage::sessions::StoredSessionStateSnapshot>, +) -> 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, + action: &'static str, +) { + if let Err(error) = result { + warn!(%session_id, %error, action, "session lineage cleanup failed"); + } +} + +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() @@ -153,22 +374,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); } } @@ -182,6 +388,7 @@ struct SessionRuntime { shell_tx: flume::Sender, shell_rx: flume::Receiver, last_status: SessionStatus, + direct_bootstrap_active: bool, } impl SessionRuntime { @@ -212,11 +419,45 @@ struct SpawnCtx { available_models: Arc>>, storage_writer: Arc, picker: Arc, + hydrated_roots: RefCell>, } impl SpawnCtx { - fn spawn_runtime(&self, session: AppSession) -> SessionRuntime { + 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) + .map_err(|error| eyre!("invalid session identity: {error}"))?; + if let Some(handle) = &self.lua_event_handle { + 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); let handles = AgentHandles::spawn( @@ -227,7 +468,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(), @@ -252,19 +493,20 @@ 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); } + let last_status = SessionStatus::of(&app); let (shell_tx, shell_rx) = flume::unbounded::(); - SessionRuntime { + Ok(SessionRuntime { app, handles, shell_tx, shell_rx, - last_status: SessionStatus::Idle, - } + last_status, + direct_bootstrap_active, + }) } } @@ -272,6 +514,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, @@ -294,6 +537,7 @@ pub(crate) struct EventLoop<'t> { struct SubmissionPersistence { session_id: n00nId, dispatch: SubmissionDispatch, + execution_started: bool, result: Result<(), SessionError>, } @@ -464,7 +708,7 @@ impl<'t> EventLoop<'t> { mut model, mut needs_login, commands, - sessions, + mut sessions, focused, mut startup_warnings, storage, @@ -510,6 +754,83 @@ 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(|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) { + 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 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( + live_sessions, + 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, @@ -529,12 +850,13 @@ impl<'t> EventLoop<'t> { available_models: bg.available, storage_writer, picker, + hydrated_roots: RefCell::new(HashSet::new()), }; 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")); } @@ -558,6 +880,7 @@ impl<'t> EventLoop<'t> { terminal, sessions: runtimes, focused, + lineage, ctx, input: InputReader::spawn()?, warn_rx: bg.warn_rx, @@ -764,19 +1087,165 @@ 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.checkpoint_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) { + 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); + return; + } + if self.sessions[idx].direct_bootstrap_active { + match &envelope.event { + n00n_agent::AgentEvent::ToolDone(done) => { + 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] + .app + .state + .session + .meta + .direct_output + .is_none() => + { + self.sessions[idx].app.state.session.meta.direct_output = + Some(bounded_direct_output(message, &self.ctx.config)); + self.sessions[idx] + .app + .state + .session + .meta + .direct_output_is_error = true; + } + _ => {} + } + } + 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 { .. } + | n00n_agent::AgentEvent::QueueItemConsumed { .. } => { + 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 { + self.sessions[idx].direct_bootstrap_active = false; + self.sessions[idx] + .app + .state + .session + .meta + .queued_direct_tools + .clear(); + 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); } + fn capture_plugin_state(&mut self, idx: usize) -> std::result::Result<(), String> { + let Some(handle) = self.ctx.lua_event_handle.clone() else { + return Ok(()); + }; + let root_id = self.sessions[idx] + .app + .state + .session + .meta + .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())? + { + Arc::unwrap_or_clone(root) + } else { + AppSession::load(root_id, &self.ctx.storage).map_err(|error| error.to_string())? + }; + 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(()) + } + fn drain_channels(&mut self) -> Result<()> { // Leftovers beyond the budget are picked up right after the next draw. let mut scheduler = DrainScheduler::default(); @@ -886,6 +1355,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, @@ -907,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) => { @@ -915,24 +1393,34 @@ 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_for_delete(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)) - } - Err(e) => Err(e.to_string()), - }; - let _ = reply_tx.send(reply); - }); + self.lineage.remove_sessions(&targets); + targets.reverse(); + delete_sessions_sequentially(&self.ctx.storage_writer, targets, reply_tx); } SessionRequest::Live => { let list: Vec<_> = self @@ -959,12 +1447,26 @@ 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 paused_team = paused_team_run(&history); + 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.app.state.session.meta.direct_output_is_error) + .filter(|_| direct_output.is_some()); + 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, @@ -972,6 +1474,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, })) @@ -985,61 +1488,236 @@ impl<'t> EventLoop<'t> { prompt, focus, parent_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 execution_active = prompt.is_some() || bootstrap.is_some(); + let reservation = self + .lineage + .reserve_new(caller, explicit_parent, execution_active) + .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; + warn_lineage_cleanup( + self.lineage.release(reservation), + caller, + "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) => { + warn_lineage_cleanup( + self.lineage.release(reservation), + caller, + "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 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; + 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 + .queue + .push(QueueItem::DirectTool { + run_id, + tool: bootstrap.tool, + input: bootstrap.input, + }); + Ok(json!("started")) + } else if let Some(prompt) = prompt { + self.lineage + .set_execution_active(id, 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; + 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(); + warn_lineage_cleanup( + self.lineage.rollback_new(id), + id, + "roll back new session", + ); + 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, + host_control, } => { - 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 explicit_target = id.as_deref().map(parse_session_id).transpose()?; + 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(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) => { + let meta = &mut self.sessions[idx].app.state.session.meta; + meta.lifecycle = StoredSessionLifecycle::Running; + meta.direct_paused_team = None; + Ok(state) + } + Err(error) => { + if activated { + warn_lineage_cleanup( + self.lineage.set_execution_active(target, false), + target, + "roll back prompt activity", + ); + } + Err(error) + } + } + })(); + let _ = reply_tx.send(reply); + } + SessionRequest::Cancel { + id, + caller_id, + host_control, + } => { + let reply = (|| { + let requested = parse_session_id(&id)?; + 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) + .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 = match self + .ctx + .storage_writer + .latest_snapshot(session_id) + .map_err(|error| error.to_string())? + { + Some(session) => Arc::unwrap_or_clone(session), + None => AppSession::load(session_id, &self.ctx.storage) + .map_err(|error| error.to_string())?, + }; + 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 + || self.sessions[idx] + .app + .state + .session + .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); + 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; + } + 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}")); } - let actions = self.sessions[idx].app.cancel_current_run(); - self.dispatch(idx, actions); Ok(json!(true)) - }); + })(); let _ = reply_tx.send(reply); } SessionRequest::Focus { id } => { @@ -1107,6 +1785,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; } @@ -1135,15 +1816,42 @@ 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 = 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; + 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 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 runtime = match self.ctx.spawn_runtime(session) { + Ok(runtime) => runtime, + Err(error) => { + warn_lineage_cleanup( + self.lineage.remove_runtime(id), + id, + "roll back runtime activation", + ); + return Err(error.to_string()); + } + }; + let idx = self.push_runtime(runtime); self.set_focus(idx); Ok(()) } @@ -1270,6 +1978,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(); @@ -1282,6 +1991,7 @@ impl<'t> EventLoop<'t> { self.ctx.ui_config.tool_output_lines, &permissions, &mut rt.app, + Some(identity), lua_handle, ); } @@ -1294,34 +2004,80 @@ 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); + rt.app.save_session(); + 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; - 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 + .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; - 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 @@ -1329,15 +2085,36 @@ impl<'t> EventLoop<'t> { let _ = completion_tx.send(SubmissionPersistence { session_id, dispatch: *dispatch, + execution_started, result, }); }); } Action::CancelAgent { run_id } => { - let _ = self.sessions[idx] + let id = self.sessions[idx].id(); + if let Err(error) = self.sessions[idx] .handles .cmd_tx - .try_send(AgentCommand::Cancel { run_id }); + .try_send(AgentCommand::Cancel { run_id }) + { + 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(); + self.sessions[idx].app.state.session.meta.direct_paused_team = None; + 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] @@ -1345,8 +2122,34 @@ 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"); + self.sessions[idx].app.status = Status::error(error.to_string()); + 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"); + self.sessions[idx].app.status = Status::error(error.to_string()); + 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" + ); + self.sessions[idx].app.status = Status::error(error.to_string()); + 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; @@ -1378,7 +2181,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 @@ -1550,6 +2360,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(..) { @@ -1562,6 +2377,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()); } @@ -1648,16 +2470,23 @@ fn scroll_delta(kind: MouseEventKind, lines: u32) -> i32 { #[cfg(test)] mod tests { use super::{ - DRAIN_BUDGET, DrainScheduler, TEAM_TOOL_NAME, complete_model_fetch_with, - draw_then_post_terminal, paused_team_run, should_save_periodically, + 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::{agent::ModelSlot, components::Status}; + use crate::{AppSession, agent::ModelSlot, components::Status}; use arc_swap::ArcSwap; + use n00n_agent::AgentConfig; use n00n_providers::{ AgentError, ContentBlock, Message, Model, Role, provider::unconfigured_provider, }; - use n00n_storage::id::n00nId; + use n00n_storage::{ + id::{SessionRef, n00nId}, + sessions::{StoredDelivery, StoredDirectTool, StoredQueuedMessage, StoredSessionLifecycle}, + }; use ratatui::{ Terminal, backend::{Backend, ClearType, TestBackend, WindowSize}, @@ -1760,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 { @@ -1809,6 +2656,111 @@ 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 { + 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"); + 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!({}), + }]; + 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; + 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()); + } + struct FailingBackend(TestBackend); fn infallible(result: Result) -> T { 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..88e63cad6 --- /dev/null +++ b/n00n-ui/src/session_lineage.rs @@ -0,0 +1,1100 @@ +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)] +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) 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)] +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 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}")] + 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 { + root_session_id: n00nId, + parent_id: Option, + runtime_present: bool, + execution_active: bool, + deleted: bool, +} + +#[derive(Debug, Clone, Copy)] +struct PendingReservation { + caller: n00nId, + parent: n00nId, + root: n00nId, + depth: usize, + execution_active: bool, +} + +#[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, +} + +impl SessionLineageGuard { + pub(crate) fn from_live( + sessions: impl IntoIterator, + limits: LineageLimits, + ) -> Result { + let mut guard = Self { + limits, + sessions: HashMap::new(), + children: HashMap::new(), + lineage_cache: HashMap::new(), + reservations: HashMap::new(), + next_reservation_id: 1, + }; + for session in sessions { + if guard + .sessions + .insert( + session.id, + SessionNode { + root_session_id: session.root_session_id, + parent_id: session.parent_id, + runtime_present: session.runtime_present, + execution_active: session.execution_active, + deleted: false, + }, + ) + .is_some() + { + return Err(LineageError::DuplicateSession(session.id)); + } + } + 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) + } + + 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)); + } + 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))? + .runtime_present = true; + self.sessions + .get_mut(&session.id) + .ok_or(LineageError::UnknownSession(session.id))? + .execution_active = session.execution_active; + return Ok(()); + } + + self.sessions.insert( + session.id, + SessionNode { + root_session_id: session.root_session_id, + parent_id: session.parent_id, + runtime_present: true, + execution_active: session.execution_active, + deleted: false, + }, + ); + if let Err(error) = self.rebuild_topology() { + self.sessions.remove(&session.id); + 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(()) + } + + pub(crate) fn remove_runtime(&mut self, id: n00nId) -> Result<(), LineageError> { + let node = self + .sessions + .get_mut(&id) + .ok_or(LineageError::UnknownSession(id))?; + node.runtime_present = false; + node.execution_active = false; + 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, + active: bool, + ) -> Result { + let node = self + .sessions + .get_mut(&id) + .ok_or(LineageError::UnknownSession(id))?; + if active && (!node.runtime_present || node.deleted) { + 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.runtime_present || node.deleted { + 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, + execution_active: bool, + ) -> Result { + let caller_lineage = self.lineage(caller)?; + let parent = match explicit_parent { + Some(p) => p, + None => 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, + }); + } + 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, + }); + } + + 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, + execution_active, + }, + ); + 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 { + root_session_id: pending.root, + parent_id: Some(pending.parent), + runtime_present: true, + execution_active: pending.execution_active, + deleted: false, + }, + ); + if let Err(error) = self.rebuild_topology() { + self.sessions.remove(&child_id); + 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(()) + } + + pub(crate) fn release(&mut self, reservation: NewReservation) -> Result<(), LineageError> { + self.reservations + .remove(&reservation.id) + .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); + self.rebuild_topology()?; + Ok(()) + } + + 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)); + } + let mut pending = self + .children + .get(&parent) + .into_iter() + .flat_map(|children| children.iter().copied()) + .collect::>(); + let mut descendants = Vec::new(); + while let Some(id) = pending.pop() { + 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()); + } + } + 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, + explicit_target: Option, + ) -> Result { + let caller_lineage = self.lineage(caller)?; + let target = match explicit_target { + Some(t) => t, + None => caller, + }; + let target_node = self + .sessions + .get(&target) + .ok_or(LineageError::UnknownSession(target))?; + if !target_node.runtime_present || target_node.deleted { + return Err(LineageError::TargetNotLive(target)); + } + 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) + } + + pub(crate) fn descendant_counts(&self, root: n00nId) -> Result { + let mut total = 0; + let mut active = 0; + for id in self.descendants_of(root)? { + let node = self + .sessions + .get(&id) + .ok_or(LineageError::UnknownSession(id))?; + if node.deleted { + continue; + } + total += 1; + if node.execution_active { + active += 1; + } + } + let reserved = self + .reservations + .values() + .filter(|reservation| reservation.root == root) + .count(); + Ok(DescendantCounts { + total, + active, + reserved, + }) + } + + 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() { + 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 node = self + .sessions + .get(&id) + .ok_or(LineageError::UnknownSession(id))?; + let cached = self + .lineage_cache + .get(&id) + .ok_or(LineageError::UnknownSession(id))?; + Ok(SessionLineage { + caller: id, + root: cached.root, + parent: node.parent_id, + depth: cached.depth, + }) + } + + 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 current = start; + loop { + path.push(current); + let parent = self + .sessions + .get(¤t) + .ok_or(LineageError::UnknownSession(current))? + .parent_id; + let Some(parent) = parent else { + return Ok(path); + }; + 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 +} + +#[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, + root_session_id: match parent_id { + Some(p) => p, + None => id, + }, + parent_id, + runtime_present: true, + execution_active: parent_id.is_some(), + } + } + + 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, true) + .expect("root A capacity"); + guard.commit_new(reservation, child_a).expect("child A"); + assert!(matches!( + guard.reserve_new(root_a, None, true), + Err(LineageError::TotalDescendantsExceeded { .. }) + )); + + 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); + } + + #[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), true), + Err(LineageError::ParentMismatch) + )); + assert!(matches!( + guard.reserve_new(id(99), None, true), + 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) + )); + assert!(matches!( + guard.authorize_prompt(sibling, Some(root)), + 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, true).expect("depth one"); + guard + .commit_new(reservation, grandchild) + .expect("grandchild"); + assert!(matches!( + guard.reserve_new(grandchild, None, true), + 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, 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); + 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 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); + 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, true).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, true).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, true), + 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, 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] + ); + 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 ea74dedd7..5b832a285 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)] @@ -114,6 +115,10 @@ enum Op { generation: u64, done: DeleteCallback, }, + Latest { + id: n00nId, + done: flume::Sender>>, + }, #[cfg(test)] Pause { entered: flume::Sender<()>, @@ -197,6 +202,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| Arc::clone(&snapshot.session)); + let _ = done.send(session); + } #[cfg(test)] Op::Pause { entered, release } => { let _ = entered.send(()); @@ -242,6 +255,29 @@ impl StorageWriter { self.enqueue_reserved_snapshot(generation, session); } + 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())?; + 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. pub fn persist( &self, @@ -262,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); @@ -650,6 +707,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", @@ -879,6 +944,51 @@ 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) => 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 065d2ceed..6c7b2b461 100644 --- a/plugins/lib/n00n/subagent.lua +++ b/plugins/lib/n00n/subagent.lua @@ -8,6 +8,38 @@ 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" } + +-- 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 + copy[index] = name + end + return copy +end + +local function excluded_tools(opts) + local excluded = {} + local seen = {} + if opts.allow_orchestration ~= true 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 +57,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 +158,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 +212,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..126c4c9f0 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 = subagent.orchestration_tools() 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: " .. n00n.ui.truncate_text(input.description or input.prompt or "background task", 60).head + 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 7508f8506..6e4b1ad7c 100644 --- a/plugins/team/init.lua +++ b/plugins/team/init.lua @@ -810,16 +810,22 @@ 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: " .. n00n.ui.truncate_text(input.goal or "", 60).head + 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 }) + + 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/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/configuration/_index.md b/site/docs/content/configuration/_index.md index 21f2725b5..104f512bc 100644 --- a/site/docs/content/configuration/_index.md +++ b/site/docs/content/configuration/_index.md @@ -107,9 +107,14 @@ 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) | +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 6194bde7e..e7b85c149 100644 --- a/site/docs/content/lua-api/_index.md +++ b/site/docs/content/lua-api/_index.md @@ -2806,7 +2806,13 @@ 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) for a + + 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. @@ -6003,6 +6009,12 @@ function M.make_local_tool(schema, on_submit) -- 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. -- Returns (result | nil, err, cost, usage, model_spec) -- @@ -6020,6 +6032,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) diff --git a/src/cmd/tui_bridge.rs b/src/cmd/tui_bridge.rs index 1aa3027d0..2462388b3 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 { @@ -130,6 +133,8 @@ fn message_one( text: text.to_owned(), steer: opts.steer, control: opts.control, + caller_id: None, + host_control: true, }, ) .map_err(|e| map_not_found(id, e))?; @@ -137,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(|| { @@ -150,6 +157,8 @@ fn resume_one(tx: &flume::Sender, id: &str) -> ControlResult<()> { text: prompt, steer: true, control: true, + caller_id: None, + host_control: true, }, ) .map_err(|e| map_not_found(id, e))?; @@ -179,8 +188,17 @@ 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))?; + 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))?; Ok(()) } @@ -333,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, @@ -343,12 +361,26 @@ 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(()) } + #[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(); @@ -361,10 +393,18 @@ mod tests { text, steer, control, + caller_id, + host_control, } => { - if id.as_deref() != Some("sess-1") || text != "hi" || !steer || !control { + 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}" + "unexpected prompt id={id:?} text={text:?} steer={steer} control={control} caller_id={caller_id:?}" ))); return; } @@ -379,7 +419,7 @@ mod tests { let backend = tui_backend(tx); backend .message( - "sess-1", + "00000000-0000-7000-8000-000000000001", "hi", &MessageOpts { steer: true, @@ -399,10 +439,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); @@ -412,10 +454,17 @@ mod tests { text, steer, control, + caller_id, + host_control, } => { - if id.as_deref() != Some("sess-1") || !steer || !control { + 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}" + "unexpected prompt id={id:?} steer={steer} control={control} caller_id={caller_id:?}" ))); return; } @@ -437,7 +486,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(()) } @@ -454,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, @@ -494,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(()) } 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 { .. }