Architecture: Agent State Management and ContextVar Issues

[purvanshjoshi]

Description

The agent state management has several architectural issues that affect reliability and scalability:

1. ContextVar Usage Issues:

   • current_run_id = ContextVar('current_run_id') may not work correctly in async contexts
   • ContextVar values are not properly propagated across async boundaries
   • Potential race conditions in concurrent agent execution

2. State Management Problems:

   • AgentState TypedDict is mutable and shared across agents
   • No proper isolation between agent executions
   • State mutations can cause unexpected behavior

3. Error Handling Gaps:

   • No rollback mechanism for failed agent operations
   • No proper cleanup of partial state on errors
   • No timeout handling for long-running agent tasks

Proposed Fix

1. Improve ContextVar usage:
   `python
   from contextvars import ContextVar
   import asyncio

current_run_id: ContextVar[str] = ContextVar('current_run_id', default='')

async def set_run_id(run_id: str):
current_run_id.set(run_id)
# Ensure context is properly propagated
task = asyncio.current_task()
if task:
task.context = {**getattr(task, 'context', {}), 'run_id': run_id}
`

2. Implement proper state isolation:
   `python
   from dataclasses import dataclass, field
   from typing import Optional

@DataClass
class AgentState:
repo_name: str
target_issue: Optional[int] = None
architect_directive: str = ''
idea: str = ''
pm_decision: str = ''
code: str = ''
review: str = ''
issue_number: int = 0
pr_number: int = 0
branch_name: str = ''
iteration: int = 0
log_messages: list = field(default_factory=list)

def copy(self):
    return AgentState(**self.__dict__)

`

3. Add error handling and rollback:
   python async def safe_agent_execution(agent_func, state): try: return await agent_func(state) except Exception as e: # Log error and return safe state logger.error(f'Agent execution failed: {e}') return state.copy() finally: # Cleanup resources await cleanup_resources()

Acceptance Criteria

1. ContextVar is properly used in async contexts
2. Agent state is properly isolated between executions
3. Error handling includes rollback mechanisms
4. Timeout handling is implemented for all agents
5. No emojis in any codebase changes or commits

Impact

HIGH - These issues affect:

• System reliability and stability
• Concurrent execution safety
• Error recovery and debugging
• Production deployment readiness