feat: consolidated streaming and timeout improvements (#90)

* feat: consolidated streaming and timeout improvements

## Summary

This PR consolidates **3 feature PRs** for streaming and timeout improvements.

### Included PRs:
- #25: Add per-chunk timeout for SSE streaming to prevent hangs
- #29: Add early JSON validation for tool call arguments
- #34: Add per-tool timeout in batch execution

### Key Changes:
- Added CHUNK_TIMEOUT_SECS constant for SSE streaming (60s)
- Wrapped SSE event iteration with tokio::time::timeout
- Added validate_arguments() method to StreamToolCall
- Added is_valid_complete() helper and complete_tool_call_validated() method
- Added DEFAULT_TOOL_TIMEOUT_SECS constant (60 seconds) for batch
- Added tool_timeout_secs field to BatchToolArgs for configuration
- Applied individual timeout to each tool execution in execute_parallel()

### Files Modified:
- src/cortex-engine/src/client/cortex.rs
- src/cortex-engine/src/streaming.rs
- src/cortex-engine/src/tools/handlers/batch.rs
- src/cortex-engine/src/tools/unified_executor.rs

Closes #25, #29, #34

* 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

Author: echobt

src/cortex-engine/src/client/cortex.rs

@@ -5,12 +5,15 @@
 //! - Responses API (streaming SSE)
 //! - Credit system with price verification
 
+use std::time::Duration;
+
 use async_trait::async_trait;
 use eventsource_stream::Eventsource;
 use futures::StreamExt;
 use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
+use tokio::time::timeout;
 use tokio_stream::wrappers::ReceiverStream;
 
 use super::{
@@ -22,6 +25,11 @@ use crate::error::{CortexError, Result};
 
 const DEFAULT_CORTEX_URL: &str = "https://api.cortex.foundation";
 
+/// Timeout in seconds for receiving individual SSE chunks during streaming.
+/// If no data is received within this duration, the connection is terminated
+/// to prevent indefinite hangs when connections stall mid-stream.
+const CHUNK_TIMEOUT_SECS: u64 = 60;
+
 /// Pricing information for a model.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PricingInfo {
@@ -567,8 +575,26 @@ impl ModelClient for CortexClient {
             let mut stream = std::pin::pin!(stream);
             let mut accumulated_text = String::new();
             let mut usage = TokenUsage::default();
-
-            while let Some(event_result) = stream.next().await {
+            let chunk_timeout = Duration::from_secs(CHUNK_TIMEOUT_SECS);
+
+            loop {
+                // Apply per-chunk timeout to prevent indefinite hangs when connections stall
+                let event_result = match timeout(chunk_timeout, stream.next()).await {
+                    Ok(Some(result)) => result,
+                    Ok(None) => break, // Stream ended normally
+                    Err(_) => {
+                        // Timeout elapsed - no data received within CHUNK_TIMEOUT_SECS
+                        let _ = tx
+                            .send(Err(CortexError::BackendError {
+                                message: format!(
+                                    "SSE chunk timeout - no data received for {} seconds",
+                                    CHUNK_TIMEOUT_SECS
+                                ),
+                            }))
+                            .await;
+                        break;
+                    }
+                };
                 match event_result {
                     Ok(event) => {
                         if event.data.is_empty() || event.data == "[DONE]" {

src/cortex-engine/src/streaming.rs

@@ -143,6 +143,18 @@ impl StreamContent {
         }
     }
 
+    /// Complete a tool call and validate its arguments.
+    /// Returns Ok(()) if the tool call was found and arguments are valid JSON.
+    /// Returns Err if tool call not found or arguments are invalid JSON.
+    pub fn complete_tool_call_validated(&mut self, id: &str) -> Result<(), String> {
+        if let Some(tc) = self.tool_calls.iter_mut().find(|tc| tc.id == id) {
+            tc.complete = true;
+            tc.validate_arguments()
+        } else {
+            Err(format!("Tool call with id '{}' not found", id))
+        }
+    }
+
     /// Check if has content.
     pub fn has_content(&self) -> bool {
         !self.text.is_empty() || !self.tool_calls.is_empty()
@@ -171,6 +183,23 @@ impl StreamToolCall {
             None
         }
     }
+
+    /// Validate that arguments contain valid JSON.
+    /// Returns Ok(()) if valid, Err with details if invalid.
+    pub fn validate_arguments(&self) -> Result<(), String> {
+        if self.arguments.trim().is_empty() {
+            return Ok(()); // Empty is valid (no arguments)
+        }
+        serde_json::from_str::<serde_json::Value>(&self.arguments)
+            .map(|_| ())
+            .map_err(|e| format!("Invalid JSON in tool call arguments: {}", e))
+    }
+
+    /// Check if arguments are complete and valid JSON.
+    /// Returns true only if complete and valid.
+    pub fn is_valid_complete(&self) -> bool {
+        self.complete && self.validate_arguments().is_ok()
+    }
 }
 
 /// Token counts.

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

@@ -23,6 +23,10 @@ pub const MAX_BATCH_SIZE: usize = 10;
 /// Default timeout for batch execution in seconds.
 pub const DEFAULT_BATCH_TIMEOUT_SECS: u64 = 300;
 
+/// Default timeout for individual tool execution in seconds.
+/// This prevents a single tool from consuming the entire batch timeout.
+pub const DEFAULT_TOOL_TIMEOUT_SECS: u64 = 60;
+
 /// Tools that cannot be called within a batch (prevent recursion).
 /// Note: Task is now allowed to enable parallel task execution.
 pub const DISALLOWED_TOOLS: &[&str] = &["Batch", "batch", "Agent", "agent"];
@@ -45,6 +49,10 @@ pub struct BatchToolArgs {
     /// Optional timeout in seconds for the entire batch (default: 300s).
     #[serde(default)]
     pub timeout_secs: Option<u64>,
+    /// Optional timeout in seconds for individual tools (default: 60s).
+    /// This prevents a single tool from consuming the entire batch timeout.
+    #[serde(default)]
+    pub tool_timeout_secs: Option<u64>,
 }
 
 /// Result of a single tool call within the batch.
@@ -158,7 +166,7 @@ impl BatchToolHandler {
         &self,
         calls: Vec<BatchToolCall>,
         context: &ToolContext,
-        timeout_duration: Duration,
+        tool_timeout: Duration,
     ) -> BatchResult {
         let start_time = Instant::now();
         let results = Arc::new(Mutex::new(Vec::with_capacity(calls.len())));
@@ -176,9 +184,9 @@ impl BatchToolHandler {
                 async move {
                     let call_start = Instant::now();
 
-                    // Execute with per-call timeout (use batch timeout for each call)
+                    // Execute with per-tool timeout to prevent single tools from blocking others
                     let execution_result = timeout(
-                        timeout_duration,
+                        tool_timeout,
                         executor.execute_tool(&call.tool, call.arguments, &ctx),
                     )
                     .await;
@@ -202,19 +210,23 @@ impl BatchToolHandler {
                             duration_ms,
                         },
                         Ok(Err(e)) => BatchCallResult {
-                            tool: tool_name,
+                            tool: tool_name.clone(),
                             index,
                             success: false,
                             result: None,
                             error: Some(format!("Execution error: {}", e)),
                             duration_ms,
                         },
                         Err(_) => BatchCallResult {
-                            tool: tool_name,
+                            tool: tool_name.clone(),
                             index,
                             success: false,
                             result: None,
-                            error: Some(format!("Timeout after {}s", timeout_duration.as_secs())),
+                            error: Some(format!(
+                                "Tool '{}' timed out after {}s",
+                                tool_name,
+                                tool_timeout.as_secs()
+                            )),
                             duration_ms,
                         },
                     };
@@ -328,14 +340,30 @@ impl ToolHandler for BatchToolHandler {
         // Validate calls
         self.validate_calls(&args.calls)?;
 
-        // Determine timeout
-        let timeout_secs = args.timeout_secs.unwrap_or(DEFAULT_BATCH_TIMEOUT_SECS);
-        let timeout_duration = Duration::from_secs(timeout_secs);
-
-        // Execute all tools in parallel
-        let batch_result = self
-            .execute_parallel(args.calls, context, timeout_duration)
-            .await;
+        // Determine overall batch timeout (wraps around entire parallel execution)
+        let batch_timeout_secs = args.timeout_secs.unwrap_or(DEFAULT_BATCH_TIMEOUT_SECS);
+        let batch_timeout = Duration::from_secs(batch_timeout_secs);
+
+        // Determine per-tool timeout (prevents single tool from blocking others)
+        let tool_timeout_secs = args.tool_timeout_secs.unwrap_or(DEFAULT_TOOL_TIMEOUT_SECS);
+        let tool_timeout = Duration::from_secs(tool_timeout_secs);
+
+        // Execute all tools in parallel with overall batch timeout
+        let batch_result = match timeout(
+            batch_timeout,
+            self.execute_parallel(args.calls, context, tool_timeout),
+        )
+        .await
+        {
+            Ok(result) => result,
+            Err(_) => {
+                // Batch-level timeout exceeded
+                return Ok(ToolResult::error(format!(
+                    "Batch execution timed out after {}s. Consider using a longer timeout_secs or reducing the number of tools.",
+                    batch_timeout_secs
+                )));
+            }
+        };
 
         // Format output
         let output = self.format_result(&batch_result);
@@ -384,6 +412,12 @@ pub fn batch_tool_definition() -> ToolDefinition {
                     "description": "Optional timeout in seconds for the entire batch execution (default: 300)",
                     "minimum": 1,
                     "maximum": 600
+                },
+                "tool_timeout_secs": {
+                    "type": "integer",
+                    "description": "Optional timeout in seconds for individual tool execution (default: 60). Prevents a single tool from consuming the entire batch timeout.",
+                    "minimum": 1,
+                    "maximum": 300
                 }
             }
         }),
@@ -409,6 +443,7 @@ pub async fn execute_batch(
             })
             .collect(),
         timeout_secs: None,
+        tool_timeout_secs: None,
     };
 
     let arguments = serde_json::to_value(args)
@@ -652,4 +687,58 @@ mod tests {
             elapsed.as_millis()
         );
     }
+
+    #[tokio::test]
+    async fn test_batch_timeout() {
+        // Create an executor with a slow tool
+        struct SlowExecutor;
+
+        #[async_trait]
+        impl BatchToolExecutor for SlowExecutor {
+            async fn execute_tool(
+                &self,
+                _name: &str,
+                _arguments: Value,
+                _context: &ToolContext,
+            ) -> Result<ToolResult> {
+                // Sleep longer than batch timeout
+                tokio::time::sleep(Duration::from_secs(5)).await;
+                Ok(ToolResult::success("Done"))
+            }
+
+            fn has_tool(&self, _name: &str) -> bool {
+                true
+            }
+        }
+
+        let executor: Arc<dyn BatchToolExecutor> = Arc::new(SlowExecutor);
+        let handler = BatchToolHandler::new(executor);
+        let context = ToolContext::new(PathBuf::from("."));
+
+        // Use a very short batch timeout (1 second) to test timeout behavior
+        let args = json!({
+            "calls": [
+                {"tool": "SlowTool", "arguments": {}}
+            ],
+            "timeout_secs": 1
+        });
+
+        let start = Instant::now();
+        let result = handler.execute(args, &context).await;
+        let elapsed = start.elapsed();
+
+        assert!(result.is_ok());
+        let tool_result = result.unwrap();
+
+        // Should timeout quickly (around 1 second)
+        assert!(
+            elapsed.as_secs() < 3,
+            "Batch should have timed out quickly, but took {}s",
+            elapsed.as_secs()
+        );
+
+        // Should have timed out
+        assert!(!tool_result.success);
+        assert!(tool_result.output.contains("timed out"));
+    }
 }

src/cortex-engine/src/tools/unified_executor.rs

@@ -466,6 +466,7 @@ impl UnifiedToolExecutor {
                 Ok(BatchToolArgs {
                     calls,
                     timeout_secs: arguments.get("timeout_secs").and_then(|v| v.as_u64()),
+                    tool_timeout_secs: arguments.get("tool_timeout_secs").and_then(|v| v.as_u64()),
                 })
             } else {
                 Err(CortexError::InvalidInput(