fix: Address 3 critical production robustness gaps

Fixes #1370

## Gap 1: Multi-Agent Concurrency Safety
- 1a: Added thread-safe cost tracking with _cost_lock in Agent class
- 1b: Fixed parallel workflow shared state by using copy.deepcopy() instead of shallow copy
- 1c: Confirmed SessionDeduplicationCache already has proper thread safety

## Gap 2: Session Isolation
- 2a: Fixed session key collision by including session_id in agent keys
- 2b: Added session TTL support with is_expired(), close(), time_to_expiry() methods
- 2c: Routed chat history through SessionStore first, memory as fallback

## Gap 3: Error Propagation
- 3a: Wrapped tool exceptions in ToolExecutionError for better observability
- 3b: Added exponential backoff and is_retryable checks to workflow retry logic
- 3c: Implemented cross-step handoff cycle detection with visited-set tracking

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: MervinPraison <MervinPraison@users.noreply.github.com>

Author: praisonai-triage-agent[bot]

src/praisonai-agents/praisonaiagents/agent/agent.py

@@ -1563,6 +1563,8 @@ def __init__(
         # Token budget guard (zero overhead when _max_budget is None)
         self._max_budget = _max_budget
         self._on_budget_exceeded = _on_budget_exceeded
+        # Thread-safe cost/token tracking (Gap 1a fix)
+        self._cost_lock = threading.Lock()
         self._total_cost = 0.0
         self._total_tokens_in = 0
         self._total_tokens_out = 0
@@ -1906,7 +1908,9 @@ def thinking_budget(self, value: Optional[int]) -> None:
     @property
     def total_cost(self) -> float:
         """Cumulative USD cost of all LLM calls in this agent run."""
-        return self._total_cost
+        # Thread-safe cost reading (Gap 1a fix)
+        with self._cost_lock:
+            return self._total_cost
 
     @property
     def cost_summary(self) -> dict:
@@ -1915,12 +1919,14 @@ def cost_summary(self) -> dict:
         Returns:
             dict with keys: tokens_in, tokens_out, cost, llm_calls
         """
-        return {
-            "tokens_in": self._total_tokens_in,
-            "tokens_out": self._total_tokens_out,
-            "cost": self._total_cost,
-            "llm_calls": self._llm_call_count,
-        }
+        # Thread-safe cost reading (Gap 1a fix)
+        with self._cost_lock:
+            return {
+                "tokens_in": self._total_tokens_in,
+                "tokens_out": self._total_tokens_out,
+                "cost": self._total_cost,
+                "llm_calls": self._llm_call_count,
+            }
 
     @property
     def context_manager(self) -> Optional[Any]:

src/praisonai-agents/praisonaiagents/agent/chat_mixin.py

@@ -677,26 +677,31 @@ def _chat_completion(self, messages, temperature=1.0, tools=None, stream=True, r
             )
 
             # Budget tracking & enforcement (zero overhead when _max_budget is None)
-            self._total_cost += _cost_usd
-            self._total_tokens_in += _prompt_tokens
-            self._total_tokens_out += _completion_tokens
-            self._llm_call_count += 1
-            if self._max_budget and self._total_cost >= self._max_budget:
+            # Thread-safe cost tracking (Gap 1a fix)
+            with self._cost_lock:
+                self._total_cost += _cost_usd
+                self._total_tokens_in += _prompt_tokens
+                self._total_tokens_out += _completion_tokens
+                self._llm_call_count += 1
+                budget_exceeded = self._max_budget and self._total_cost >= self._max_budget
+                current_cost = self._total_cost
+            
+            if budget_exceeded:
                 if self._on_budget_exceeded == "stop":
                     raise BudgetExceededError(
-                        f"Agent '{self.name}' exceeded budget: ${self._total_cost:.4f} >= ${self._max_budget:.4f}",
+                        f"Agent '{self.name}' exceeded budget: ${current_cost:.4f} >= ${self._max_budget:.4f}",
                         budget_type="cost",
                         limit=self._max_budget,
-                        used=self._total_cost,
+                        used=current_cost,
                         agent_id=self.name
                     )
                 elif self._on_budget_exceeded == "warn":
                     logging.warning(
-                        f"[budget] {self.name}: ${self._total_cost:.4f} exceeded "
+                        f"[budget] {self.name}: ${current_cost:.4f} exceeded "
                         f"${self._max_budget:.4f} budget"
                     )
                 elif callable(self._on_budget_exceeded):
-                    self._on_budget_exceeded(self._total_cost, self._max_budget)
+                    self._on_budget_exceeded(current_cost, self._max_budget)
 
             # Trigger AFTER_LLM hook
             from ..hooks import HookEvent, AfterLLMInput

src/praisonai-agents/praisonaiagents/agent/tool_execution.py

@@ -290,8 +290,15 @@ def _execute_tool_with_context(self, function_name, arguments, state, tool_call_
             _duration_ms = (_time.time() - _tool_start_time) * 1000
             _trace_emitter.tool_call_end(self.name, function_name, None, _duration_ms, str(e))
 
-            # Trigger OnError hook if needed (optional future step)
-            raise
+            # Gap 3a fix: Wrap exceptions in ToolExecutionError for better observability
+            from ..errors import ToolExecutionError
+            is_retryable = not isinstance(e, (ValueError, TypeError, AttributeError))
+            raise ToolExecutionError(
+                f"Tool '{function_name}' failed: {e}",
+                tool_name=function_name,
+                agent_id=self.name,
+                is_retryable=is_retryable,
+            ) from e
 
     def _trigger_after_agent_hook(self, prompt, response, start_time, tools_used=None):
         """Trigger AFTER_AGENT hook and return response."""

src/praisonai-agents/praisonaiagents/session.py

@@ -51,7 +51,8 @@ def __init__(
         agent_url: Optional[str] = None,
         memory_config: Optional[Dict[str, Any]] = None,
         knowledge_config: Optional[Dict[str, Any]] = None,
-        timeout: int = 30
+        timeout: int = 30,
+        session_ttl: Optional[int] = None  # Gap 2b: TTL in seconds
     ):
         """
         Initialize a new session with optional persistence or remote agent connectivity.
@@ -63,6 +64,7 @@ def __init__(
             memory_config: Configuration for memory system (defaults to RAG)
             knowledge_config: Configuration for knowledge base system  
             timeout: HTTP timeout for remote agent calls (default: 30 seconds)
+            session_ttl: Time-to-live in seconds after which session expires (Gap 2b)
         """
         self.session_id = session_id or str(uuid.uuid4())[:8]
         self.user_id = user_id or "default_user"
@@ -110,6 +112,10 @@ def __init__(
             self._agents_instance = None
             self._agents = {}  # Track agents and their chat histories
 
+        # Gap 2b: Session TTL and cleanup support  
+        self.session_ttl = session_ttl
+        self._created_at = time.time()  # Track creation time for TTL
+
     def _get_session_dir(self):
         """Return session-specific directory using paths.py."""
         from pathlib import Path
@@ -186,8 +192,8 @@ def Agent(
 
         agent = Agent(**agent_kwargs)
 
-        # Create a unique key for this agent (using name and role)
-        agent_key = f"{name}:{role}"
+        # Create a unique key for this agent (Gap 2a fix: include session_id for proper isolation)
+        agent_key = f"{self.session_id}:{name}:{role}"
 
         # Restore chat history if it exists from previous sessions
         if agent_key in self._agents:
@@ -270,7 +276,7 @@ def restore_state(self) -> Dict[str, Any]:
 
     def _restore_agent_chat_history(self, agent_key: str) -> List[Dict[str, Any]]:
         """
-        Restore agent chat history from memory.
+        Restore agent chat history from SessionStore first, then memory fallback (Gap 2c fix).
         
         Args:
             agent_key: Unique identifier for the agent
@@ -281,7 +287,18 @@ def _restore_agent_chat_history(self, agent_key: str) -> List[Dict[str, Any]]:
         if self.is_remote:
             return []
 
-        # Search for agent chat history in memory
+        # Gap 2c: Try SessionStore first for clean separation
+        try:
+            from .session.store import get_default_session_store
+            session_store = get_default_session_store()
+            session_id = f"{self.session_id}_{agent_key}"
+            chat_history = session_store.get_chat_history(session_id)
+            if chat_history:
+                return chat_history
+        except ImportError:
+            pass
+        
+        # Fallback: Search for agent chat history in memory (backward compatibility)
         results = self.memory.search_short_term(
             query="Agent chat history for",
             limit=10
@@ -350,10 +367,10 @@ def _save_agent_chat_histories(self) -> None:
                 chat_history = agent_data.get("chat_history")
 
             if chat_history is not None:
-                # G-2 FIX: Try SessionStore first for clean separation
+                # G-2 FIX: Use SessionStore for clean separation (Gap 2c fix)
                 session_store = None
                 try:
-                    from .session import get_default_session_store
+                    from .session.store import get_default_session_store
                     session_store = get_default_session_store()
                 except ImportError:
                     pass
@@ -580,6 +597,50 @@ def send_message(self, message: str, **kwargs) -> str:
         """
         return self.chat(message, **kwargs)
 
+    def is_expired(self) -> bool:
+        """Check if the session has expired based on TTL (Gap 2b fix)."""
+        if self.session_ttl is None:
+            return False
+        return time.time() - self._created_at > self.session_ttl
+
+    def close(self) -> None:
+        """Close and cleanup the session (Gap 2b fix)."""
+        if self.is_remote:
+            return  # No cleanup needed for remote sessions
+        
+        # Clear memory
+        if self._memory:
+            try:
+                # Clear short-term and long-term memory for this session
+                # Note: This is a basic implementation - specific memory backends
+                # might need more sophisticated cleanup
+                self._memory = None
+            except Exception:
+                pass  # Ignore cleanup errors
+        
+        # Clear knowledge
+        if self._knowledge:
+            try:
+                self._knowledge = None
+            except Exception:
+                pass  # Ignore cleanup errors
+        
+        # Clear agents
+        self._agents.clear()
+        
+        # Clean up session directory (optional - commented out for safety)
+        # import shutil
+        # session_dir = self._get_session_dir()
+        # if session_dir.exists():
+        #     shutil.rmtree(session_dir)
+
+    def time_to_expiry(self) -> Optional[float]:
+        """Get seconds until session expires, or None if no TTL set (Gap 2b)."""
+        if self.session_ttl is None:
+            return None
+        elapsed = time.time() - self._created_at
+        return max(0, self.session_ttl - elapsed)
+
     def __str__(self) -> str:
         if self.is_remote:
             return f"Session(id='{self.session_id}', user='{self.user_id}', remote_agent='{self.agent_url}')"

src/praisonai-agents/praisonaiagents/workflows/workflows.py

@@ -585,6 +585,8 @@ class AgentFlow:
     _reflection_config: Optional[Any] = field(default=None, repr=False)
     # Execution history for debugging (only populated when history=True)
     _execution_history: List[Dict[str, Any]] = field(default_factory=list, repr=False)
+    # Gap 3c: Cross-step handoff cycle detection
+    _handoff_chain: List[str] = field(default_factory=list, repr=False)
 
     def __post_init__(self):
         """Resolve consolidated params to internal values."""
@@ -899,6 +901,30 @@ def from_template(
                 "Install with: pip install praisonai"
             )
 
+    def _check_handoff_cycle(self, step: Any) -> None:
+        """Check for cross-step handoff cycles (Gap 3c fix)."""
+        step_id = getattr(step, 'name', str(step))
+        
+        # If step involves handoff (basic check for agent steps)
+        if hasattr(step, 'agent') and step.agent:
+            # Track this step in handoff chain
+            if step_id in self._handoff_chain:
+                # Cycle detected!
+                cycle_path = self._handoff_chain[self._handoff_chain.index(step_id):] + [step_id]
+                from ..errors import HandoffCycleError
+                raise HandoffCycleError(
+                    f"Cross-step handoff cycle detected: {' -> '.join(cycle_path)}",
+                    cycle_path=cycle_path
+                )
+            
+            # Add to chain (limit chain length to prevent memory issues)
+            self._handoff_chain.append(step_id)
+            if len(self._handoff_chain) > 100:  # Reasonable limit
+                self._handoff_chain = self._handoff_chain[-50:]  # Keep last 50
+        else:
+            # Non-agent step, reset chain 
+            self._handoff_chain.clear()
+    
     def run(
         self,
         input: str = "",
@@ -1096,6 +1122,9 @@ def run(
                 except Exception as e:
                     logger.error(f"should_run failed for {step.name}: {e}")
 
+            # Gap 3c: Check for cross-step handoff cycles
+            self._check_handoff_cycle(step)
+            
             # Execute step with retry and guardrail support
             output = None
             stop = False
@@ -1273,6 +1302,25 @@ def run(
                 except Exception as e:
                     step_error = e
                     output = f"Error: {e}"
+                    
+                    # Gap 3b fix: Check if error is retryable and implement exponential backoff
+                    is_retryable = getattr(e, 'is_retryable', True)  # Default to retryable
+                    if not is_retryable:
+                        # Non-retryable error - break out of retry loop immediately
+                        if verbose:
+                            print(f"❌ {step.name} failed with non-retryable error: {e}")
+                        break
+                    
+                    retry_count += 1
+                    if retry_count <= max_retries:
+                        # Exponential backoff: wait 2^(retry_count-1) seconds
+                        import time
+                        backoff_seconds = 2 ** (retry_count - 1)
+                        if verbose:
+                            print(f"🔄 {step.name} failed (attempt {retry_count}/{max_retries}), retrying in {backoff_seconds}s: {e}")
+                        time.sleep(backoff_seconds)
+                        continue  # Retry
+                    
                     if self.on_step_error:
                         try:
                             self.on_step_error(step.name, e)
@@ -2306,7 +2354,7 @@ def execute_with_branch(step=step, idx=idx, opt_prev=optimized_previous):
                     emitter.set_branch(f"parallel_{idx}")
                     try:
                         return self._execute_single_step_internal(
-                            step, opt_prev, input, all_variables.copy(), model, False, idx, stream, depth=depth+1
+                            step, opt_prev, input, copy.deepcopy(all_variables), model, False, idx, stream, depth=depth+1
                         )
                     finally:
                         emitter.clear_branch()