@@ -340,14 +340,30 @@ impl ToolHandler for BatchToolHandler {
340340 // Validate calls
341341 self . validate_calls ( & args. calls ) ?;
342342
343+ // Determine overall batch timeout (wraps around entire parallel execution)
344+ let batch_timeout_secs = args. timeout_secs . unwrap_or ( DEFAULT_BATCH_TIMEOUT_SECS ) ;
345+ let batch_timeout = Duration :: from_secs ( batch_timeout_secs) ;
346+
343347 // Determine per-tool timeout (prevents single tool from blocking others)
344348 let tool_timeout_secs = args. tool_timeout_secs . unwrap_or ( DEFAULT_TOOL_TIMEOUT_SECS ) ;
345349 let tool_timeout = Duration :: from_secs ( tool_timeout_secs) ;
346350
347- // Execute all tools in parallel
348- let batch_result = self
349- . execute_parallel ( args. calls , context, tool_timeout)
350- . await ;
351+ // Execute all tools in parallel with overall batch timeout
352+ let batch_result = match timeout (
353+ batch_timeout,
354+ self . execute_parallel ( args. calls , context, tool_timeout) ,
355+ )
356+ . await
357+ {
358+ Ok ( result) => result,
359+ Err ( _) => {
360+ // Batch-level timeout exceeded
361+ return Ok ( ToolResult :: error ( format ! (
362+ "Batch execution timed out after {}s. Consider using a longer timeout_secs or reducing the number of tools." ,
363+ batch_timeout_secs
364+ ) ) ) ;
365+ }
366+ } ;
351367
352368 // Format output
353369 let output = self . format_result ( & batch_result) ;
@@ -671,4 +687,58 @@ mod tests {
671687 elapsed. as_millis( )
672688 ) ;
673689 }
690+
691+ #[ tokio:: test]
692+ async fn test_batch_timeout ( ) {
693+ // Create an executor with a slow tool
694+ struct SlowExecutor ;
695+
696+ #[ async_trait]
697+ impl BatchToolExecutor for SlowExecutor {
698+ async fn execute_tool (
699+ & self ,
700+ _name : & str ,
701+ _arguments : Value ,
702+ _context : & ToolContext ,
703+ ) -> Result < ToolResult > {
704+ // Sleep longer than batch timeout
705+ tokio:: time:: sleep ( Duration :: from_secs ( 5 ) ) . await ;
706+ Ok ( ToolResult :: success ( "Done" ) )
707+ }
708+
709+ fn has_tool ( & self , _name : & str ) -> bool {
710+ true
711+ }
712+ }
713+
714+ let executor: Arc < dyn BatchToolExecutor > = Arc :: new ( SlowExecutor ) ;
715+ let handler = BatchToolHandler :: new ( executor) ;
716+ let context = ToolContext :: new ( PathBuf :: from ( "." ) ) ;
717+
718+ // Use a very short batch timeout (1 second) to test timeout behavior
719+ let args = json ! ( {
720+ "calls" : [
721+ { "tool" : "SlowTool" , "arguments" : { } }
722+ ] ,
723+ "timeout_secs" : 1
724+ } ) ;
725+
726+ let start = Instant :: now ( ) ;
727+ let result = handler. execute ( args, & context) . await ;
728+ let elapsed = start. elapsed ( ) ;
729+
730+ assert ! ( result. is_ok( ) ) ;
731+ let tool_result = result. unwrap ( ) ;
732+
733+ // Should timeout quickly (around 1 second)
734+ assert ! (
735+ elapsed. as_secs( ) < 3 ,
736+ "Batch should have timed out quickly, but took {}s" ,
737+ elapsed. as_secs( )
738+ ) ;
739+
740+ // Should have timed out
741+ assert ! ( !tool_result. success) ;
742+ assert ! ( tool_result. output. contains( "timed out" ) ) ;
743+ }
674744}
0 commit comments