Skip to content

Add generic tool policy middleware #2137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
senamakel merged 3 commits into from
May 20, 2026
Merged

Add generic tool policy middleware #2137

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3294b3e
Add generic tool policy middleware
vaddisrinivas May 18, 2026
3dbbcce
Merge branch 'main' into pr/2137
senamakel May 19, 2026
85a9875
fix(agent): use tracing::debug for policy-denial log (addresses @code…
senamakel May 19, 2026
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
16 changes: 16 additions & 0 deletions src/openhuman/agent/harness/session/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(),
Expand Down
157 changes: 92 additions & 65 deletions src/openhuman/agent/harness/session/turn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::provider::{
Expand Down Expand Up @@ -1044,76 +1045,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()
),
Comment on lines +1055 to +1069

@coderabbitai coderabbitai Bot May 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Keep raw policy reasons out of the tool result.

reason is arbitrary policy-provided text, but this branch turns it into tool output that gets persisted and fed back to the model on the next iteration. That makes internal enforcement details or user/session-specific context prompt-visible. Return a generic denial message here and keep the detailed reason only in logs/telemetry.

Suggested change 
             if let ToolPolicyDecision::Deny { reason } =
                 self.tool_policy.check(&policy_request).await
             {
+                let policy_name = self.tool_policy.name().to_string();
                 tracing::debug!(
                     tool = call.name.as_str(),
-                    policy = self.tool_policy.name(),
+                    policy = %policy_name,
                     reason = %reason,
                     "[agent_loop] tool denied by policy"
                 );
                 (
-                    format!(
-                        "Tool '{}' denied by policy '{}': {reason}",
-                        call.name,
-                        self.tool_policy.name()
-                    ),
+                    format!("Tool '{}' denied by policy '{}'", call.name, policy_name),
                     false,
                 )
             } else {
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/openhuman/agent/harness/session/turn.rs` around lines 1055 - 1069, The
branch handling ToolPolicyDecision::Deny currently injects the raw
policy-provided reason into the tool result; instead, keep the detailed reason
only in logs/telemetry (leave the tracing::debug call with reason intact) and
return a generic denial string for the tool output generated in the code that
formats ("Tool '{}' denied by policy '{}'...") so that
self.tool_policy.check(&policy_request).await's reason is not persisted or
exposed to the model; update the formatted message produced for call.name and
self.tool_policy.name to a non-sensitive, generic message (e.g., "Tool denied by
policy") while retaining the existing debug logging that references reason and
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)
Expand Down
83 changes: 83 additions & 0 deletions src/openhuman/agent/harness/session/turn_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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 {
Expand Down
7 changes: 7 additions & 0 deletions src/openhuman/agent/harness/session/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
1 change: 1 addition & 0 deletions src/openhuman/agent/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down
Loading
Loading