feat: apply max budget enforcement to all safe-output tools

Previously only update-work-item and comment-on-work-item had the max
frontmatter option and budget enforcement. This change:

- Adds max field (default: 1) to CreateWorkItemConfig, CreatePrConfig,
  CreateWikiPageConfig, and UpdateWikiPageConfig
- Replaces per-tool budget tracking in execute.rs with a generic
  HashMap-based approach that automatically applies to all budgeted tools
- Adds MaxConfig helper for extracting max from any tool config JSON
- Adds extract_entry_context() for human-readable log context across
  all tool types (work item IDs, titles, wiki paths)
- Updates AGENTS.md to document max for all safe-output tools
- Adds tests for new budget enforcement, context extraction, and
  MaxConfig deserialization

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: jamesadevine

AGENTS.md

@@ -800,6 +800,7 @@ Creates an Azure DevOps work item.
 - `assignee` - User to assign (email or display name)
 - `tags` - List of tags to apply
 - `custom-fields` - Map of custom field reference names to values (e.g., `Custom.MyField: "value"`)
+- `max` - Maximum number of create-work-item outputs allowed per run (default: 1)
 - `artifact-link` - Configuration for GitHub Copilot artifact linking:
   - `enabled` - Whether to add an artifact link (default: false)
   - `repository` - Repository name override (defaults to BUILD_REPOSITORY_NAME)
@@ -897,6 +898,7 @@ Note: The source branch name is auto-generated from a sanitized version of the P
 - `reviewers` - List of reviewer emails to add
 - `labels` - List of labels to apply
 - `work-items` - List of work item IDs to link
+- `max` - Maximum number of create-pull-request outputs allowed per run (default: 1)
 
 **Multi-repository support:**
 When `workspace: root` and multiple repositories are checked out, agents can create PRs for any allowed repository:
@@ -974,6 +976,7 @@ safe-outputs:
     path-prefix: "/agent-output"    # Optional — prepended to the agent-supplied path (restricts write scope)
     title-prefix: "[Agent] "        # Optional — prepended to the last path segment (the page title)
     comment: "Created by agent"     # Optional — default commit comment when agent omits one
+    max: 1                          # Maximum number of create-wiki-page outputs allowed per run (default: 1)
 ```
 
 Note: `wiki-name` is required. If it is not set, execution fails with an explicit error message.
@@ -995,6 +998,7 @@ safe-outputs:
     path-prefix: "/agent-output"    # Optional — prepended to the agent-supplied path (restricts write scope)
     title-prefix: "[Agent] "        # Optional — prepended to the last path segment (the page title)
     comment: "Updated by agent"     # Optional — default commit comment when agent omits one
+    max: 1                          # Maximum number of update-wiki-page outputs allowed per run (default: 1)
 ```
 
 Note: `wiki-name` is required. If it is not set, execution fails with an explicit error message.

src/execute.rs

@@ -5,13 +5,15 @@
 
 use anyhow::{Result, bail};
 use log::{debug, error, info, warn};
+use serde::Deserialize;
 use serde_json::Value;
+use std::collections::HashMap;
 use std::path::Path;
 
 use crate::ndjson::{self, SAFE_OUTPUT_FILENAME};
 use crate::tools::{
-    CommentOnWorkItemConfig, CommentOnWorkItemResult, CreatePrResult, CreateWikiPageResult, CreateWorkItemResult, ExecutionContext, ExecutionResult,
-    Executor, UpdateWikiPageResult, UpdateWorkItemConfig, UpdateWorkItemResult,
+    CreatePrResult, CreateWikiPageResult, CreateWorkItemResult, CommentOnWorkItemResult,
+    ExecutionContext, ExecutionResult, Executor, UpdateWikiPageResult, UpdateWorkItemResult,
 };
 
 // Re-export memory types for use by main.rs
@@ -87,15 +89,22 @@ pub async fn execute_safe_outputs(
         }
     }
 
-    // Fetch the update-work-item max once; used to skip excess entries without aborting the batch
-    let update_wi_config: UpdateWorkItemConfig = ctx.get_tool_config("update-work-item");
-    let max_update_wi = update_wi_config.max as usize;
-    let mut update_wi_executed: usize = 0;
-
-    // Fetch the comment-on-work-item max once; same skip-and-continue pattern
-    let comment_wi_config: CommentOnWorkItemConfig = ctx.get_tool_config("comment-on-work-item");
-    let max_comment_wi = comment_wi_config.max as usize;
-    let mut comment_wi_executed: usize = 0;
+    // Build budget map: tool_name → (executed_count, max_allowed).
+    // All safe-output tools that perform side-effects are budgeted. The `max` field
+    // is extracted generically from each tool's config JSON (defaulting to 1).
+    let budgeted_tools = [
+        "create-work-item",
+        "create-pull-request",
+        "update-work-item",
+        "comment-on-work-item",
+        "create-wiki-page",
+        "update-wiki-page",
+    ];
+    let mut budgets: HashMap<&str, (usize, usize)> = HashMap::new();
+    for tool_name in &budgeted_tools {
+        let max_config: MaxConfig = ctx.get_tool_config(tool_name);
+        budgets.insert(tool_name, (0, max_config.max as usize));
+    }
 
     let mut results = Vec::new();
     for (i, entry) in entries.iter().enumerate() {
@@ -107,35 +116,18 @@ pub async fn execute_safe_outputs(
             entry_json
         );
 
-        // Enforce update-work-item max: skip excess entries rather than aborting the whole batch.
+        // Generic budget enforcement: skip excess entries rather than aborting the whole batch.
         // Budget is consumed before execution so that failed attempts (target policy rejection,
         // network errors) still count — this prevents unbounded retries against a failing endpoint.
-        if entry.get("name").and_then(|n| n.as_str()) == Some("update-work-item") {
-            let wi_id = entry
-                .get("id")
-                .and_then(|v| v.as_u64())
-                .map(|id| format!(" (work item #{})", id))
-                .unwrap_or_default();
-            if let Some(result) = check_budget(entries.len(), i, "update-work-item", &wi_id, update_wi_executed, max_update_wi) {
-                results.push(result);
-                continue;
-            }
-            update_wi_executed += 1;
-        }
-
-        // Enforce comment-on-work-item max: same skip-and-continue pattern as update-work-item.
-        // Budget is consumed before execution so that failed attempts still count.
-        if entry.get("name").and_then(|n| n.as_str()) == Some("comment-on-work-item") {
-            let wi_id = entry
-                .get("work_item_id")
-                .and_then(|v| v.as_i64())
-                .map(|id| format!(" (work item #{})", id))
-                .unwrap_or_default();
-            if let Some(result) = check_budget(entries.len(), i, "comment-on-work-item", &wi_id, comment_wi_executed, max_comment_wi) {
-                results.push(result);
-                continue;
+        if let Some(tool_name) = entry.get("name").and_then(|n| n.as_str()) {
+            if let Some((executed, max)) = budgets.get_mut(tool_name) {
+                let context_id = extract_entry_context(entry);
+                if let Some(result) = check_budget(entries.len(), i, tool_name, &context_id, *executed, *max) {
+                    results.push(result);
+                    continue;
+                }
+                *executed += 1;
             }
-            comment_wi_executed += 1;
         }
 
         match execute_safe_output(entry, ctx).await {
@@ -284,6 +276,42 @@ pub async fn execute_safe_output(
     Ok((tool_name.to_string(), result))
 }
 
+/// Helper struct for extracting the `max` field from any tool's config JSON.
+/// All safe-output tool configs use `max` with a default of 1.
+#[derive(Deserialize)]
+struct MaxConfig {
+    #[serde(default = "default_max")]
+    max: u32,
+}
+
+fn default_max() -> u32 {
+    1
+}
+
+impl Default for MaxConfig {
+    fn default() -> Self {
+        Self { max: default_max() }
+    }
+}
+
+/// Extract a human-readable context identifier from a safe-output entry for log messages.
+fn extract_entry_context(entry: &Value) -> String {
+    if let Some(id) = entry.get("id").and_then(|v| v.as_u64()) {
+        return format!(" (work item #{})", id);
+    }
+    if let Some(id) = entry.get("work_item_id").and_then(|v| v.as_i64()) {
+        return format!(" (work item #{})", id);
+    }
+    if let Some(title) = entry.get("title").and_then(|v| v.as_str()) {
+        let truncated = if title.len() > 40 { &title[..40] } else { title };
+        return format!(" (\"{}\")", truncated);
+    }
+    if let Some(path) = entry.get("path").and_then(|v| v.as_str()) {
+        return format!(" (path: {})", path);
+    }
+    String::new()
+}
+
 /// Returns `Some(result)` when the budget for `tool_name` is exhausted so the caller can push the
 /// result and `continue` to the next entry. Returns `None` when a budget slot is still available
 /// and the caller should proceed with execution.
@@ -735,4 +763,153 @@ mod tests {
         let r = result.unwrap();
         assert!(r.message.contains("(work item #42)"));
     }
+
+    // --- extract_entry_context unit tests ---
+
+    #[test]
+    fn test_extract_entry_context_with_id() {
+        let entry = serde_json::json!({"name": "update-work-item", "id": 42});
+        assert_eq!(extract_entry_context(&entry), " (work item #42)");
+    }
+
+    #[test]
+    fn test_extract_entry_context_with_work_item_id() {
+        let entry = serde_json::json!({"name": "comment-on-work-item", "work_item_id": 99});
+        assert_eq!(extract_entry_context(&entry), " (work item #99)");
+    }
+
+    #[test]
+    fn test_extract_entry_context_with_title() {
+        let entry = serde_json::json!({"name": "create-work-item", "title": "Fix the bug"});
+        assert_eq!(extract_entry_context(&entry), " (\"Fix the bug\")");
+    }
+
+    #[test]
+    fn test_extract_entry_context_with_path() {
+        let entry = serde_json::json!({"name": "create-wiki-page", "path": "/Overview/NewPage"});
+        assert_eq!(extract_entry_context(&entry), " (path: /Overview/NewPage)");
+    }
+
+    #[test]
+    fn test_extract_entry_context_empty() {
+        let entry = serde_json::json!({"name": "noop"});
+        assert_eq!(extract_entry_context(&entry), "");
+    }
+
+    // --- MaxConfig unit tests ---
+
+    #[test]
+    fn test_max_config_default() {
+        let config = MaxConfig::default();
+        assert_eq!(config.max, 1);
+    }
+
+    #[test]
+    fn test_max_config_from_json_with_max() {
+        let json = serde_json::json!({"max": 5, "other_field": true});
+        let config: MaxConfig = serde_json::from_value(json).unwrap();
+        assert_eq!(config.max, 5);
+    }
+
+    #[test]
+    fn test_max_config_from_json_without_max() {
+        let json = serde_json::json!({"other_field": true});
+        let config: MaxConfig = serde_json::from_value(json).unwrap();
+        assert_eq!(config.max, 1);
+    }
+
+    // --- Generic budget enforcement for all tool types ---
+
+    #[tokio::test]
+    async fn test_budget_enforcement_create_work_item_max() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let safe_output_path = temp_dir.path().join(SAFE_OUTPUT_FILENAME);
+
+        // Write 3 create-work-item entries + 1 noop; max set to 2
+        let ndjson = r#"{"name":"create-work-item","title":"First item","description":"A description that is definitely longer than thirty characters."}
+{"name":"create-work-item","title":"Second item","description":"A description that is definitely longer than thirty characters."}
+{"name":"create-work-item","title":"Third item","description":"A description that is definitely longer than thirty characters."}
+{"name":"noop","context":"still runs"}
+"#;
+        tokio::fs::write(&safe_output_path, ndjson).await.unwrap();
+
+        let mut tool_configs = HashMap::new();
+        tool_configs.insert("create-work-item".to_string(), serde_json::json!({"max": 2}));
+
+        let ctx = ExecutionContext {
+            ado_org_url: Some("https://dev.azure.com/org".to_string()),
+            ado_organization: Some("org".to_string()),
+            ado_project: Some("Proj".to_string()),
+            access_token: Some("token".to_string()),
+            working_directory: PathBuf::from("."),
+            source_directory: PathBuf::from("."),
+            tool_configs,
+            repository_id: None,
+            repository_name: None,
+            allowed_repositories: HashMap::new(),
+        };
+
+        let results = execute_safe_outputs(temp_dir.path(), &ctx).await;
+        assert!(results.is_ok(), "Batch should not abort when max is exceeded");
+        let results = results.unwrap();
+        assert_eq!(results.len(), 4, "Expected 4 results");
+
+        // Only 1 should be skipped (max=2 allows first 2, third is skipped)
+        let skipped: Vec<_> = results
+            .iter()
+            .filter(|r| r.message.contains("maximum create-work-item count"))
+            .collect();
+        assert_eq!(skipped.len(), 1, "Expected 1 skipped entry, got: {:?}", skipped);
+
+        // noop still runs
+        assert!(results[3].success, "noop should still succeed");
+    }
+
+    #[tokio::test]
+    async fn test_budget_enforcement_mixed_tools_independent_budgets() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let safe_output_path = temp_dir.path().join(SAFE_OUTPUT_FILENAME);
+
+        // Mix of tools: each has max=1 (default), so only the first of each type should pass budget
+        let ndjson = r#"{"name":"create-work-item","title":"WI 1","description":"A description that is definitely longer than thirty characters."}
+{"name":"create-work-item","title":"WI 2","description":"A description that is definitely longer than thirty characters."}
+{"name":"create-wiki-page","path":"/Page1","content":"Some valid wiki content here."}
+{"name":"create-wiki-page","path":"/Page2","content":"Some valid wiki content here."}
+{"name":"noop","context":"always runs"}
+"#;
+        tokio::fs::write(&safe_output_path, ndjson).await.unwrap();
+
+        let ctx = ExecutionContext {
+            ado_org_url: Some("https://dev.azure.com/org".to_string()),
+            ado_organization: Some("org".to_string()),
+            ado_project: Some("Proj".to_string()),
+            access_token: Some("token".to_string()),
+            working_directory: PathBuf::from("."),
+            source_directory: PathBuf::from("."),
+            tool_configs: HashMap::new(), // defaults: max=1 for all
+            repository_id: None,
+            repository_name: None,
+            allowed_repositories: HashMap::new(),
+        };
+
+        let results = execute_safe_outputs(temp_dir.path(), &ctx).await.unwrap();
+        assert_eq!(results.len(), 5);
+
+        // Second create-work-item should be skipped
+        let cwi_skipped: Vec<_> = results
+            .iter()
+            .filter(|r| r.message.contains("maximum create-work-item count"))
+            .collect();
+        assert_eq!(cwi_skipped.len(), 1, "Expected 1 skipped create-work-item");
+
+        // Second create-wiki-page should be skipped
+        let cwp_skipped: Vec<_> = results
+            .iter()
+            .filter(|r| r.message.contains("maximum create-wiki-page count"))
+            .collect();
+        assert_eq!(cwp_skipped.len(), 1, "Expected 1 skipped create-wiki-page");
+
+        // noop always runs
+        assert!(results[4].success, "noop should still succeed");
+    }
 }

src/tools/create_pr.rs

@@ -261,6 +261,10 @@ pub struct CreatePrConfig {
     /// Work item IDs to link to the PR
     #[serde(default, rename = "work-items")]
     pub work_items: Vec<i32>,
+
+    /// Maximum number of create-pull-request outputs allowed per pipeline run (default: 1)
+    #[serde(default = "default_max")]
+    pub max: u32,
 }
 
 fn default_target_branch() -> String {
@@ -271,6 +275,10 @@ fn default_true() -> bool {
     true
 }
 
+fn default_max() -> u32 {
+    1
+}
+
 impl Default for CreatePrConfig {
     fn default() -> Self {
         Self {
@@ -281,6 +289,7 @@ impl Default for CreatePrConfig {
             reviewers: Vec::new(),
             labels: Vec::new(),
             work_items: Vec::new(),
+            max: default_max(),
         }
     }
 }
@@ -1119,6 +1128,7 @@ mod tests {
         assert!(!config.auto_complete);
         assert!(config.delete_source_branch);
         assert!(config.squash_merge);
+        assert_eq!(config.max, 1);
     }
 
     #[test]

src/tools/create_wiki_page.rs

@@ -123,6 +123,14 @@ pub struct CreateWikiPageConfig {
     /// Default commit comment used when the agent does not supply one.
     #[serde(default)]
     pub comment: Option<String>,
+
+    /// Maximum number of create-wiki-page outputs allowed per pipeline run (default: 1)
+    #[serde(default = "default_max")]
+    pub max: u32,
+}
+
+fn default_max() -> u32 {
+    1
 }
 
 impl Default for CreateWikiPageConfig {
@@ -133,6 +141,7 @@ impl Default for CreateWikiPageConfig {
             path_prefix: None,
             title_prefix: None,
             comment: None,
+            max: default_max(),
         }
     }
 }
@@ -510,6 +519,7 @@ mod tests {
         assert!(config.path_prefix.is_none());
         assert!(config.title_prefix.is_none());
         assert!(config.comment.is_none());
+        assert_eq!(config.max, 1);
     }
 
     #[test]