Add generic tool policy middleware (#2137)

Co-authored-by: Steven Enamakel <enamakel@tinyhumans.ai>

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<dyn crate::openhuman::agent::tool_policy::ToolPolicy>,
+    ) -> 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(),

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)

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<AtomicUsize>,
+}
+
+#[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<ToolResult> {
+        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<dyn Memory> =
+        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 {

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<Arc<dyn crate::openhuman::agent::harness::payload_summarizer::PayloadSummarizer>>,
+    /// 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<dyn ToolPolicy>,
     /// 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<Arc<dyn crate::openhuman::agent::harness::payload_summarizer::PayloadSummarizer>>,
+    /// Optional pre-execution tool policy. Defaults to allow-all.
+    pub(super) tool_policy: Option<Arc<dyn ToolPolicy>>,
     /// 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.

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::{

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<String>) -> 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);
+    }
+}