Python: Refactor State class - rename, simplify, and implement superstep caching #3525
Description
Summary
The SharedState class should be renamed to State and refactored. Key questions around why it uses async methods for what is essentially a dictionary, and implementing proper superstep boundary semantics.
Files
| File | Line | Concern |
|---|---|---|
_shared_state.py |
26 | Could simply be a dictionary |
_shared_state.py |
27-28 | Local vs remote-backed shared state |
_shared_state.py |
29-31 | Commit semantics at superstep boundaries |
_shared_state.py |
32 | Document locking behavior |
_shared_state.py |
33 | Rename to just "State" |
_shared_state.py |
48 | No default value for get() |
Context
Current Implementation:
class SharedState:
"""A class to manage shared state in a workflow."""
def __init__(self) -> None:
self._state: dict[str, Any] = {}
self._shared_state_lock = asyncio.Lock()Key Questions:
-
Why async dict? Current implementation uses async methods (
async def get,async def set) with anasyncio.Lock. Why not just use a regular dict? Need to investigate if async is actually required or if this is over-engineering. -
Superstep Caching: State changes should be cached in a temporary dict during execution, then committed to the actual state dict at superstep boundaries. This ensures consistency and predictable behavior.
Get without default:
async def get(self, key: str) -> Any:
"""Get a value from the shared state."""
async with self._shared_state_lock:
return await self.get_within_hold(key)The get_within_hold raises KeyError if key not found, but there is no get(key, default) signature.
Action Items
- Rename: Rename
SharedStatetoState - Investigate async: Determine why we are using async methods for a dict - simplify to sync if not needed
- Superstep caching: Implement temp dict caching with commit at superstep boundaries:
- Cache state changes in interim dict during execution
- Commit cached changes to actual state at superstep boundary
- Distribution: Design interface for pluggable backends (in-memory, Redis, etc.)
- Documentation: Clarify locking behavior and superstep semantics
- API: Add
get(key, default=None)overload for convenience