diff --git a/src/openhuman/agent/harness/session/builder.rs b/src/openhuman/agent/harness/session/builder.rs index ac4afcd97..42aee9ee9 100644 --- a/src/openhuman/agent/harness/session/builder.rs +++ b/src/openhuman/agent/harness/session/builder.rs @@ -87,6 +87,7 @@ impl AgentBuilder { omit_profile: None, omit_memory_md: None, payload_summarizer: None, + tool_policy: None, archivist_hook: None, unified_compaction_enabled: true, } @@ -326,6 +327,18 @@ impl AgentBuilder { self } + /// Installs pre-execution policy middleware for tool calls. + /// + /// The default policy allows all calls. Custom policies can deny a call + /// before `Tool::execute_with_options` runs. + pub fn tool_policy( + mut self, + policy: Arc, + ) -> Self { + self.tool_policy = Some(policy); + self + } + /// Attach the production [`ArchivistHook`] instance so the session /// turn loop can call [`ArchivistHook::flush_open_segment`] at /// session-wind-down time, guaranteeing the trailing open segment is @@ -553,6 +566,9 @@ impl AgentBuilder { omit_profile: self.omit_profile.unwrap_or(true), omit_memory_md: self.omit_memory_md.unwrap_or(true), payload_summarizer: self.payload_summarizer, + tool_policy: self.tool_policy.unwrap_or_else(|| { + Arc::new(crate::openhuman::agent::tool_policy::AllowAllToolPolicy) + }), last_seen_integrations_hash: 0, archivist_hook: self.archivist_hook, synthesized_tool_names: std::collections::HashSet::new(), diff --git a/src/openhuman/agent/harness/session/turn.rs b/src/openhuman/agent/harness/session/turn.rs index 2dc7e49e4..18269235b 100644 --- a/src/openhuman/agent/harness/session/turn.rs +++ b/src/openhuman/agent/harness/session/turn.rs @@ -25,6 +25,7 @@ use crate::openhuman::agent::harness; use crate::openhuman::agent::hooks::{self, ToolCallRecord, TurnContext}; use crate::openhuman::agent::memory_loader::collect_recall_citations; use crate::openhuman::agent::progress::AgentProgress; +use crate::openhuman::agent::tool_policy::{ToolPolicyDecision, ToolPolicyRequest}; use crate::openhuman::context::prompt::{LearnedContextData, PromptContext, PromptTool}; use crate::openhuman::context::{ReductionOutcome, ARCHIVIST_EXTRACTION_PROMPT}; use crate::openhuman::inference::model_context::context_window_for_model; @@ -1080,76 +1081,102 @@ impl Agent { false, ) } else if let Some(tool) = self.tools.iter().find(|t| t.name() == call.name) { - // Per-call options: ask the tool for markdown output when the - // context manager is configured to prefer it. Tools that - // implement `execute_with_options` will populate - // `markdown_formatted`; others fall through to the default - // implementation which forwards to `execute`. - let prefer_markdown = self.context.prefer_markdown_tool_output(); - let options = ToolCallOptions { prefer_markdown }; - let outcome = tool - .execute_with_options(call.arguments.clone(), options) - .await; - match outcome { - Ok(r) => { - if !r.is_error { - let mut output = r.output_for_llm(prefer_markdown); - if prefer_markdown && r.markdown_formatted.is_some() { - log::debug!( - "[agent_loop] tool={} returned markdown payload bytes={}", - call.name, - output.len() - ); - } - // Issue #574 — if a payload summarizer is wired - // in (orchestrator session only) and the output - // exceeds the configured threshold, hand it to - // the summarizer sub-agent before it enters - // history. On any failure or below-threshold - // payload, leave `output` untouched and let the - // existing tool_result_budget_bytes truncation - // pipeline handle it downstream. - if let Some(ps) = self.payload_summarizer.as_ref() { - log::debug!( - "[agent_loop] payload_summarizer intercepting tool={} bytes={}", - call.name, - output.len() - ); - match ps.maybe_summarize(&call.name, None, &output).await { - Ok(Some(payload)) => { - log::info!( - "[agent_loop] payload_summarizer compressed tool={} {}->{} bytes", - call.name, - payload.original_bytes, - payload.summary_bytes - ); - output = payload.summary; - } - Ok(None) => { - log::debug!( - "[agent_loop] payload_summarizer pass-through tool={} bytes={}", - call.name, - output.len() - ); - } - Err(e) => { - log::warn!( - "[agent_loop] payload_summarizer error tool={} err={} (passing raw payload through)", - call.name, - e - ); + let policy_request = ToolPolicyRequest { + tool_name: call.name.clone(), + arguments: call.arguments.clone(), + session_id: self.event_session_id().to_string(), + channel: self.event_channel().to_string(), + agent_definition_id: self.agent_definition_id.to_string(), + }; + if let ToolPolicyDecision::Deny { reason } = + self.tool_policy.check(&policy_request).await + { + tracing::debug!( + tool = call.name.as_str(), + policy = self.tool_policy.name(), + reason = %reason, + "[agent_loop] tool denied by policy" + ); + ( + format!( + "Tool '{}' denied by policy '{}': {reason}", + call.name, + self.tool_policy.name() + ), + false, + ) + } else { + // Per-call options: ask the tool for markdown output when the + // context manager is configured to prefer it. Tools that + // implement `execute_with_options` will populate + // `markdown_formatted`; others fall through to the default + // implementation which forwards to `execute`. + let prefer_markdown = self.context.prefer_markdown_tool_output(); + let options = ToolCallOptions { prefer_markdown }; + let outcome = tool + .execute_with_options(call.arguments.clone(), options) + .await; + match outcome { + Ok(r) => { + if !r.is_error { + let mut output = r.output_for_llm(prefer_markdown); + if prefer_markdown && r.markdown_formatted.is_some() { + log::debug!( + "[agent_loop] tool={} returned markdown payload bytes={}", + call.name, + output.len() + ); + } + // Issue #574 — if a payload summarizer is wired + // in (orchestrator session only) and the output + // exceeds the configured threshold, hand it to + // the summarizer sub-agent before it enters + // history. On any failure or below-threshold + // payload, leave `output` untouched and let the + // existing tool_result_budget_bytes truncation + // pipeline handle it downstream. + if let Some(ps) = self.payload_summarizer.as_ref() { + log::debug!( + "[agent_loop] payload_summarizer intercepting tool={} bytes={}", + call.name, + output.len() + ); + match ps.maybe_summarize(&call.name, None, &output).await { + Ok(Some(payload)) => { + log::info!( + "[agent_loop] payload_summarizer compressed tool={} {}->{} bytes", + call.name, + payload.original_bytes, + payload.summary_bytes + ); + output = payload.summary; + } + Ok(None) => { + log::debug!( + "[agent_loop] payload_summarizer pass-through tool={} bytes={}", + call.name, + output.len() + ); + } + Err(e) => { + log::warn!( + "[agent_loop] payload_summarizer error tool={} err={} (passing raw payload through)", + call.name, + e + ); + } } } + (output, true) + } else { + ( + format!("Error: {}", r.output_for_llm(prefer_markdown)), + false, + ) } - (output, true) - } else { - ( - format!("Error: {}", r.output_for_llm(prefer_markdown)), - false, - ) } + Err(e) => (format!("Error executing {}: {e}", call.name), false), } - Err(e) => (format!("Error executing {}: {e}", call.name), false), } } else { (format!("Unknown tool: {}", call.name), false) diff --git a/src/openhuman/agent/harness/session/turn_tests.rs b/src/openhuman/agent/harness/session/turn_tests.rs index b5fcdbc89..91456c96f 100644 --- a/src/openhuman/agent/harness/session/turn_tests.rs +++ b/src/openhuman/agent/harness/session/turn_tests.rs @@ -3,12 +3,14 @@ use crate::core::event_bus::{global, init_global, DomainEvent}; use crate::openhuman::agent::dispatcher::XmlToolDispatcher; use crate::openhuman::agent::hooks::{PostTurnHook, TurnContext}; use crate::openhuman::agent::memory_loader::MemoryLoader; +use crate::openhuman::agent::tool_policy::{ToolPolicy, ToolPolicyDecision, ToolPolicyRequest}; use crate::openhuman::inference::provider::{ChatRequest, ChatResponse, Provider}; use crate::openhuman::memory::Memory; use crate::openhuman::tools::Tool; use crate::openhuman::tools::ToolResult; use async_trait::async_trait; use std::collections::HashSet; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use tokio::sync::Mutex as AsyncMutex; use tokio::sync::Notify; @@ -106,6 +108,47 @@ impl Tool for EchoTool { } } +struct CountingTool { + calls: Arc, +} + +#[async_trait] +impl Tool for CountingTool { + fn name(&self) -> &str { + "counting" + } + + fn description(&self) -> &str { + "counting" + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type":"object"}) + } + + async fn execute(&self, _args: serde_json::Value) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(ToolResult::success("counting-output")) + } +} + +struct DenyCountingPolicy; + +#[async_trait] +impl ToolPolicy for DenyCountingPolicy { + fn name(&self) -> &str { + "deny_counting" + } + + async fn check(&self, request: &ToolPolicyRequest) -> ToolPolicyDecision { + assert_eq!(request.tool_name, "counting"); + assert_eq!(request.session_id, "turn-test-session"); + assert_eq!(request.channel, "turn-test-channel"); + assert_eq!(request.agent_definition_id, "main"); + ToolPolicyDecision::deny("locked by test policy") + } +} + struct LongTool; #[async_trait] @@ -378,6 +421,46 @@ async fn execute_tool_call_reports_unknown_tool() { assert!(!record.success); } +#[tokio::test] +async fn execute_tool_call_denies_by_policy_before_tool_runs() { + let workspace = tempfile::TempDir::new().expect("temp workspace"); + let workspace_path = workspace.path().to_path_buf(); + std::mem::forget(workspace); + let memory_cfg = crate::openhuman::config::MemoryConfig { + backend: "none".into(), + ..crate::openhuman::config::MemoryConfig::default() + }; + let mem: Arc = + Arc::from(crate::openhuman::memory::create_memory(&memory_cfg, &workspace_path).unwrap()); + let calls = Arc::new(AtomicUsize::new(0)); + + let agent = Agent::builder() + .provider(Box::new(DummyProvider)) + .tools(vec![Box::new(CountingTool { + calls: Arc::clone(&calls), + })]) + .memory(mem) + .tool_dispatcher(Box::new(XmlToolDispatcher)) + .workspace_dir(workspace_path) + .event_context("turn-test-session", "turn-test-channel") + .tool_policy(Arc::new(DenyCountingPolicy)) + .build() + .unwrap(); + let call = ParsedToolCall { + name: "counting".into(), + arguments: serde_json::json!({ "value": 1 }), + tool_call_id: Some("policy-1".into()), + }; + + let (result, record) = agent.execute_tool_call(&call, 0).await; + assert!(!result.success); + assert!(result.output.contains("denied by policy 'deny_counting'")); + assert!(result.output.contains("locked by test policy")); + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert_eq!(record.name, "counting"); + assert!(!record.success); +} + #[tokio::test] async fn turn_runs_full_tool_cycle_with_context_and_hooks() { let provider_impl = Arc::new(SequenceProvider { diff --git a/src/openhuman/agent/harness/session/types.rs b/src/openhuman/agent/harness/session/types.rs index a85dc7ef9..d93e3b91c 100644 --- a/src/openhuman/agent/harness/session/types.rs +++ b/src/openhuman/agent/harness/session/types.rs @@ -11,6 +11,7 @@ use crate::openhuman::agent::harness::archivist::ArchivistHook; use crate::openhuman::agent::hooks::PostTurnHook; use crate::openhuman::agent::memory_loader::MemoryLoader; use crate::openhuman::agent::progress::AgentProgress; +use crate::openhuman::agent::tool_policy::ToolPolicy; use crate::openhuman::context::prompt::SystemPromptBuilder; use crate::openhuman::context::ContextManager; use crate::openhuman::inference::provider::{ChatMessage, ConversationMessage, Provider}; @@ -153,6 +154,10 @@ pub struct Agent { /// summarizer sub-agent before they enter agent history. pub(super) payload_summarizer: Option>, + /// Pre-execution policy hook for tool calls in this session. The + /// default policy allows all calls so existing agents keep their + /// behaviour unless a caller opts into stricter policy. + pub(super) tool_policy: Arc, /// Hash of the Composio connection set this Agent last reconciled /// against. Compared at top-of-turn to a fresh hash computed from /// [`crate::openhuman::composio::cached_active_integrations`]; on @@ -232,6 +237,8 @@ pub struct AgentBuilder { /// to a `SubagentPayloadSummarizer` instance. pub(super) payload_summarizer: Option>, + /// Optional pre-execution tool policy. Defaults to allow-all. + pub(super) tool_policy: Option>, /// Optional reference to the production `ArchivistHook`. Set when /// `config.learning.episodic_capture_enabled` is true. Used to call /// `flush_open_segment` at the closest available session-end signal. diff --git a/src/openhuman/agent/mod.rs b/src/openhuman/agent/mod.rs index 3cca398c0..4fe3b6641 100644 --- a/src/openhuman/agent/mod.rs +++ b/src/openhuman/agent/mod.rs @@ -41,6 +41,7 @@ pub mod prompts; mod schemas; pub mod stop_hooks; pub mod task_board; +pub mod tool_policy; pub mod tree_loader; pub mod triage; pub use schemas::{ diff --git a/src/openhuman/agent/tool_policy.rs b/src/openhuman/agent/tool_policy.rs new file mode 100644 index 000000000..3028ad65c --- /dev/null +++ b/src/openhuman/agent/tool_policy.rs @@ -0,0 +1,76 @@ +//! Generic pre-execution policy hook for agent tool calls. +//! +//! The default policy preserves existing behaviour. Callers that need a +//! narrower runtime can install a custom policy through `AgentBuilder` and +//! deny a tool before any side effect reaches the tool implementation. + +use async_trait::async_trait; + +/// Snapshot of the tool call and session context a policy can inspect. +#[derive(Debug, Clone)] +pub struct ToolPolicyRequest { + pub tool_name: String, + pub arguments: serde_json::Value, + pub session_id: String, + pub channel: String, + pub agent_definition_id: String, +} + +/// Decision returned by a [`ToolPolicy`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ToolPolicyDecision { + Allow, + Deny { reason: String }, +} + +impl ToolPolicyDecision { + pub fn deny(reason: impl Into) -> Self { + Self::Deny { + reason: reason.into(), + } + } +} + +/// Policy middleware invoked before an agent executes a tool. +#[async_trait] +pub trait ToolPolicy: Send + Sync { + /// Stable policy name for logs and user-visible denial messages. + fn name(&self) -> &str; + + /// Inspect a tool call and decide whether it can execute. + async fn check(&self, request: &ToolPolicyRequest) -> ToolPolicyDecision; +} + +/// Default policy used when no caller installs a stricter one. +#[derive(Debug, Default)] +pub struct AllowAllToolPolicy; + +#[async_trait] +impl ToolPolicy for AllowAllToolPolicy { + fn name(&self) -> &str { + "allow_all" + } + + async fn check(&self, _request: &ToolPolicyRequest) -> ToolPolicyDecision { + ToolPolicyDecision::Allow + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn allow_all_policy_allows_every_call() { + let policy = AllowAllToolPolicy; + let request = ToolPolicyRequest { + tool_name: "echo".into(), + arguments: serde_json::json!({ "value": 1 }), + session_id: "session".into(), + channel: "chat".into(), + agent_definition_id: "orchestrator".into(), + }; + + assert_eq!(policy.check(&request).await, ToolPolicyDecision::Allow); + } +}