Merge pull request #134 from TransformerOptimus/feat/bench-runner-tool-policy

Bench-runner tool fixes: timeouts, ignore-list, result caps (ToolPolicy)

Author: adithyn7

crates/agent/src/agent/config.rs

@@ -52,6 +52,9 @@ pub struct AgentConfig {
     /// before editing it (keyed by `(session_id, turn)`), enabling per-turn undo.
     /// `None` disables checkpoint capture (bench/tests).
     pub checkpoint_dir: Option<PathBuf>,
+    /// Per-binary tool behavior knobs (timeouts, ignore-list, result caps).
+    /// Default = permissive app behavior; `bench-runner` sets `ToolPolicy::bench()`.
+    pub tool_policy: crate::tool::ToolPolicy,
 }
 
 impl AgentConfig {
@@ -71,6 +74,7 @@ impl AgentConfig {
             subagents: None,
             subagent_inheritance: None,
             checkpoint_dir: None,
+            tool_policy: crate::tool::ToolPolicy::default(),
         }
     }
 

crates/agent/src/agent/loop_.rs

@@ -483,6 +483,7 @@ impl AgentLoop {
                 let approval_ref = self.approval_handler.clone();
                 let checkpoint_dir = self.config.checkpoint_dir.clone();
                 let checkpoint_turn = iteration + self.iteration_offset;
+                let tool_policy = self.config.tool_policy.clone();
 
                 join_set.spawn(async move {
                     let result = execute_tool_call_impl(
@@ -497,6 +498,7 @@ impl AgentLoop {
                         approval_ref.as_deref(),
                         checkpoint_dir,
                         checkpoint_turn,
+                        tool_policy,
                     )
                     .await;
                     (idx, tool_call_id, result)
@@ -1037,6 +1039,7 @@ async fn execute_tool_call_impl(
     approval_handler: Option<&dyn ApprovalHandler>,
     checkpoint_dir: Option<std::path::PathBuf>,
     checkpoint_turn: u32,
+    tool_policy: crate::tool::ToolPolicy,
 ) -> ToolResult {
     // Parse arguments first — if this fails, emit a basic ToolStart before the error ToolEnd
     let args: serde_json::Value = match serde_json::from_str(arguments_json) {
@@ -1206,6 +1209,7 @@ async fn execute_tool_call_impl(
         tool_call_id: tool_call_id.to_string(),
         checkpoint_dir,
         checkpoint_turn,
+        policy: tool_policy,
     };
 
     let mut result = match tool.execute(args, &ctx).await {

crates/agent/src/skills/tool.rs

@@ -132,6 +132,7 @@ mod tests {
             tool_call_id: "tc_1".into(),
             checkpoint_dir: None,
             checkpoint_turn: 0,
+            policy: crate::tool::ToolPolicy::default(),
         };
         let result = tool
             .execute(json!({"name": "nope"}), &ctx)
@@ -154,6 +155,7 @@ mod tests {
             tool_call_id: "tc_1".into(),
             checkpoint_dir: None,
             checkpoint_turn: 0,
+            policy: crate::tool::ToolPolicy::default(),
         };
         let result = tool
             .execute(json!({"name": "hello"}), &ctx)
@@ -195,6 +197,7 @@ mod tests {
             tool_call_id: "tc_1".into(),
             checkpoint_dir: None,
             checkpoint_turn: 0,
+            policy: crate::tool::ToolPolicy::default(),
         };
         let result = tool
             .execute(json!({"name": "meta"}), &ctx)
@@ -219,6 +222,7 @@ mod tests {
             tool_call_id: "tc_1".into(),
             checkpoint_dir: None,
             checkpoint_turn: 0,
+            policy: crate::tool::ToolPolicy::default(),
         };
         let err = tool.execute(json!({}), &ctx).await.unwrap_err();
         assert!(err.0.contains("Missing"));

crates/agent/src/subagents/tool.rs

@@ -245,6 +245,10 @@ impl Tool for SpawnSubagentTool {
             subagents: None,
             subagent_inheritance: None, // depth-1 enforced
             checkpoint_dir: self.inherit.checkpoint_dir.clone(),
+            // Subagents are only spawned by the desktop app, which runs the permissive
+            // default policy; the headless bench-runner never spawns them. Default keeps
+            // app behavior unchanged.
+            tool_policy: crate::tool::ToolPolicy::default(),
         };
 
         // Child event channel + drain. Per DEC-1 (in docs/subagent-bugs-and-fixes.md),

crates/agent/src/tool/bash.rs

@@ -56,11 +56,18 @@ impl Tool for BashTool {
             .and_then(|v| v.as_str())
             .unwrap_or("Running command");
 
-        let timeout_ms = args
+        let mut timeout_ms = args
             .get("timeout")
             .and_then(|v| v.as_u64())
             .unwrap_or(DEFAULT_TIMEOUT_MS);
 
+        // Bench policy clamps the model-supplied timeout to a hard ceiling so a single
+        // runaway command can't burn the whole per-cell wall. No-op under the default
+        // (app) policy, which keeps the timeout fully model-controlled.
+        if let Some(ceiling) = ctx.policy.bash_timeout_ceiling_ms {
+            timeout_ms = timeout_ms.min(ceiling);
+        }
+
         // Emit status event
         let _ = ctx.event_tx.send(AgentEvent::ToolStatus {
             session_id: ctx.session_id.clone(),
@@ -231,6 +238,7 @@ mod tests {
             tool_call_id: "tc_1".into(),
             checkpoint_dir: None,
             checkpoint_turn: 0,
+            policy: crate::tool::ToolPolicy::default(),
         };
 
         let tool = BashTool;
@@ -254,6 +262,34 @@ mod tests {
         assert!(result.output.contains("cancelled"));
     }
 
+    #[tokio::test]
+    async fn test_bench_policy_clamps_timeout() {
+        // Under bench policy the ceiling is 300s. A command that exceeds a *small*
+        // model-supplied timeout still times out (the clamp is a min(), so the
+        // smaller value wins) — this proves the clamp path doesn't break the
+        // existing timeout behavior.
+        let dir = tempdir().unwrap();
+        let tool = BashTool;
+        let ctx = ToolContext::test_context_bench(dir.path());
+
+        let result = tool
+            .execute(
+                json!({ "command": "sleep 10", "description": "sleep", "timeout": 100 }),
+                &ctx,
+            )
+            .await
+            .unwrap();
+
+        assert!(result.is_error);
+        assert!(result.output.contains("timed out"));
+    }
+
+    #[test]
+    fn test_bench_policy_ceiling_value() {
+        assert_eq!(crate::tool::ToolPolicy::bench().bash_timeout_ceiling_ms, Some(300_000));
+        assert_eq!(crate::tool::ToolPolicy::default().bash_timeout_ceiling_ms, None);
+    }
+
     #[tokio::test]
     async fn test_working_dir() {
         let dir = tempdir().unwrap();

crates/agent/src/tool/codebase_graph.rs

@@ -7,6 +7,9 @@ use crate::context_engine::ContextEngineApi;
 use crate::error::ToolError;
 use super::{Tool, ToolContext, ToolResult};
 
+/// Max rows rendered per graph section (callers / deps / related) under bench policy.
+const GRAPH_SECTION_CAP: usize = 50;
+
 pub struct CodebaseGraphTool {
     context_engine: Arc<dyn ContextEngineApi>,
 }
@@ -73,7 +76,13 @@ impl Tool for CodebaseGraphTool {
             ));
         }
 
-        let _ = ctx; // ToolContext not needed for graph queries (no working-copy overlay)
+        // Bench policy: cap each rendered section so a high-degree node doesn't dump
+        // hundreds of rows into the context. No-op (unbounded) under the app policy.
+        let section_cap = if ctx.policy.codebase_result_caps {
+            GRAPH_SECTION_CAP
+        } else {
+            usize::MAX
+        };
 
         log::info!(
             "[codebase_graph] query={:?}, function_name={:?}, file_path={:?}, query_type={:?}",
@@ -165,50 +174,47 @@ impl Tool for CodebaseGraphTool {
 
         if show_callers && !callers.is_empty() {
             output.push_str("## Blast Radius (upstream callers)\n\n");
-            for (i, item) in callers.iter().enumerate() {
-                output.push_str(&format!(
-                    "{}. {} ({}, depth: {})\n",
-                    i + 1,
-                    item.name,
-                    item.file_path,
-                    item.depth,
-                ));
-            }
+            render_section(&mut output, &callers, section_cap);
             output.push('\n');
         }
 
         if show_deps && !dependencies.is_empty() {
             output.push_str("## Dependencies (downstream)\n\n");
-            for (i, item) in dependencies.iter().enumerate() {
-                output.push_str(&format!(
-                    "{}. {} ({}, depth: {})\n",
-                    i + 1,
-                    item.name,
-                    item.file_path,
-                    item.depth,
-                ));
-            }
+            render_section(&mut output, &dependencies, section_cap);
             output.push('\n');
         }
 
         if !other.is_empty() && query_type.is_none() {
             output.push_str("## Related\n\n");
-            for (i, item) in other.iter().enumerate() {
-                output.push_str(&format!(
-                    "{}. {} ({}, depth: {})\n",
-                    i + 1,
-                    item.name,
-                    item.file_path,
-                    item.depth,
-                ));
-            }
+            render_section(&mut output, &other, section_cap);
             output.push('\n');
         }
 
         Ok(ToolResult::success(output))
     }
 }
 
+/// Render up to `cap` graph rows; append a `… and N more` line when the section
+/// is truncated so the model knows results were withheld.
+fn render_section(
+    output: &mut String,
+    items: &[&crate::context_engine::GraphResultItem],
+    cap: usize,
+) {
+    for (i, item) in items.iter().take(cap).enumerate() {
+        output.push_str(&format!(
+            "{}. {} ({}, depth: {})\n",
+            i + 1,
+            item.name,
+            item.file_path,
+            item.depth,
+        ));
+    }
+    if items.len() > cap {
+        output.push_str(&format!("… and {} more\n", items.len() - cap));
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -405,6 +411,57 @@ mod tests {
         assert!(result.output.contains("'query' or 'function_name'"));
     }
 
+    #[tokio::test]
+    async fn test_bench_policy_caps_section() {
+        let results: Vec<GraphResultItem> = (0..60)
+            .map(|i| GraphResultItem {
+                name: format!("caller_{i}"),
+                file_path: format!("src/f{i}.go"),
+                chunk_id: format!("g{i}"),
+                depth: 1,
+                direction: Some("caller".into()),
+            })
+            .collect();
+        let mock = MockContextEngine::with_graph_results(results);
+        let tool = CodebaseGraphTool::new(Arc::new(mock));
+        let dir = tempfile::tempdir().unwrap();
+        let ctx = ToolContext::test_context_bench(dir.path());
+
+        let result = tool
+            .execute(json!({"function_name": "f", "query_type": "blast_radius"}), &ctx)
+            .await
+            .unwrap();
+        assert!(!result.is_error);
+        assert!(result.output.contains("caller_49"), "row 50 (index 49) should render");
+        assert!(!result.output.contains("caller_50"), "row 51 must be capped under bench policy");
+        assert!(result.output.contains("… and 10 more"));
+    }
+
+    #[tokio::test]
+    async fn test_default_policy_no_section_cap() {
+        let results: Vec<GraphResultItem> = (0..60)
+            .map(|i| GraphResultItem {
+                name: format!("caller_{i}"),
+                file_path: format!("src/f{i}.go"),
+                chunk_id: format!("g{i}"),
+                depth: 1,
+                direction: Some("caller".into()),
+            })
+            .collect();
+        let mock = MockContextEngine::with_graph_results(results);
+        let tool = CodebaseGraphTool::new(Arc::new(mock));
+        let dir = tempfile::tempdir().unwrap();
+        let ctx = ToolContext::test_context(dir.path());
+
+        let result = tool
+            .execute(json!({"function_name": "f", "query_type": "blast_radius"}), &ctx)
+            .await
+            .unwrap();
+        assert!(!result.is_error);
+        assert!(result.output.contains("caller_59"), "default policy must render all rows");
+        assert!(!result.output.contains("more"));
+    }
+
     #[tokio::test]
     async fn test_graph_function_name_only() {
         let mock = MockContextEngine::with_graph_results(make_caller_results());