Beyond Aggregate Risk: Role-Stratified Conformal Risk Control for LLM Tool Calls
2026-07-27 • Machine Learning
Machine LearningArtificial IntelligenceComputation and Language
AI summaryⓘ
The authors study how language-model agents use tools by filling in different fields that carry varying risks. They point out that current methods check the whole action at once, which can miss dangerous mistakes in small but important parts. To fix this, they introduce a new way to check each field separately based on its role, ensuring each part meets safety limits individually. Their approach better controls risks across different scenarios and changes in the environment. They conclude that safety checks should focus on each role in a tool call, not just the entire action.
language-model agentsstructured tool callsrisk controlconformal calibrationper-field detectorsemantic argument rolesfinite-sample guaranteedistribution shiftmodel robustness
Authors
Md Ashikur Rahman, Md Arifur Rahman, Niamul Hassan Samin, Khandaker Rifah Tasnia, Sifat Rahman Ahona, Juena Ahmed Noshin
Abstract
Language-model agents act through structured tool calls whose arguments carry different risks. Untrusted content may safely influence an email body but should not determine a recipient, account, command, or credential. Existing statistical methods typically control risk over the entire action, allowing failures in rare, high-risk fields to be obscured by benign arguments. We introduce role-stratified per-field conformal risk control, a calibration layer that wraps any per-field detector and sets separate thresholds and risk budgets for semantic argument roles. For a role with prevalence $p_r$, aggregate-only certification must use an effective budget of $αp_r$ to guarantee role-specific risk $α$, whereas role-stratified calibration certifies each sufficiently sampled role directly with a finite-sample guarantee; rarer roles are handled by pooled certification. Across AgentDojo and InjecAgent with six language models, the empirical utility gap tracks this predicted price of coarseness, and our method achieves the most consistent role-specific budget compliance under model and attack transfer, detector noise, gradual drift, unseen tool suites, and adaptive attacks. It provides formal per-role guarantees under exchangeability or after recalibration, and empirical compliance under frozen distribution shift. These results suggest that structured tool calls should be certified at the semantic-role level, not the whole action.