Skip to content
Merged
Show file tree
Hide file tree
Changes from 13 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ n00n-markdown = { path = "n00n-markdown" }
n00n-daemon = { path = "n00n-daemon" }
n00n-token-profile = { path = "n00n-token-profile" }
serde = { version = "1", features = ["derive", "rc"] }
serde_json = "1"
serde_json = { version = "1", features = ["arbitrary_precision"] }
Comment thread
coderabbitai[bot] marked this conversation as resolved.
sonic-rs = "0.5.8"
toon-format = { version = "0.5", default-features = false }
tiktoken-rs = "0.12"
Expand Down
1 change: 1 addition & 0 deletions changelog.d/320.added.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Added host-owned, root- and session-scoped Lua plugin state with bounded snapshots, lifecycle restoration, and cleanup.
97 changes: 63 additions & 34 deletions n00n-agent/src/agent/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,16 @@ use crate::cancel::{CancelMap, CancelToken, PreDispatchGate};
use crate::mcp::McpSession;
use crate::permissions::{PermissionAnswer, PermissionManager};
use crate::tools::{
ActiveTools, Deadline, FileReadTracker, LocalTools, ToolAudience, ToolContext, ToolFilter,
ToolRegistry,
ActiveTools, Deadline, FileReadTracker, LocalTools, SessionIdentity, ToolAudience, ToolContext,
ToolFilter, ToolRegistry,
};
use crate::{
AgentConfig, AgentError, AgentEvent, AgentInput, AgentMode, EventSender, ExtractedCommand,
FusionContinuation, FusionPhase, FusionRequestDecision, FusionState, InterruptPoint,
InterruptSource, ToolDoneEvent, TurnCompleteEvent,
};
use n00n_config::{ToolKey, ToolOutputLines};
#[cfg(test)]
use n00n_storage::id::SessionRef;

use crate::tokenize::{
Expand Down Expand Up @@ -167,7 +168,7 @@ pub struct AgentParams {
pub config: Arc<AgentConfig>,
pub tool_output_lines: ToolOutputLines,
pub permissions: Arc<PermissionManager>,
pub session_id: Option<SessionRef>,
pub identity: Option<SessionIdentity>,
pub timeouts: n00n_providers::Timeouts,
pub openai_options: OpenAiOptions,
pub file_tracker: Arc<FileReadTracker>,
Expand Down Expand Up @@ -214,7 +215,7 @@ pub struct Agent<'h> {
thinking_empty_retried: bool,
permissions: Arc<PermissionManager>,
opts: RequestOptions,
session_id: Option<SessionRef>,
identity: Option<SessionIdentity>,
timeouts: n00n_providers::Timeouts,
openai_options: OpenAiOptions,
file_tracker: Arc<FileReadTracker>,
Expand Down Expand Up @@ -273,7 +274,7 @@ impl<'h> Agent<'h> {
post_tool_empty_retried: false,
thinking_empty_retried: false,
opts: RequestOptions::default(),
session_id: params.session_id,
identity: params.identity,
file_tracker: params.file_tracker,
prompt_slots: params.prompt_slots,
subagent_cancels: params.subagent_cancels,
Expand Down Expand Up @@ -502,7 +503,7 @@ impl<'h> Agent<'h> {
event_tx: &self.event_tx,
cancel: &self.cancel,
opts,
session_id: self.session_id.as_ref(),
session_id: self.identity.as_ref().map(SessionIdentity::session_id),
})
.await
}
Expand Down Expand Up @@ -849,20 +850,29 @@ impl<'h> Agent<'h> {
}

fn effective_tool_filter(&self) -> ToolFilter {
if !self.allow_dynamic_mcp_tools
|| !self.tool_filter.matches(crate::mcp::TOOL_SEARCH_TOOL_NAME)
{
return self.tool_filter.clone();
}
let Some(mcp) = self.mcp.as_ref() else {
return self.tool_filter.clone();
};
// Always include tool_search when MCP is present so agents with
// ToolFilter::Only can still discover and run MCP tools via search.
let mut filter = self.tool_filter.clone();
let tool_search = crate::mcp::TOOL_SEARCH_TOOL_NAME;
if crate::tools::is_tool_enabled(&self.config.disabled_tools, tool_search) {
if !filter.matches(tool_search) {
filter = filter.including([tool_search.to_owned()]);
}
} else {
filter = filter.excluding(&[tool_search]);
}
if !self.allow_dynamic_mcp_tools {
return filter;
}
let capability_exclusions = crate::tools::capability_exclusions(&self.model);
let names = mcp.loaded_tool_names().into_iter().filter(|name| {
crate::tools::is_tool_enabled(&self.config.disabled_tools, name)
&& !capability_exclusions.contains(&name.as_str())
});
self.tool_filter.clone().including(names)
filter.including(names)
}

fn tool_context(&self) -> ToolContext {
Expand All @@ -886,6 +896,7 @@ impl<'h> Agent<'h> {
prompt_slots: Arc::clone(&self.prompt_slots),
opts: self.opts.clone(),
subagent_cancels: Arc::clone(&self.subagent_cancels),
identity: self.identity.clone(),
registry: Arc::clone(&self.registry),
workflow: self.workflow,
audience: self.audience,
Expand Down Expand Up @@ -1144,7 +1155,7 @@ impl<'h> Agent<'h> {
&self.event_tx,
&self.cancel,
CompactionTrigger::Auto,
self.session_id.as_ref(),
self.identity.as_ref().map(SessionIdentity::session_id),
&cwd,
None,
)
Expand Down Expand Up @@ -1440,33 +1451,33 @@ mod tests {
}

#[test]
fn dynamic_mcp_filter_preserves_only_without_tool_search() {
fn dynamic_mcp_filter_includes_tool_search_with_mcp() {
let mut history = History::new(Vec::new());
let (mut agent, _) = make_agent(MockProvider::new(Vec::new()), &mut history);
let mcp = crate::mcp::stub_session(&[("srv.fetch_issue", "Fetch a GitHub issue")]);
agent.tool_filter = ToolFilter::Only(vec!["read".into()]);
agent = agent
.with_mcp(Some(mcp.clone()))
.with_dynamic_mcp_tools(true);
agent = agent.with_mcp(Some(mcp)).with_dynamic_mcp_tools(false);

mcp.search_tools("issue").unwrap();
let effective_filter = agent.effective_tool_filter();
assert!(!effective_filter.matches("tool_search"));
assert!(effective_filter.matches("tool_search"));
assert!(effective_filter.matches("read"));
assert!(!effective_filter.matches("write"));
assert!(!effective_filter.matches("srv__fetch_issue"));
let mut definitions = json!([
{"name": "read"},
{"name": "write"},
]);
mcp.extend_tools(&mut definitions);
filter_provider_tools(&mut definitions, &effective_filter, &AgentMode::Build);
let names = definitions
.as_array()
.unwrap()
.iter()
.filter_map(|definition| definition["name"].as_str())
.collect::<Vec<_>>();
}

assert_eq!(names, ["read"]);
#[test]
fn dynamic_mcp_filter_keeps_disabled_tool_search_blocked() {
let mut history = History::new(Vec::new());
let (mut agent, _) = make_agent(MockProvider::new(Vec::new()), &mut history);
let mcp = crate::mcp::stub_session(&[("srv.fetch_issue", "Fetch a GitHub issue")]);
let mut config = (*agent.config).clone();
config
.disabled_tools
.push(crate::mcp::TOOL_SEARCH_TOOL_NAME.into());
agent.config = Arc::new(config);
agent = agent.with_mcp(Some(mcp)).with_dynamic_mcp_tools(true);

assert!(!agent.effective_tool_filter().matches("tool_search"));
}

#[test]
Expand All @@ -1489,6 +1500,24 @@ mod tests {
assert!(!agent.effective_tool_filter().matches(DISABLED_MCP_TOOL));
}

#[test]
fn dynamic_mcp_filter_includes_loaded_tools_with_flag() {
let mut history = History::new(Vec::new());
let (mut agent, _) = make_agent(MockProvider::new(Vec::new()), &mut history);
let mcp = crate::mcp::stub_session(&[("srv.fetch_issue", "Fetch a GitHub issue")]);
agent.tool_filter = ToolFilter::Only(vec!["read".into()]);
agent = agent
.with_mcp(Some(mcp.clone()))
.with_dynamic_mcp_tools(true);

mcp.search_tools("issue").unwrap();
let effective_filter = agent.effective_tool_filter();
assert!(effective_filter.matches("tool_search"));
assert!(effective_filter.matches("srv__fetch_issue"));
assert!(effective_filter.matches("read"));
assert!(!effective_filter.matches("write"));
}

#[test]
fn estimate_message_tokens_empty_is_zero() {
assert_eq!(estimate_message_tokens(&[], ""), 0);
Expand Down Expand Up @@ -2141,7 +2170,7 @@ mod tests {
},
std::path::PathBuf::from("/tmp"),
)),
session_id: None,
identity: None,
timeouts: n00n_providers::Timeouts::default(),
openai_options: OpenAiOptions::default(),
file_tracker: FileReadTracker::fresh(),
Expand Down Expand Up @@ -3354,7 +3383,7 @@ mod tests {
},
std::path::PathBuf::from("/tmp"),
)),
session_id: None,
identity: None,
timeouts: n00n_providers::Timeouts::default(),
openai_options: OpenAiOptions::default(),
file_tracker: FileReadTracker::fresh(),
Expand Down
8 changes: 5 additions & 3 deletions n00n-agent/src/headless.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,9 @@ use crate::cancel::{CancelMap, CancelToken};
use crate::permissions::PermissionManager;
use crate::prompt::ResolvedSlots;
use crate::template;
use crate::tools::{DescriptionContext, FileReadTracker, ToolAudience, ToolFilter, ToolRegistry};
use crate::tools::{
DescriptionContext, FileReadTracker, SessionIdentity, ToolAudience, ToolFilter, ToolRegistry,
};
use crate::{
Agent, AgentConfig, AgentEvent, AgentInput, AgentMode, AgentParams, AgentRunParams, Envelope,
EventSender, ImageSource, McpHandle, McpSession, PermissionsConfig, ToolOutput,
Expand Down Expand Up @@ -263,7 +265,7 @@ pub fn spawn(params: HeadlessParams) -> HeadlessHandle {
params.permissions_config,
working_dir_path,
)),
session_id: Some(session_ref_clone.clone()),
identity: Some(SessionIdentity::root(session_ref_clone.clone())),
timeouts: params.timeouts,
openai_options: params.openai_options,
file_tracker: FileReadTracker::fresh(),
Expand Down Expand Up @@ -518,7 +520,7 @@ pub fn spawn_interactive(params: InteractiveParams) -> InteractiveHandle {
config: Arc::clone(&params.config),
tool_output_lines: ToolOutputLines::default(),
permissions: Arc::clone(&permissions),
session_id: Some(session_ref_clone.clone()),
identity: Some(SessionIdentity::root(session_ref_clone.clone())),
timeouts: params.timeouts,
openai_options: params.openai_options,
file_tracker: Arc::clone(&file_tracker),
Expand Down
70 changes: 69 additions & 1 deletion n00n-agent/src/tools/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,40 @@ pub fn timeout_annotation(secs: u64) -> String {
pub type LocalToolFn = Arc<dyn Fn(&Value) -> Result<String, String> + Send + Sync>;
pub type LocalTools = Arc<HashMap<String, LocalToolFn>>;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SessionIdentity {
session_id: SessionRef,
root_session_id: SessionRef,
}

impl SessionIdentity {
#[must_use]
pub fn root(session_id: SessionRef) -> Self {
Self {
root_session_id: session_id.clone(),
session_id,
}
}

#[must_use]
pub fn child(session_id: SessionRef, root_session_id: SessionRef) -> Self {
Self {
session_id,
root_session_id,
}
}

#[must_use]
pub fn session_id(&self) -> &SessionRef {
&self.session_id
}

#[must_use]
pub fn root_session_id(&self) -> &SessionRef {
&self.root_session_id
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

#[derive(Clone)]
pub struct ToolContext {
pub provider: Arc<dyn Provider>,
Expand All @@ -321,6 +355,8 @@ pub struct ToolContext {
pub prompt_slots: Arc<crate::prompt::ResolvedSlots>,
pub opts: RequestOptions,
pub subagent_cancels: Arc<CancelMap<String>>,
/// Immutable session and root identity of the agent executing this tool.
pub identity: Option<SessionIdentity>,
pub registry: Arc<ToolRegistry>,
pub tool_filter: ToolFilter,
pub workflow: bool,
Expand Down Expand Up @@ -588,6 +624,7 @@ pub fn interpreter_ctx(
prompt_slots: Arc::new(crate::prompt::ResolvedSlots::default()),
opts: RequestOptions::default(),
subagent_cancels: Arc::new(CancelMap::new()),
identity: None,
registry,
tool_filter: ToolFilter::All,
workflow: false,
Expand Down Expand Up @@ -629,7 +666,8 @@ pub mod test_support {

use super::{
AgentMode, Arc, CancelToken, DescriptionContext, FileReadTracker, LazyLock,
PermissionManager, ToolContext, ToolRegistry, Value, interpreter_ctx, registry,
PermissionManager, SessionIdentity, ToolContext, ToolRegistry, Value, interpreter_ctx,
registry,
};

pub const GUARDED_TOOL_NAME: &str = "guarded_mock";
Expand Down Expand Up @@ -705,6 +743,9 @@ pub mod test_support {
Arc::new(ToolRegistry::new()),
);
ctx.tool_use_id = tool_use_id.map(String::from);
ctx.identity = Some(SessionIdentity::root(
n00n_storage::id::SessionRef::generate(),
));
ctx
}

Expand Down Expand Up @@ -745,6 +786,33 @@ mod tests {

const LINE_LIMIT: usize = 500;

#[test]
fn root_session_identity_uses_one_id() {
let session_id = n00n_storage::id::SessionRef::generate();
let identity = SessionIdentity::root(session_id.clone());

assert_eq!(identity.session_id(), &session_id);
assert_eq!(identity.root_session_id(), &session_id);
}

#[test]
fn descendant_session_identities_inherit_root_and_remain_distinct() {
let root = SessionIdentity::root(n00n_storage::id::SessionRef::generate());
let child = SessionIdentity::child(
n00n_storage::id::SessionRef::generate(),
root.root_session_id().clone(),
);
let grandchild = SessionIdentity::child(
n00n_storage::id::SessionRef::generate(),
child.root_session_id().clone(),
);

assert_ne!(child.session_id(), root.session_id());
assert_ne!(grandchild.session_id(), child.session_id());
assert_eq!(child.root_session_id(), root.root_session_id());
assert_eq!(grandchild.root_session_id(), root.root_session_id());
}

#[test_case(true ; "vision_model_keeps_view_image")]
#[test_case(false ; "text_only_model_loses_view_image")]
fn from_config_gates_view_image_on_vision(vision: bool) {
Expand Down
Loading