agnt5 0.2.8a3__cp310-abi3-macosx_11_0_arm64.whl → 0.2.8a4__cp310-abi3-macosx_11_0_arm64.whl

agnt5/__init__.py

@@ -7,7 +7,6 @@ with built-in durability guarantees and state management.
 
 from ._compat import _import_error, _rust_available
 from .agent import Agent, AgentContext, AgentRegistry, AgentResult, Handoff, agent, handoff
-from .agent_session import AgentSession
 from .client import Client, RunError
 from .context import Context
 from .function import FunctionContext
@@ -15,7 +14,7 @@ from .workflow import WorkflowContext
 from .entity import (
     Entity,
     EntityRegistry,
-    EntityStateManager,
+    EntityStateAdapter,
     EntityType,
     create_entity_context,
     with_entity_context,
@@ -55,7 +54,7 @@ __all__ = [
     "Entity",
     "EntityType",
     "EntityRegistry",
-    "EntityStateManager",
+    "EntityStateAdapter",
     "with_entity_context",
     "create_entity_context",
     "workflow",
@@ -67,7 +66,6 @@ __all__ = [
     "Agent",
     "AgentRegistry",
     "AgentResult",
-    "AgentSession",
     "Handoff",
     "handoff",
     # Types

agnt5/agent.py

@@ -79,42 +79,46 @@ class AgentContext(Context):
         self._agent_name = agent_name
         self._session_id = session_id or run_id
 
-        # Determine state manager based on parent context
-        from .entity import EntityStateManager, _get_state_manager
+        # Determine state adapter based on parent context
+        from .entity import EntityStateAdapter, _get_state_adapter
 
         if state_manager:
-            # Explicit state manager provided
-            self._state_manager = state_manager
-            logger.debug(f"AgentContext using provided state manager")
+            # Explicit state adapter provided (parameter name kept for backward compat)
+            self._state_adapter = state_manager
+            logger.debug(f"AgentContext using provided state adapter")
         elif parent_context:
-            # Try to inherit state manager from parent
+            # Try to inherit state adapter from parent
             try:
                 # Check if parent is WorkflowContext or AgentContext
                 if hasattr(parent_context, '_workflow_entity'):
-                    # WorkflowContext - get state manager from worker context
-                    self._state_manager = _get_state_manager()
+                    # WorkflowContext - get state adapter from worker context
+                    self._state_adapter = _get_state_adapter()
                     logger.debug(f"AgentContext inheriting state from WorkflowContext")
-                elif hasattr(parent_context, '_state_manager'):
-                    # Parent AgentContext - share state manager
-                    self._state_manager = parent_context._state_manager
+                elif hasattr(parent_context, '_state_adapter'):
+                    # Parent AgentContext - share state adapter
+                    self._state_adapter = parent_context._state_adapter
                     logger.debug(f"AgentContext inheriting state from parent AgentContext")
+                elif hasattr(parent_context, '_state_manager'):
+                    # Backward compatibility: parent has old _state_manager
+                    self._state_adapter = parent_context._state_manager
+                    logger.debug(f"AgentContext inheriting state from parent (legacy)")
                 else:
-                    # FunctionContext or base Context - create new state manager
-                    self._state_manager = EntityStateManager()
-                    logger.debug(f"AgentContext created new state manager (parent has no state)")
+                    # FunctionContext or base Context - create new state adapter
+                    self._state_adapter = EntityStateAdapter()
+                    logger.debug(f"AgentContext created new state adapter (parent has no state)")
             except RuntimeError as e:
-                # _get_state_manager() failed (not in worker context) - create standalone
-                self._state_manager = EntityStateManager()
-                logger.debug(f"AgentContext created standalone state manager (not in worker context)")
+                # _get_state_adapter() failed (not in worker context) - create standalone
+                self._state_adapter = EntityStateAdapter()
+                logger.debug(f"AgentContext created standalone state adapter (not in worker context)")
         else:
             # Try to get from worker context first
             try:
-                self._state_manager = _get_state_manager()
-                logger.debug(f"AgentContext got state manager from worker context")
+                self._state_adapter = _get_state_adapter()
+                logger.debug(f"AgentContext got state adapter from worker context")
             except RuntimeError as e:
-                # Standalone - create new state manager
-                self._state_manager = EntityStateManager()
-                logger.debug(f"AgentContext created standalone state manager")
+                # Standalone - create new state adapter
+                self._state_adapter = EntityStateAdapter()
+                logger.debug(f"AgentContext created standalone state adapter")
 
         # Conversation key for state storage (used for in-memory state)
         self._conversation_key = f"agent:{agent_name}:{self._session_id}:messages"
@@ -127,24 +131,23 @@ class AgentContext(Context):
         """
         Get state interface for agent state management.
 
+        Note: This is a simplified in-memory state interface for agent-specific data.
+        Conversation history is managed separately via get_conversation_history() and
+        save_conversation_history() which use the Rust-backed persistence layer.
+
         Returns:
-            EntityState instance for state operations
+            Dict-like object for state operations
 
         Example:
-            # Store conversation history
-            messages = ctx.state.get(f"agent:{agent_name}:{session_id}:messages", [])
-            messages.append({"role": "user", "content": "Hello"})
-            ctx.state.set(f"agent:{agent_name}:{session_id}:messages", messages)
-
-            # Store agent-specific data
-            ctx.state.set("research_results", data)
+            # Store agent-specific data (in-memory only)
+            ctx.state["research_results"] = data
+            ctx.state["iteration_count"] = 5
         """
-        from .entity import EntityState
-
-        # Use agent's conversation key as the state key
-        state_key = ("agent", self._conversation_key)
-        state_dict = self._state_manager.get_or_create_state(state_key)
-        return EntityState(state_dict)
+        # Simple dict-based state for agent-specific data
+        # This is in-memory only and not persisted to platform
+        if not hasattr(self, '_agent_state'):
+            self._agent_state = {}
+        return self._agent_state
 
     @property
     def session_id(self) -> str:
@@ -155,45 +158,30 @@ class AgentContext(Context):
         """
         Retrieve conversation history from state, loading from database if needed.
 
+        Uses the EntityStateAdapter which delegates to Rust core for cache-first loading.
+
         Returns:
             List of Message objects from conversation history
         """
-        # Try to load from database first if not in memory
         entity_type = "AgentSession"
-        entity_key = self._entity_key  # Use entity_key without :messages suffix
-        state_key = (entity_type, entity_key)
-
-        # Check if we need to load from platform
-        if state_key not in self._state_manager._states and self._state_manager._rust_manager:
-            try:
-                result = await self._state_manager._rust_manager.py_load_state(
-                    entity_type, entity_key
-                )
-                found, state_json_bytes, version = result
-                if found:
-                    import json
-                    state_json = state_json_bytes.decode('utf-8') if isinstance(state_json_bytes, bytes) else state_json_bytes
-                    session_data = json.loads(state_json)
-
-                    # Extract messages from session object (might be old format or new format)
-                    if isinstance(session_data, dict) and "messages" in session_data:
-                        # New format with session metadata
-                        messages_data = session_data["messages"]
-                    elif isinstance(session_data, list):
-                        # Old format - just messages array
-                        messages_data = session_data
-                    else:
-                        messages_data = []
-
-                    # Load messages into in-memory state (store just messages, not full session object)
-                    self._state_manager.load_state_from_platform(state_key, json.dumps(messages_data), version)
-                    logger.info(f"Loaded conversation history from database: {entity_key} (version {version})")
-            except Exception as e:
-                logger.warning(f"Failed to load conversation history from database: {e}")
-
-        # Get from in-memory state
-        messages_data = self.state.get(self._conversation_key, [])
-        logger.debug(f"Loaded {len(messages_data)} messages from conversation history")
+        entity_key = self._entity_key
+
+        # Load session data via adapter (Rust handles cache + platform load)
+        session_data = await self._state_adapter.load_state(entity_type, entity_key)
+
+        # Extract messages from session object
+        if isinstance(session_data, dict) and "messages" in session_data:
+            # New format with session metadata
+            messages_data = session_data["messages"]
+            logger.debug(f"Loaded {len(messages_data)} messages from session {entity_key}")
+        elif isinstance(session_data, list):
+            # Old format - just messages array
+            messages_data = session_data
+            logger.debug(f"Loaded {len(messages_data)} messages (legacy format)")
+        else:
+            # No messages found
+            messages_data = []
+            logger.debug(f"No conversation history found for {entity_key}")
 
         # Convert dict representations back to Message objects
         messages = []
@@ -221,6 +209,8 @@ class AgentContext(Context):
         """
         Save conversation history to state and persist to database.
 
+        Uses the EntityStateAdapter which delegates to Rust core for version-checked saves.
+
         Args:
             messages: List of Message objects to persist
         """
@@ -234,46 +224,126 @@ class AgentContext(Context):
                 "content": msg.content
             })
 
-        # Save to in-memory state
-        self.state.set(self._conversation_key, messages_data)
+        import time
+        entity_type = "AgentSession"
+        entity_key = self._entity_key
 
-        # Persist to database via Rust EntityStateManager
-        if self._state_manager._rust_manager:
-            try:
-                import json
-                import time
-                entity_type = "AgentSession"
-                entity_key = self._entity_key  # Use entity_key without :messages suffix
-                state_key = (entity_type, entity_key)
-
-                # Get current version
-                expected_version = self._state_manager._versions.get(state_key, 0)
-
-                # Build session object with metadata for the sessions API
-                now = time.time()
-                session_data = {
-                    "session_id": self._session_id,
-                    "agent_name": self._agent_name,
-                    "created_at": now if expected_version == 0 else None,  # Only set on first save
-                    "last_message_time": now,
-                    "message_count": len(messages_data),
-                    "messages": messages_data,
-                    "metadata": {}
-                }
+        # Load current state with version for optimistic locking
+        current_state, current_version = await self._state_adapter.load_with_version(
+            entity_type, entity_key
+        )
 
-                # Serialize to JSON bytes
-                state_json = json.dumps(session_data).encode('utf-8')
+        # Build session object with metadata
+        now = time.time()
 
-                # Save to platform
-                new_version = await self._state_manager._rust_manager.py_save_state(
-                    entity_type, entity_key, state_json, expected_version
-                )
+        # Get custom metadata from instance variable or preserve from loaded state
+        custom_metadata = getattr(self, '_custom_metadata', current_state.get("metadata", {}))
+
+        session_data = {
+            "session_id": self._session_id,
+            "agent_name": self._agent_name,
+            "created_at": current_state.get("created_at", now),  # Preserve existing or set new
+            "last_message_time": now,
+            "message_count": len(messages_data),
+            "messages": messages_data,
+            "metadata": custom_metadata  # Save custom metadata
+        }
+
+        # Save to platform via adapter (Rust handles optimistic locking)
+        try:
+            new_version = await self._state_adapter.save_state(
+                entity_type,
+                entity_key,
+                session_data,
+                current_version
+            )
+            logger.info(
+                f"Persisted conversation history: {entity_key} (version {current_version} -> {new_version})"
+            )
+        except Exception as e:
+            logger.error(f"Failed to persist conversation history to database: {e}")
+            # Don't fail - conversation is still in memory for this execution
+
+    async def get_metadata(self) -> Dict[str, Any]:
+        """
+        Get conversation session metadata.
+
+        Returns session metadata including:
+        - created_at: Timestamp of first message (float, Unix timestamp)
+        - last_activity: Timestamp of last message (float, Unix timestamp)
+        - message_count: Number of messages in conversation (int)
+        - custom: Dict of user-provided custom metadata
 
-                # Update version tracking
-                self._state_manager._versions[state_key] = new_version
-                logger.info(f"Persisted conversation history to database: {entity_key} (version {expected_version} -> {new_version})")
-            except Exception as e:
-                logger.error(f"Failed to persist conversation history to database: {e}")
+        Returns:
+            Dictionary with metadata. If no conversation exists yet, returns defaults.
+
+        Example:
+            ```python
+            metadata = await context.get_metadata()
+            print(f"Session created: {metadata['created_at']}")
+            print(f"User ID: {metadata['custom'].get('user_id')}")
+            ```
+        """
+        entity_type = "AgentSession"
+        entity_key = self._entity_key
+
+        # Load session data
+        session_data = await self._state_adapter.load_state(entity_type, entity_key)
+
+        if not session_data:
+            # No conversation exists yet - return defaults
+            return {
+                "created_at": None,
+                "last_activity": None,
+                "message_count": 0,
+                "custom": getattr(self, '_custom_metadata', {})
+            }
+
+        messages = session_data.get("messages", [])
+
+        # Derive timestamps from messages if available
+        created_at = session_data.get("created_at")
+        last_activity = session_data.get("last_message_time")
+
+        return {
+            "created_at": created_at,
+            "last_activity": last_activity,
+            "message_count": len(messages),
+            "custom": session_data.get("metadata", {})
+        }
+
+    def update_metadata(self, **kwargs) -> None:
+        """
+        Update custom session metadata.
+
+        Metadata will be persisted alongside conversation history on next save.
+        Use this to store application-specific data like user_id, preferences, etc.
+
+        Args:
+            **kwargs: Key-value pairs to store as metadata
+
+        Example:
+            ```python
+            # Store user identification and preferences
+            context.update_metadata(
+                user_id="user-123",
+                subscription_tier="premium",
+                preferences={"theme": "dark", "language": "en"}
+            )
+
+            # Later retrieve it
+            metadata = await context.get_metadata()
+            user_id = metadata["custom"]["user_id"]
+            ```
+
+        Note:
+            - Metadata is merged with existing metadata (doesn't replace)
+            - Changes persist on next save_conversation_history() call
+            - Use simple JSON-serializable types (str, int, float, dict, list)
+        """
+        if not hasattr(self, '_custom_metadata'):
+            self._custom_metadata = {}
+        self._custom_metadata.update(kwargs)
 
 
 class Handoff:

agnt5/entity.py

@@ -15,156 +15,245 @@ from ._telemetry import setup_module_logger
 
 logger = setup_module_logger(__name__)
 
-# Context variable for worker-scoped state manager
+# Context variable for worker-scoped state adapter
 # This is set by Worker before entity execution and accessed by Entity instances
-_entity_state_manager_ctx: contextvars.ContextVar[Optional["EntityStateManager"]] = \
-    contextvars.ContextVar('_entity_state_manager', default=None)
+_entity_state_adapter_ctx: contextvars.ContextVar[Optional["EntityStateAdapter"]] = \
+    contextvars.ContextVar('_entity_state_adapter', default=None)
 
 # Global entity registry
 _ENTITY_REGISTRY: Dict[str, "EntityType"] = {}
 
 
-class EntityStateManager:
+class EntityStateAdapter:
     """
-    Worker-scoped state and lock management for entities.
-
-    This class provides isolated state management per Worker instance,
-    replacing the global dict approach. Each Worker gets its own state manager,
-    which provides:
-    - State storage per entity (type, key)
-    - Single-writer locks per entity
-    - Version tracking for optimistic locking
-    - Platform state loading/saving via Rust EntityStateManager
+    Thin Python adapter providing Pythonic interface to Rust EntityStateManager core.
+
+    This adapter provides language-specific concerns only:
+    - Worker-local asyncio.Lock for coarse-grained coordination
+    - Type conversions between Python dict and JSON bytes
+    - Pythonic async/await API over Rust core
+
+    All business logic (caching, version tracking, retry logic, gRPC) lives in the Rust core.
+    This keeps the Python layer simple (~150 LOC) and enables sharing business logic across SDKs.
     """
 
-    def __init__(self, rust_entity_state_manager=None):
+    def __init__(self, rust_core=None):
         """
-        Initialize empty state manager.
+        Initialize entity state adapter.
 
         Args:
-            rust_entity_state_manager: Optional Rust EntityStateManager for gRPC communication.
-                                      TODO: Wire this up once PyO3 bindings are complete.
+            rust_core: Rust EntityStateManager instance (from _core module).
+                      If None, operates in standalone/testing mode with in-memory state.
         """
-        self._states: Dict[Tuple[str, str], Dict[str, Any]] = {}
-        self._locks: Dict[Tuple[str, str], asyncio.Lock] = {}
-        self._versions: Dict[Tuple[str, str], int] = {}
-        self._rust_manager = rust_entity_state_manager  # TODO: Use for load/save
-        logger.debug("Created EntityStateManager")
-
-    def get_or_create_state(self, state_key: Tuple[str, str]) -> Dict[str, Any]:
+        self._rust_core = rust_core
+        # Worker-local locks for coarse-grained coordination within this worker
+        self._local_locks: Dict[Tuple[str, str], asyncio.Lock] = {}
+
+        # Standalone mode: in-memory state storage when no Rust core
+        # This enables testing without the full platform stack
+        if rust_core is None:
+            self._standalone_states: Dict[Tuple[str, str], Dict[str, Any]] = {}
+            self._standalone_versions: Dict[Tuple[str, str], int] = {}
+            logger.debug("Created EntityStateAdapter in standalone mode (in-memory state)")
+        else:
+            logger.debug("Created EntityStateAdapter with Rust core")
+
+    def get_local_lock(self, state_key: Tuple[str, str]) -> asyncio.Lock:
         """
-        Get or create state dict for entity instance.
+        Get worker-local asyncio.Lock for single-writer guarantee within this worker.
+
+        This provides coarse-grained coordination for operations within the same worker.
+        Cross-worker conflicts are handled by the Rust core via optimistic concurrency.
 
         Args:
             state_key: Tuple of (entity_type, entity_key)
 
         Returns:
-            State dict for the entity instance
+            asyncio.Lock for this worker-local operation
         """
-        if state_key not in self._states:
-            self._states[state_key] = {}
-        return self._states[state_key]
+        if state_key not in self._local_locks:
+            self._local_locks[state_key] = asyncio.Lock()
+        return self._local_locks[state_key]
 
-    def get_or_create_lock(self, state_key: Tuple[str, str]) -> asyncio.Lock:
+    async def load_state(self, entity_type: str, entity_key: str) -> Dict[str, Any]:
         """
-        Get or create async lock for entity instance.
+        Load entity state (Rust handles cache-first logic and platform load).
+
+        In standalone mode (no Rust core), uses in-memory state storage.
 
         Args:
-            state_key: Tuple of (entity_type, entity_key)
+            entity_type: Type of entity (e.g., "ShoppingCart", "Counter")
+            entity_key: Unique key for entity instance
 
         Returns:
-            Async lock for single-writer guarantee
+            State dictionary (empty dict if not found)
         """
-        if state_key not in self._locks:
-            self._locks[state_key] = asyncio.Lock()
-        return self._locks[state_key]
+        if not self._rust_core:
+            # Standalone mode - return from in-memory storage
+            state_key = (entity_type, entity_key)
+            return self._standalone_states.get(state_key, {}).copy()
 
-    def load_state_from_platform(
+        try:
+            # Rust checks cache first, loads from platform if needed
+            state_json_bytes, version = await self._rust_core.py_get_cached_or_load(entity_type, entity_key)
+
+            # Convert bytes to dict
+            if state_json_bytes:
+                state_json = state_json_bytes.decode('utf-8') if isinstance(state_json_bytes, bytes) else state_json_bytes
+                return json.loads(state_json)
+            else:
+                return {}
+        except Exception as e:
+            logger.warning(f"Failed to load state for {entity_type}:{entity_key}: {e}")
+            return {}
+
+    async def save_state(
         self,
-        state_key: Tuple[str, str],
-        platform_state_json: str,
-        version: int = 0
-    ) -> None:
+        entity_type: str,
+        entity_key: str,
+        state: Dict[str, Any],
+        expected_version: int
+    ) -> int:
         """
-        Load state from platform for entity persistence.
+        Save entity state (Rust handles version check and platform save).
+
+        In standalone mode (no Rust core), stores in-memory with version tracking.
 
         Args:
-            state_key: Tuple of (entity_type, entity_key)
-            platform_state_json: JSON string of state from platform
-            version: Current version from platform
+            entity_type: Type of entity
+            entity_key: Unique key for entity instance
+            state: State dictionary to save
+            expected_version: Expected current version (for optimistic locking)
+
+        Returns:
+            New version number after save
+
+        Raises:
+            RuntimeError: If version conflict or platform error
         """
-        import json
-        try:
-            state = json.loads(platform_state_json)
-            self._states[state_key] = state
-            self._versions[state_key] = version
-            logger.debug(
-                f"Loaded platform state: {state_key[0]}/{state_key[1]} (version {version})"
-            )
-        except json.JSONDecodeError as e:
-            logger.warning(f"Failed to parse platform state: {e}")
-            self._states[state_key] = {}
-            self._versions[state_key] = 0
+        if not self._rust_core:
+            # Standalone mode - store in memory with version tracking
+            state_key = (entity_type, entity_key)
+            current_version = self._standalone_versions.get(state_key, 0)
+
+            # Optimistic locking check (even in standalone mode for consistency)
+            if current_version != expected_version:
+                raise RuntimeError(
+                    f"Version conflict: expected {expected_version}, got {current_version}"
+                )
 
-    def get_state_for_persistence(
-        self,
-        state_key: Tuple[str, str]
-    ) -> tuple[Dict[str, Any], int, int]:
+            # Store state and increment version
+            new_version = expected_version + 1
+            self._standalone_states[state_key] = state.copy()
+            self._standalone_versions[state_key] = new_version
+            return new_version
+
+        # Convert dict to JSON bytes
+        state_json = json.dumps(state).encode('utf-8')
+
+        # Rust handles optimistic locking and platform save
+        new_version = await self._rust_core.py_save_state(
+            entity_type,
+            entity_key,
+            state_json,
+            expected_version
+        )
+
+        return new_version
+
+    async def load_with_version(self, entity_type: str, entity_key: str) -> Tuple[Dict[str, Any], int]:
         """
-        Get state and version info for platform persistence.
+        Load entity state with version (for update operations).
+
+        In standalone mode (no Rust core), loads from in-memory storage with version.
 
         Args:
-            state_key: Tuple of (entity_type, entity_key)
+            entity_type: Type of entity
+            entity_key: Unique key for entity instance
 
         Returns:
-            Tuple of (state_dict, expected_version, new_version)
+            Tuple of (state_dict, version)
+        """
+        if not self._rust_core:
+            # Standalone mode - return from in-memory storage with version
+            state_key = (entity_type, entity_key)
+            state = self._standalone_states.get(state_key, {}).copy()
+            version = self._standalone_versions.get(state_key, 0)
+            return state, version
+
+        try:
+            state_json_bytes, version = await self._rust_core.py_get_cached_or_load(entity_type, entity_key)
+
+            if state_json_bytes:
+                state_json = state_json_bytes.decode('utf-8') if isinstance(state_json_bytes, bytes) else state_json_bytes
+                state = json.loads(state_json)
+            else:
+                state = {}
+
+            return state, version
+        except Exception as e:
+            logger.warning(f"Failed to load state with version for {entity_type}:{entity_key}: {e}")
+            return {}, 0
+
+    async def invalidate_cache(self, entity_type: str, entity_key: str) -> None:
         """
-        state_dict = self._states.get(state_key, {})
-        expected_version = self._versions.get(state_key, 0)
-        new_version = expected_version + 1
+        Invalidate cache entry for specific entity.
 
-        # Update version for next execution
-        self._versions[state_key] = new_version
+        Args:
+            entity_type: Type of entity
+            entity_key: Unique key for entity instance
+        """
+        if self._rust_core:
+            await self._rust_core.py_invalidate_cache(entity_type, entity_key)
 
-        return state_dict, expected_version, new_version
+    async def clear_cache(self) -> None:
+        """Clear entire cache (useful for testing)."""
+        if self._rust_core:
+            await self._rust_core.py_clear_cache()
 
     def clear_all(self) -> None:
-        """Clear all state, locks, and versions (for testing)."""
-        self._states.clear()
-        self._locks.clear()
-        self._versions.clear()
-        logger.debug("Cleared EntityStateManager")
+        """Clear all local locks (for testing)."""
+        self._local_locks.clear()
+        logger.debug("Cleared EntityStateAdapter local locks")
 
-    def get_state(self, entity_type: str, key: str) -> Optional[Dict[str, Any]]:
+    async def get_state(self, entity_type: str, key: str) -> Optional[Dict[str, Any]]:
         """Get state for debugging/testing."""
-        state_key = (entity_type, key)
-        return self._states.get(state_key)
+        state, _ = await self.load_with_version(entity_type, key)
+        return state if state else None
 
     def get_all_keys(self, entity_type: str) -> list[str]:
-        """Get all keys for entity type (for debugging/testing)."""
-        return [
-            key for (etype, key) in self._states.keys()
-            if etype == entity_type
-        ]
+        """
+        Get all keys for an entity type (testing/debugging only).
+
+        Only works in standalone mode. Returns empty list in production mode.
+        """
+        if not hasattr(self, '_standalone_states'):
+            return []
+
+        keys = []
+        for (etype, ekey) in self._standalone_states.keys():
+            if etype == entity_type:
+                keys.append(ekey)
+        return keys
 
 
-def _get_state_manager() -> EntityStateManager:
+def _get_state_adapter() -> EntityStateAdapter:
     """
-    Get the current entity state manager from context.
+    Get the current entity state adapter from context.
 
-    The state manager must be set by Worker before entity execution.
+    The state adapter must be set by Worker before entity execution.
     This ensures proper worker-scoped state isolation.
 
     Returns:
-        EntityStateManager instance
+        EntityStateAdapter instance
 
     Raises:
-        RuntimeError: If called outside of Worker context (state manager not set)
+        RuntimeError: If called outside of Worker context (state adapter not set)
     """
-    manager = _entity_state_manager_ctx.get()
-    if manager is None:
+    adapter = _entity_state_adapter_ctx.get()
+    if adapter is None:
         raise RuntimeError(
-            "Entity requires state manager context.\n\n"
+            "Entity requires state adapter context.\n\n"
             "In production:\n"
             "  Entities run automatically through Worker.\n\n"
             "In tests, use one of:\n"
@@ -179,7 +268,9 @@ def _get_state_manager() -> EntityStateManager:
             "        await cart.add_item(...)\n\n"
             "See: https://docs.agnt5.dev/sdk/entities#testing"
         )
-    return manager
+    return adapter
+
+
 
 
 # ============================================================================
@@ -188,7 +279,7 @@ def _get_state_manager() -> EntityStateManager:
 
 def with_entity_context(func):
     """
-    Decorator that sets up entity state manager for tests.
+    Decorator that sets up entity state adapter for tests.
 
     Usage:
         @with_entity_context
@@ -199,13 +290,13 @@ def with_entity_context(func):
     """
     @functools.wraps(func)
     async def wrapper(*args, **kwargs):
-        manager = EntityStateManager()
-        token = _entity_state_manager_ctx.set(manager)
+        adapter = EntityStateAdapter()
+        token = _entity_state_adapter_ctx.set(adapter)
         try:
             return await func(*args, **kwargs)
         finally:
-            _entity_state_manager_ctx.reset(token)
-            manager.clear_all()
+            _entity_state_adapter_ctx.reset(token)
+            adapter.clear_all()
     return wrapper
 
 
@@ -219,16 +310,16 @@ def create_entity_context():
 
         @pytest.fixture
         def entity_context():
-            manager, token = create_entity_context()
-            yield manager
+            adapter, token = create_entity_context()
+            yield adapter
             # Cleanup happens automatically
 
     Returns:
-        Tuple of (EntityStateManager, context_token)
+        Tuple of (EntityStateAdapter, context_token)
     """
-    manager = EntityStateManager()
-    token = _entity_state_manager_ctx.set(manager)
-    return manager, token
+    adapter = EntityStateAdapter()
+    token = _entity_state_adapter_ctx.set(adapter)
+    return adapter, token
 
 
 def extract_state_schema(entity_class: type) -> Optional[Dict[str, Any]]:
@@ -473,12 +564,12 @@ def _create_entity_method_wrapper(entity_type: str, method):
     """
     Create a wrapper for an entity method that provides single-writer consistency.
 
-    This wrapper:
-    1. Acquires a lock for the entity instance (single-writer guarantee)
-    2. Sets up EntityState with the state dict
-    3. Executes the method
-    4. Cleans up state reference
-    5. Handles errors appropriately
+    This wrapper implements hybrid locking:
+    1. Local lock (asyncio.Lock) for worker-scoped single-writer guarantee
+    2. Optimistic concurrency (via Rust) for cross-worker conflicts
+    3. Loads state via adapter (Rust handles cache + platform)
+    4. Executes the method with clean EntityState interface
+    5. Saves state via adapter (Rust handles version check + retry)
 
     Args:
         entity_type: Name of the entity type (class name)
@@ -489,36 +580,23 @@ def _create_entity_method_wrapper(entity_type: str, method):
     """
     @functools.wraps(method)
     async def entity_method_wrapper(self, *args, **kwargs):
-        """Execute entity method with single-writer guarantee."""
+        """Execute entity method with hybrid locking (local + optimistic)."""
         state_key = (entity_type, self._key)
 
-        # Get state manager and lock (single-writer guarantee)
-        state_manager = _get_state_manager()
-        lock = state_manager.get_or_create_lock(state_key)
+        # Get state adapter
+        adapter = _get_state_adapter()
+
+        # Local lock for worker-scoped single-writer guarantee
+        lock = adapter.get_local_lock(state_key)
 
         async with lock:
-            # Load state from platform if not in memory
-            if state_key not in state_manager._states and state_manager._rust_manager:
-                try:
-                    # Call Rust manager to load state from platform
-                    result = await state_manager._rust_manager.py_load_state(
-                        entity_type, self._key
-                    )
-                    # result is a tuple: (found, state_json, version)
-                    found, state_json_bytes, version = result
-                    if found:
-                        # Decode and load state
-                        state_json = state_json_bytes.decode('utf-8') if isinstance(state_json_bytes, bytes) else state_json_bytes
-                        state_manager.load_state_from_platform(
-                            state_key, state_json, version
-                        )
-                        logger.info(f"Loaded entity state from platform: {entity_type}:{self._key} (version {version})")
-                except Exception as e:
-                    logger.warning(f"Failed to load entity state from platform: {e}")
-                    # Continue with empty state
+            # Load state with version (Rust handles cache-first + platform load)
+            state_dict, current_version = await adapter.load_with_version(entity_type, self._key)
 
-            # Get or create state for this entity instance
-            state_dict = state_manager.get_or_create_state(state_key)
+            logger.debug(
+                "Loaded state for %s:%s (version %d)",
+                entity_type, self._key, current_version
+            )
 
             # Set up EntityState on instance for method access
             self._state = EntityState(state_dict)
@@ -529,25 +607,26 @@ def _create_entity_method_wrapper(entity_type: str, method):
                 result = await method(self, *args, **kwargs)
                 logger.debug("Completed %s:%s.%s", entity_type, self._key, method.__name__)
 
-                # Save state to platform after successful execution
-                if state_manager._rust_manager:
-                    try:
-                        import json
-                        # Get current version
-                        expected_version = state_manager._versions.get(state_key, 0)
-                        # Serialize state
-                        state_json = json.dumps(state_dict).encode('utf-8')
-                        # Call Rust manager to save state
-                        new_version = await state_manager._rust_manager.py_save_state(
-                            entity_type, self._key, state_json, expected_version
-                        )
-                        # Update version
-                        state_manager._versions[state_key] = new_version
-                        logger.info(f"Saved entity state to platform: {entity_type}:{self._key} (version {expected_version} -> {new_version})")
-                    except Exception as e:
-                        logger.error(f"Failed to save entity state to platform: {e}")
-                        # Don't fail the method execution just because persistence failed
-                        # The state is still in memory
+                # Save state after successful execution
+                # Rust handles optimistic locking (version check)
+                try:
+                    new_version = await adapter.save_state(
+                        entity_type,
+                        self._key,
+                        state_dict,
+                        current_version
+                    )
+                    logger.info(
+                        "Saved state for %s:%s (version %d -> %d)",
+                        entity_type, self._key, current_version, new_version
+                    )
+                except Exception as e:
+                    logger.error(
+                        "Failed to save state for %s:%s: %s",
+                        entity_type, self._key, e
+                    )
+                    # Don't fail the method execution just because persistence failed
+                    # The state is still in the local dict for this execution
 
                 return result
 

agnt5/worker.py

@@ -126,13 +126,16 @@ class Worker:
         self.runtime = runtime
         self.metadata = metadata or {}
 
+        # Get tenant_id from environment (required for entity state management)
+        import os
+        self._tenant_id = os.getenv("AGNT5_TENANT_ID", "default-tenant")
+
         # Import Rust worker
         try:
-            from ._core import PyWorker, PyWorkerConfig, PyComponentInfo, EntityStateManager as RustEntityStateManager
+            from ._core import PyWorker, PyWorkerConfig, PyComponentInfo
             self._PyWorker = PyWorker
             self._PyWorkerConfig = PyWorkerConfig
             self._PyComponentInfo = PyComponentInfo
-            self._RustEntityStateManager = RustEntityStateManager
         except ImportError as e:
             raise ImportError(
                 f"Failed to import Rust core worker: {e}. "
@@ -149,16 +152,17 @@ class Worker:
         # Create Rust worker instance
         self._rust_worker = self._PyWorker(self._rust_config)
 
-        # Get tenant_id for entity state manager
-        import os
-        tenant_id = os.getenv("AGNT5_TENANT_ID", "00000000-0000-0000-0000-000000000001")
+        # Create worker-scoped entity state adapter with Rust core
+        from .entity import EntityStateAdapter
+        from ._core import EntityStateManager as RustEntityStateManager
+
+        # Create Rust core for entity state management
+        rust_core = RustEntityStateManager(tenant_id=self._tenant_id)
 
-        # Create Rust entity state manager
-        self._rust_entity_state_manager = self._RustEntityStateManager(tenant_id)
+        # Create Python adapter (thin wrapper around Rust core)
+        self._entity_state_adapter = EntityStateAdapter(rust_core=rust_core)
 
-        # Create worker-scoped entity state manager with Rust manager
-        from .entity import EntityStateManager
-        self._entity_state_manager = EntityStateManager(rust_entity_state_manager=self._rust_entity_state_manager)
+        logger.info("Created EntityStateAdapter with Rust core for state management")
 
         # Component registration: auto-discover or explicit
         if auto_register:
@@ -754,7 +758,7 @@ class Worker:
         """Execute a workflow handler with automatic replay support."""
         import json
         from .workflow import WorkflowEntity, WorkflowContext
-        from .entity import _get_state_manager
+        from .entity import _get_state_adapter
         from ._core import PyExecuteComponentResponse
 
         try:
@@ -795,10 +799,15 @@ class Worker:
                 logger.debug(f"Loaded {len(completed_steps)} completed steps into workflow entity")
 
             if initial_state:
-                # Load initial state into entity's state manager
-                state_manager = _get_state_manager()
-                state_manager._states[workflow_entity._state_key] = initial_state
-                logger.debug(f"Loaded initial state with {len(initial_state)} keys into workflow entity")
+                # Load initial state into entity's state adapter
+                state_adapter = _get_state_adapter()
+                if hasattr(state_adapter, '_standalone_states'):
+                    # Standalone mode - set state directly
+                    state_adapter._standalone_states[workflow_entity._state_key] = initial_state
+                    logger.debug(f"Loaded initial state with {len(initial_state)} keys into workflow entity (standalone)")
+                else:
+                    # Production mode - state is managed by Rust core
+                    logger.debug(f"Initial state will be loaded from platform (production mode)")
 
             # Create WorkflowContext with entity and runtime_context for trace correlation
             ctx = WorkflowContext(
@@ -933,11 +942,11 @@ class Worker:
         """Execute an entity method."""
         import json
         from .context import Context
-        from .entity import EntityType, Entity, _entity_state_manager_ctx
+        from .entity import EntityType, Entity, _entity_state_adapter_ctx
         from ._core import PyExecuteComponentResponse
 
-        # Set entity state manager in context for Entity instances to access
-        _entity_state_manager_ctx.set(self._entity_state_manager)
+        # Set entity state adapter in context for Entity instances to access
+        _entity_state_adapter_ctx.set(self._entity_state_adapter)
 
         try:
             # Parse input data
@@ -952,23 +961,8 @@ class Worker:
             if not method_name:
                 raise ValueError("Entity invocation requires 'method' parameter")
 
-            # Load state from platform if provided in request metadata
-            state_key = (entity_type.name, entity_key)
-            if hasattr(request, 'metadata') and request.metadata:
-                if "entity_state" in request.metadata:
-                    platform_state_json = request.metadata["entity_state"]
-                    platform_version = int(request.metadata.get("state_version", "0"))
-
-                    # Load platform state into state manager
-                    self._entity_state_manager.load_state_from_platform(
-                        state_key,
-                        platform_state_json,
-                        platform_version
-                    )
-                    logger.info(
-                        f"Loaded entity state from platform: {entity_type.name}/{entity_key} "
-                        f"(version {platform_version})"
-                    )
+            # Note: State loading is now handled automatically by the entity method wrapper
+            # via EntityStateAdapter which uses the Rust core for cache + platform persistence
 
             # Create entity instance using the stored class reference
             entity_instance = entity_type.entity_class(key=entity_key)
@@ -979,32 +973,15 @@ class Worker:
 
             method = getattr(entity_instance, method_name)
 
-            # Execute method
+            # Execute method (entity method wrapper handles state load/save automatically)
             result = await method(**input_dict)
 
             # Serialize result
             output_data = json.dumps(result).encode("utf-8")
 
-            # Capture entity state after execution with version tracking
-            state_dict, expected_version, new_version = \
-                self._entity_state_manager.get_state_for_persistence(state_key)
-
+            # Note: State persistence is now handled automatically by the entity method wrapper
+            # via EntityStateAdapter which uses Rust core for optimistic locking + version tracking
             metadata = {}
-            if state_dict:
-                # Serialize state as JSON string for platform persistence
-                state_json = json.dumps(state_dict)
-                # Pass in metadata for Worker Coordinator to publish
-                metadata = {
-                    "entity_state": state_json,
-                    "entity_type": entity_type.name,
-                    "entity_key": entity_key,
-                    "expected_version": str(expected_version),
-                    "new_version": str(new_version),
-                }
-                logger.info(
-                    f"Captured entity state: {entity_type.name}/{entity_key} "
-                    f"(version {expected_version} → {new_version})"
-                )
 
             return PyExecuteComponentResponse(
                 invocation_id=request.invocation_id,
@@ -1039,11 +1016,11 @@ class Worker:
         import json
         import uuid
         from .agent import AgentContext
-        from .entity import _entity_state_manager_ctx
+        from .entity import _entity_state_adapter_ctx
         from ._core import PyExecuteComponentResponse
 
-        # Set entity state manager in context so AgentContext can access it
-        _entity_state_manager_ctx.set(self._entity_state_manager)
+        # Set entity state adapter in context so AgentContext can access it
+        _entity_state_adapter_ctx.set(self._entity_state_adapter)
 
         try:
             # Parse input data
@@ -1157,9 +1134,12 @@ class Worker:
         if self.metadata:
             self._rust_worker.set_service_metadata(self.metadata)
 
-        # Set entity state manager on Rust worker for database persistence
-        logger.info("Configuring entity state manager for database persistence")
-        self._rust_worker.set_entity_state_manager(self._rust_entity_state_manager)
+        # Configure entity state manager on Rust worker for database persistence
+        logger.info("Configuring Rust EntityStateManager for database persistence")
+        # Access the Rust core from the adapter
+        if hasattr(self._entity_state_adapter, '_rust_core') and self._entity_state_adapter._rust_core:
+            self._rust_worker.set_entity_state_manager(self._entity_state_adapter._rust_core)
+            logger.info("Successfully configured Rust EntityStateManager")
 
         # Get the current event loop to pass to Rust for concurrent Python async execution
         # This allows Rust to execute Python async functions on the same event loop

agnt5/workflow.py

@@ -11,7 +11,7 @@ from typing import Any, Callable, Dict, Optional, TypeVar, cast
 
 from ._schema_utils import extract_function_metadata, extract_function_schemas
 from .context import Context
-from .entity import Entity, EntityState, _get_state_manager
+from .entity import Entity, EntityState, _get_state_adapter
 from .function import FunctionContext
 from .types import HandlerFunc, WorkflowConfig
 from ._telemetry import setup_module_logger
@@ -346,10 +346,8 @@ class WorkflowEntity(Entity):
         for debugging and replay of AI workflows.
         """
         if self._state is None:
-            # Get state dict from state manager
-            state_manager = _get_state_manager()
-            state_dict = state_manager.get_or_create_state(self._state_key)
-            self._state = WorkflowState(state_dict, self)
+            # Initialize with empty state dict - will be populated by entity system
+            self._state = WorkflowState({}, self)
         return self._state
 
 

{agnt5-0.2.8a3.dist-info → agnt5-0.2.8a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agnt5
-Version: 0.2.8a3
+Version: 0.2.8a4
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: Programming Language :: Python :: 3

{agnt5-0.2.8a3.dist-info → agnt5-0.2.8a4.dist-info}/RECORD

@@ -1,16 +1,15 @@
-agnt5-0.2.8a3.dist-info/METADATA,sha256=CSD_9YZ9vmaRf7sXcVSefdJ2YeIZnwQ2PZCohtvv5TY,996
-agnt5-0.2.8a3.dist-info/WHEEL,sha256=-lwEpi49KOTCcgx48T3fLSP8Dxynwa-iRMZNo-JZaqc,103
-agnt5/__init__.py,sha256=dlJ3rUXiWIZMd0LpGGhgaItkdYlcp3yuEufQb-Yqh9A,2105
+agnt5-0.2.8a4.dist-info/METADATA,sha256=Dp8yJo_dx5B79OM0ptGgR80TjCUbyovpSc4Pt57dfuY,996
+agnt5-0.2.8a4.dist-info/WHEEL,sha256=-lwEpi49KOTCcgx48T3fLSP8Dxynwa-iRMZNo-JZaqc,103
+agnt5/__init__.py,sha256=ACkK91EPdnv5tYip09QCZ9rfV4iBKzNjGfYVLJD1XGg,2045
 agnt5/_compat.py,sha256=BGuy3v5VDOHVa5f3Z-C22iMN19lAt0mPmXwF3qSSWxI,369
-agnt5/_core.abi3.so,sha256=KfV9XpsvyVQTwFxDRQwAYbz6M-3LnIEpqUxPWLHgpbM,12223872
+agnt5/_core.abi3.so,sha256=zjKPxbDlFR0otveH4mWRbL_O6hWq6vXDDFj-WY10xRs,12598432
 agnt5/_retry_utils.py,sha256=loHsWY5BR4wZy57IzcDEjQAy88DHVwVIr25Cn1d9GPA,5801
 agnt5/_schema_utils.py,sha256=MR67RW757T4Oq2Jqf4kB61H_b51zwaf3CLWELnkngRo,9572
 agnt5/_telemetry.py,sha256=bIY9AvBRjJBTHoBPbfR6X1OgaiUf-T0vCoi0_snsWXA,5957
-agnt5/agent.py,sha256=fLX8k7cJ5-OKF9OU0Ykd__62mGFrYuToGUK0UdpNtas,41171
-agnt5/agent_session.py,sha256=mZ61mVax6Wt1ArocwdVEx4Kju6veFYtqxtSNWeKOve0,3161
+agnt5/agent.py,sha256=aBrhtPaUAHOHv3-h_Yb2UMqFHertr1P2hJ7fA_4IXcw,43225
 agnt5/client.py,sha256=kXksazgxdVXWaG9OkjJA4cWruNtcS-ENhtnkrIdw-Nk,23212
 agnt5/context.py,sha256=S2OzPkhn_jnqSWfT21mSYOux8vHaLKQxcAvggZDHQek,2378
-agnt5/entity.py,sha256=KXSDklsqe9PoIKGz4hZpcmHJFLJKGe0Am8-3FH72Fxk,26150
+agnt5/entity.py,sha256=AlHmSHVxQD5EYBvkmERKUkwv0ERrKaT8rvRK611hv_I,28941
 agnt5/exceptions.py,sha256=mZ0q-NK6OKhYxgwBJpIbgpgzk-CJaFIHDbp1EE-pS7I,925
 agnt5/function.py,sha256=f1vaAlJRwuo8cxCOGEd8XPido00mOhlPS8UJJx-6hJI,11041
 agnt5/lm.py,sha256=9dFjd6eQ3f3lFZe7H7rWZherYiP_58MT1F5xpwD8PCg,23195
@@ -18,6 +17,6 @@ agnt5/tool.py,sha256=uc4L-Q9QyLzQDe-MZKk2Wo3o5e-mK8tfaQwVDgQdouQ,13133
 agnt5/tracing.py,sha256=Mh2-OfnQM61lM_P8gxJstafdsUA8Gxoo1lP-Joxhub8,5980
 agnt5/types.py,sha256=Zb71ZMwvrt1p4SH18cAKunp2y5tao_W5_jGYaPDejQo,2840
 agnt5/version.py,sha256=rOq1mObLihnnKgKqBrwZA0zwOPudEKVFcW1a48ynkqc,573
-agnt5/worker.py,sha256=ALX8g8ij4GZ465mgEa17CO9I7Ybysv8w9uJWQRSd3g8,48685
-agnt5/workflow.py,sha256=sU8Gk7unxE_Gla7Fe-KlXfcBvYa2326GciuoR26CCr0,19585
-agnt5-0.2.8a3.dist-info/RECORD,,
+agnt5/worker.py,sha256=NflbueeL2LT8NGywTQnEv1r-N8f54AENWcZARJ5wO8o,47975
+agnt5/workflow.py,sha256=3s9CY6a4UkJZ9YyHv2SAkY3UeCVBlfVi7jxJMFi8Dhg,19488
+agnt5-0.2.8a4.dist-info/RECORD,,

agnt5/agent_session.py

@@ -1,110 +0,0 @@
-"""Agent session management for multi-turn conversations.
-
-Provides AgentSession entity to track and persist conversation metadata
-across multiple agent invocations. Session state is backed by the entity
-infrastructure for durability.
-"""
-
-from __future__ import annotations
-
-import time
-from typing import Optional, Dict, Any
-
-from .entity import Entity
-
-
-class AgentSession(Entity):
-    """
-    Entity representing an agent conversation session.
-
-    Tracks metadata about multi-turn conversations including:
-    - Session creation and last activity timestamps
-    - Message count
-    - Agent identification
-    - Custom metadata
-
-    The actual conversation messages are stored separately in AgentContext
-    using the session_id as the key. This entity only tracks session metadata.
-
-    Example:
-        ```python
-        # Create new session
-        session = AgentSession("session-123")
-        session.create(agent_name="tutor", metadata={"user_id": "user-456"})
-
-        # Update on each message
-        session.add_message()
-
-        # Get session summary
-        summary = session.get_summary()
-        ```
-    """
-
-    def __init__(self, session_id: str):
-        """
-        Initialize agent session entity.
-
-        Args:
-            session_id: Unique session identifier
-        """
-        self.session_id = session_id
-        self.agent_name: str = ""
-        self.created_at: float = 0.0
-        self.last_message_time: float = 0.0
-        self.message_count: int = 0
-        self.metadata: Dict[str, Any] = {}
-
-    def create(self, agent_name: str, metadata: Optional[Dict[str, Any]] = None) -> None:
-        """
-        Initialize a new session.
-
-        Args:
-            agent_name: Name of the agent for this session
-            metadata: Optional additional session metadata
-        """
-        current_time = time.time()
-        self.agent_name = agent_name
-        self.created_at = current_time
-        self.last_message_time = current_time
-        self.message_count = 0
-        self.metadata = metadata or {}
-
-    def add_message(self) -> None:
-        """
-        Record a new message in the session.
-
-        Updates message count and last activity timestamp.
-        """
-        self.message_count += 1
-        self.last_message_time = time.time()
-
-    def get_summary(self) -> Dict[str, Any]:
-        """
-        Get session metadata summary.
-
-        Returns:
-            Dictionary with session metadata including:
-            - session_id
-            - agent_name
-            - created_at
-            - last_message_time
-            - message_count
-            - metadata
-        """
-        return {
-            "session_id": self.session_id,
-            "agent_name": self.agent_name,
-            "created_at": self.created_at,
-            "last_message_time": self.last_message_time,
-            "message_count": self.message_count,
-            "metadata": self.metadata,
-        }
-
-    def update_metadata(self, metadata: Dict[str, Any]) -> None:
-        """
-        Update session metadata.
-
-        Args:
-            metadata: Metadata to merge into existing metadata
-        """
-        self.metadata.update(metadata)