agnt5 0.2.8a2__cp310-abi3-manylinux_2_34_x86_64.whl → 0.2.8a4__cp310-abi3-manylinux_2_34_x86_64.whl

agnt5/__init__.py

@@ -14,7 +14,7 @@ from .workflow import WorkflowContext
 from .entity import (
     Entity,
     EntityRegistry,
-    EntityStateManager,
+    EntityStateAdapter,
     EntityType,
     create_entity_context,
     with_entity_context,
@@ -54,7 +54,7 @@ __all__ = [
     "Entity",
     "EntityType",
     "EntityRegistry",
-    "EntityStateManager",
+    "EntityStateAdapter",
     "with_entity_context",
     "create_entity_context",
     "workflow",

agnt5/agent.py

@@ -79,78 +79,109 @@ 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)")
-            except RuntimeError:
-                # _get_state_manager() failed (not in worker context) - create standalone
-                self._state_manager = EntityStateManager()
-                logger.debug(f"AgentContext created standalone state manager")
+                    # 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_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:
-            # Standalone - create new state manager
-            self._state_manager = EntityStateManager()
-            logger.debug(f"AgentContext created standalone state manager")
-
-        # Conversation key for state storage
+            # Try to get from worker context first
+            try:
+                self._state_adapter = _get_state_adapter()
+                logger.debug(f"AgentContext got state adapter from worker context")
+            except RuntimeError as e:
+                # 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"
+        # Entity key for database persistence (without :messages suffix to match API expectations)
+        self._entity_key = f"agent:{agent_name}:{self._session_id}"
+        logger.debug(f"AgentContext initialized - session_id={self._session_id}")
 
     @property
     def state(self):
         """
         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:
         """Get session identifier for this agent context."""
         return self._session_id
 
-    def get_conversation_history(self) -> List[Message]:
+    async def get_conversation_history(self) -> List[Message]:
         """
-        Retrieve conversation history from state.
+        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
         """
-        messages_data = self.state.get(self._conversation_key, [])
+        entity_type = "AgentSession"
+        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 = []
@@ -174,13 +205,17 @@ class AgentContext(Context):
 
         return messages
 
-    def save_conversation_history(self, messages: List[Message]) -> None:
+    async def save_conversation_history(self, messages: List[Message]) -> None:
         """
-        Save conversation history to state.
+        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
         """
+        logger.debug(f"Saving {len(messages)} messages to conversation history")
+
         # Convert Message objects to dict for JSON serialization
         messages_data = []
         for msg in messages:
@@ -189,7 +224,126 @@ class AgentContext(Context):
                 "content": msg.content
             })
 
-        self.state.set(self._conversation_key, messages_data)
+        import time
+        entity_type = "AgentSession"
+        entity_key = self._entity_key
+
+        # Load current state with version for optimistic locking
+        current_state, current_version = await self._state_adapter.load_with_version(
+            entity_type, entity_key
+        )
+
+        # Build session object with metadata
+        now = time.time()
+
+        # 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
+
+        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:
@@ -665,11 +819,11 @@ class Agent:
 
         # Load conversation history from state (if AgentContext)
         if isinstance(context, AgentContext):
-            messages: List[Message] = context.get_conversation_history()
+            messages: List[Message] = await context.get_conversation_history()
             # Add new user message
             messages.append(Message.user(user_message))
             # Save updated conversation
-            context.save_conversation_history(messages)
+            await context.save_conversation_history(messages)
         else:
             # Fallback for non-AgentContext (shouldn't happen with code above)
             messages = [Message.user(user_message)]
@@ -795,7 +949,7 @@ class Agent:
                                     )
                                     # Save conversation before returning
                                     if isinstance(context, AgentContext):
-                                        context.save_conversation_history(messages)
+                                        await context.save_conversation_history(messages)
                                     # Return immediately with handoff result
                                     return AgentResult(
                                         output=result["output"],
@@ -835,7 +989,7 @@ class Agent:
                     self.logger.debug(f"Agent completed after {iteration + 1} iterations")
                     # Save conversation before returning
                     if isinstance(context, AgentContext):
-                        context.save_conversation_history(messages)
+                        await context.save_conversation_history(messages)
                     return AgentResult(
                         output=response.text,
                         tool_calls=all_tool_calls,
@@ -847,7 +1001,7 @@ class Agent:
             final_output = messages[-1].content if messages else "No output generated"
             # Save conversation before returning
             if isinstance(context, AgentContext):
-                context.save_conversation_history(messages)
+                await context.save_conversation_history(messages)
             return AgentResult(
                 output=final_output,
                 tool_calls=all_tool_calls,

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,26 +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:
-            # TODO: Load state from platform if not in memory
-            # if state_key not in state_manager._states and state_manager._rust_manager:
-            #     result = await state_manager._rust_manager.load_state(
-            #         tenant_id, entity_type, self._key
-            #     )
-            #     if result.found:
-            #         state_manager.load_state_from_platform(
-            #             state_key, result.state_json, result.version
-            #         )
-
-            # Get or create state for this entity instance
-            state_dict = state_manager.get_or_create_state(state_key)
+            # Load state with version (Rust handles cache-first + platform load)
+            state_dict, current_version = await adapter.load_with_version(entity_type, self._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)
@@ -519,16 +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__)
 
-                # TODO: Save state to platform after successful execution
-                # if state_manager._rust_manager:
-                #     state_dict, expected_version, new_version = \
-                #         state_manager.get_state_for_persistence(state_key)
-                #     import json
-                #     state_json = json.dumps(state_dict).encode('utf-8')
-                #     save_result = await state_manager._rust_manager.save_state(
-                #         tenant_id, entity_type, self._key, state_json, expected_version
-                #     )
-                #     state_manager._versions[state_key] = save_result.new_version
+                # 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,6 +126,10 @@ 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
@@ -148,9 +152,17 @@ class Worker:
         # Create Rust worker instance
         self._rust_worker = self._PyWorker(self._rust_config)
 
-        # Create worker-scoped entity state manager
-        from .entity import EntityStateManager
-        self._entity_state_manager = EntityStateManager()
+        # 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 Python adapter (thin wrapper around Rust core)
+        self._entity_state_adapter = EntityStateAdapter(rust_core=rust_core)
+
+        logger.info("Created EntityStateAdapter with Rust core for state management")
 
         # Component registration: auto-discover or explicit
         if auto_register:
@@ -746,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:
@@ -787,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(
@@ -925,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
@@ -944,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)
@@ -971,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,
@@ -1027,11 +1012,16 @@ class Worker:
             )
 
     async def _execute_agent(self, agent, input_data: bytes, request):
-        """Execute an agent."""
+        """Execute an agent with session support for multi-turn conversations."""
         import json
-        from .context import Context
+        import uuid
+        from .agent import AgentContext
+        from .entity import _entity_state_adapter_ctx
         from ._core import PyExecuteComponentResponse
 
+        # Set entity state adapter in context so AgentContext can access it
+        _entity_state_adapter_ctx.set(self._entity_state_adapter)
+
         try:
             # Parse input data
             input_dict = json.loads(input_data.decode("utf-8")) if input_data else {}
@@ -1041,16 +1031,30 @@ class Worker:
             if not user_message:
                 raise ValueError("Agent invocation requires 'message' parameter")
 
-            # Create context with runtime_context for trace correlation
-            ctx = Context(
-                run_id=f"{self.service_name}:{agent.name}",
+            # Extract or generate session_id for multi-turn conversation support
+            # If session_id is provided, the agent will load previous conversation history
+            # If not provided, a new session is created with auto-generated ID
+            session_id = input_dict.get("session_id")
+
+            if not session_id:
+                session_id = str(uuid.uuid4())
+                logger.info(f"Created new agent session: {session_id}")
+            else:
+                logger.info(f"Using existing agent session: {session_id}")
+
+            # Create AgentContext with session support for conversation persistence
+            # AgentContext automatically loads/saves conversation history based on session_id
+            ctx = AgentContext(
+                run_id=request.invocation_id,
+                agent_name=agent.name,
+                session_id=session_id,
                 runtime_context=request.runtime_context,
             )
 
-            # Execute agent
+            # Execute agent - conversation history is automatically included
             agent_result = await agent.run(user_message, context=ctx)
 
-            # Build response
+            # Build response with agent output and tool calls
             result = {
                 "output": agent_result.output,
                 "tool_calls": agent_result.tool_calls,
@@ -1059,13 +1063,16 @@ class Worker:
             # Serialize result
             output_data = json.dumps(result).encode("utf-8")
 
+            # Return session_id in metadata so UI can persist it
+            metadata = {"session_id": session_id}
+
             return PyExecuteComponentResponse(
                 invocation_id=request.invocation_id,
                 success=True,
                 output_data=output_data,
                 state_update=None,
                 error_message=None,
-                metadata=None,
+                metadata=metadata,
                 is_chunk=False,
                 done=True,
                 chunk_index=0,
@@ -1127,6 +1134,13 @@ class Worker:
         if self.metadata:
             self._rust_worker.set_service_metadata(self.metadata)
 
+        # 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
         # without spawn_blocking overhead, enabling true concurrency

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.8a2.dist-info → agnt5-0.2.8a4.dist-info}/METADATA

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

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

@@ -1,15 +1,15 @@
-agnt5-0.2.8a2.dist-info/METADATA,sha256=1Lj5nOz64iKpWdMoxTRlt_DB4QavSlWDrqqjZXzH11c,996
-agnt5-0.2.8a2.dist-info/WHEEL,sha256=AdMozAxftELsa3nYun92mL1tYO-R1ewuDPju53zvoK0,107
-agnt5/__init__.py,sha256=Cscbhye6pA8Jp-sKBGfP4kYXElUdF6aOSHVf-7ph4Dg,2045
+agnt5-0.2.8a4.dist-info/METADATA,sha256=Dp8yJo_dx5B79OM0ptGgR80TjCUbyovpSc4Pt57dfuY,996
+agnt5-0.2.8a4.dist-info/WHEEL,sha256=AdMozAxftELsa3nYun92mL1tYO-R1ewuDPju53zvoK0,107
+agnt5/__init__.py,sha256=ACkK91EPdnv5tYip09QCZ9rfV4iBKzNjGfYVLJD1XGg,2045
 agnt5/_compat.py,sha256=BGuy3v5VDOHVa5f3Z-C22iMN19lAt0mPmXwF3qSSWxI,369
-agnt5/_core.abi3.so,sha256=8fEfT968EuzIbg2BlRpUuzmnKbeB5VCwhI7L9BeSOes,15141600
+agnt5/_core.abi3.so,sha256=CNmJC1AYLy8dABLguVfaWM6YqaTswjRnUZwM_07MXyk,15859720
 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=cz9gDB6c-eRJhBihEIuvTnNBwonpH6G8lbm44j_DKgA,36704
+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=jMnSRTv6MNlK05cJ0FWYQR89ZTz_ywtVuwv-Sjr2Jfc,24925
+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
@@ -17,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=gIbYOdmOczNAqCgErzfLqIukbCpOutdOTZSWv_BatkU,46777
-agnt5/workflow.py,sha256=sU8Gk7unxE_Gla7Fe-KlXfcBvYa2326GciuoR26CCr0,19585
-agnt5-0.2.8a2.dist-info/RECORD,,
+agnt5/worker.py,sha256=NflbueeL2LT8NGywTQnEv1r-N8f54AENWcZARJ5wO8o,47975
+agnt5/workflow.py,sha256=3s9CY6a4UkJZ9YyHv2SAkY3UeCVBlfVi7jxJMFi8Dhg,19488
+agnt5-0.2.8a4.dist-info/RECORD,,