use std::collections::{HashMap, HashSet};
use crate::cm_config::AgentConfig;
use crate::cm_tools::tool_dispatch::HandlerLookupTable;
use crate::cm_types::ToolCall;
use crate::cm_agent::plan_artifact::PlanStepExecutorKind;
use crate::cm_agent::step_executor_policy::{
executor_kind_tool_denied_body, tool_allowed_for_step_executor_kind,
};
use crate::cm_agent::turn_tool_policy::{
tool_allowed_for_turn, tool_calls_allow_parallel_for_role, turn_tool_denied_message,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExecuteToolsBatchOutcome {
Finished,
AbortedSse,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolBatchExecutionMode {
ParallelReadonlyBatch,
Serial,
}
pub struct ToolBatchModeParams<'a> {
pub force_serial: bool,
pub workspace_is_set: bool,
pub handler_lookup: &'a HandlerLookupTable,
pub cfg: &'a AgentConfig,
pub tool_calls: &'a [ToolCall],
pub turn_allow: Option<&'a HashSet<String>>,
}
pub fn replay_force_serial_from_env() -> bool {
std::env::var("CM_REPLAY_FORCE_SERIAL")
.ok()
.map(|v| v.trim().to_ascii_lowercase())
.is_some_and(|v| matches!(v.as_str(), "1" | "true" | "yes" | "on"))
}
pub type ParallelPrefetchFailureKey = (String, String);
pub fn resolve_tool_batch_execution_mode(
params: &ToolBatchModeParams<'_>,
) -> ToolBatchExecutionMode {
if params.force_serial || !params.workspace_is_set {
return ToolBatchExecutionMode::Serial;
}
if tool_calls_allow_parallel_for_role(
params.handler_lookup,
params.cfg,
params.tool_calls,
params.turn_allow,
) {
ToolBatchExecutionMode::ParallelReadonlyBatch
} else {
ToolBatchExecutionMode::Serial
}
}
pub fn dedup_readonly_tool_calls_count(tool_calls: &[ToolCall]) -> usize {
let mut seen: HashSet<(&str, &str)> = HashSet::with_capacity(tool_calls.len());
for tc in tool_calls {
seen.insert((tc.function.name.as_str(), tc.function.arguments.as_str()));
}
seen.len()
}
pub type ParallelPrefetchFailures = HashMap<ParallelPrefetchFailureKey, String>;
pub struct ToolPolicyEarlyDenyParams<'a> {
pub cfg: &'a AgentConfig,
pub name: &'a str,
pub step_executor_constraint: Option<PlanStepExecutorKind>,
pub tools_defs: &'a [crate::cm_types::Tool],
pub turn_allow: Option<&'a HashSet<String>>,
}
pub fn tool_policy_early_deny_message(p: &ToolPolicyEarlyDenyParams<'_>) -> Option<String> {
if let Some(k) = p.step_executor_constraint
&& !tool_allowed_for_step_executor_kind(p.cfg, p.name, k)
{
return Some(executor_kind_tool_denied_body(
p.cfg,
p.tools_defs,
p.name,
k,
));
}
if !tool_allowed_for_turn(p.name, p.turn_allow) {
return Some(turn_tool_denied_message(p.name));
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cm_config::load_config;
use crate::cm_types::{FunctionCall, ToolCall};
fn test_cfg() -> AgentConfig {
load_config(None).expect("embed default")
}
fn tc(name: &str, args: &str, id: &str) -> ToolCall {
ToolCall {
id: id.to_string(),
typ: "function".to_string(),
function: FunctionCall {
name: name.to_string(),
arguments: args.to_string(),
},
}
}
#[test]
fn dedup_counts_unique_name_args_pairs() {
let calls = vec![
tc("read_file", r#"{"path":"a"}"#, "1"),
tc("read_file", r#"{"path":"a"}"#, "2"),
tc("read_file", r#"{"path":"b"}"#, "3"),
];
assert_eq!(dedup_readonly_tool_calls_count(&calls), 2);
assert_eq!(dedup_readonly_tool_calls_count(&[]), 0);
}
#[test]
fn force_serial_always_serial_mode() {
let cfg = test_cfg();
let lookup = HandlerLookupTable::default_dispatch();
let calls = vec![tc("read_file", r#"{"path":"a"}"#, "1")];
let mode = resolve_tool_batch_execution_mode(&ToolBatchModeParams {
force_serial: true,
workspace_is_set: true,
handler_lookup: &lookup,
cfg: &cfg,
tool_calls: &calls,
turn_allow: None,
});
assert_eq!(mode, ToolBatchExecutionMode::Serial);
}
#[test]
fn early_deny_turn_allow_blocks_tool() {
let cfg = test_cfg();
let mut allow = HashSet::new();
allow.insert("read_file".to_string());
let msg = tool_policy_early_deny_message(&ToolPolicyEarlyDenyParams {
cfg: &cfg,
name: "run_command",
step_executor_constraint: None,
tools_defs: &[],
turn_allow: Some(&allow),
});
assert!(msg.is_some());
assert!(msg.unwrap().contains("run_command"));
}
}