Skip to content

Commit aab0d99

Browse files
committed
fix(batch): implement timeout_secs parameter for overall batch timeout
Address Greptile review feedback: The timeout_secs parameter was documented and accepted but never used. Now it properly wraps the entire parallel execution with a batch-level timeout, separate from the per-tool timeout_secs. - Add batch-level timeout wrapper around execute_parallel - Return descriptive error message when batch times out - Add test for batch timeout behavior
1 parent 9c20808 commit aab0d99

File tree

1 file changed

+74
-4
lines changed
  • src/cortex-engine/src/tools/handlers

1 file changed

+74
-4
lines changed

src/cortex-engine/src/tools/handlers/batch.rs

Lines changed: 74 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)