feat(session): add dirty flag to skip unnecessary agent state persistence (#1803)

Co-authored-by: Strands Agent <217235299+strands-agent@users.noreply.github.com>

Author: Unshure

src/strands/interrupt.py

@@ -52,10 +52,12 @@ class _InterruptState:
     interrupts: dict[str, Interrupt] = field(default_factory=dict)
     context: dict[str, Any] = field(default_factory=dict)
     activated: bool = False
+    _version: int = field(default=0, compare=False, repr=False)
 
     def activate(self) -> None:
         """Activate the interrupt state."""
         self.activated = True
+        self._version += 1
 
     def deactivate(self) -> None:
         """Deacitvate the interrupt state.
@@ -65,6 +67,7 @@ def deactivate(self) -> None:
         self.interrupts = {}
         self.context = {}
         self.activated = False
+        self._version += 1
 
     def resume(self, prompt: "AgentInput") -> None:
         """Configure the interrupt state if resuming from an interrupt event.
@@ -100,10 +103,27 @@ def resume(self, prompt: "AgentInput") -> None:
             self.interrupts[interrupt_id].response = interrupt_response
 
         self.context["responses"] = contents
+        self._version += 1
+
+    def _get_version(self) -> int:
+        """Get the current version number of the interrupt state.
+
+        The version is incremented each time activate(), deactivate(), or resume() is called.
+        Consumers can compare versions to detect changes without requiring
+        explicit dirty flag clearing.
+
+        Returns:
+            The current version number.
+        """
+        return self._version
 
     def to_dict(self) -> dict[str, Any]:
         """Serialize to dict for session management."""
-        return asdict(self)
+        return {
+            "interrupts": {k: v.to_dict() for k, v in self.interrupts.items()},
+            "context": self.context,
+            "activated": self.activated,
+        }
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "_InterruptState":

src/strands/session/repository_session_manager.py

@@ -1,5 +1,6 @@
 """Repository session manager implementation."""
 
+import copy
 import logging
 from typing import TYPE_CHECKING, Any
 
@@ -59,6 +60,9 @@ def __init__(
         # Keep track of the latest message of each agent in case we need to redact it.
         self._latest_agent_message: dict[str, SessionMessage | None] = {}
 
+        # Track the previously synced internal state for each agent to detect changes.
+        self._last_synced_internal_state: dict[str, dict[str, Any]] = {}
+
     def append_message(self, message: Message, agent: "Agent", **kwargs: Any) -> None:
         """Append a message to the agent's session.
 
@@ -95,15 +99,66 @@ def redact_latest_message(self, redact_message: Message, agent: "Agent", **kwarg
     def sync_agent(self, agent: "Agent", **kwargs: Any) -> None:
         """Serialize and update the agent into the session repository.
 
+        Only updates the agent if state has been modified or internal state has changed.
+        This optimization reduces unnecessary I/O operations when the agent processes
+        messages without modifying its state.
+
         Args:
             agent: Agent to sync to the session.
             **kwargs: Additional keyword arguments for future extensibility.
         """
+        # Get current versions and conversation manager state
+        current_state_version = agent.state._get_version()
+        current_interrupt_state_version = agent._interrupt_state._get_version()
+        current_conversation_manager_state = agent.conversation_manager.get_state()
+
+        # Check if we have a previous state to compare against
+        last_synced = self._last_synced_internal_state.get(agent.agent_id)
+
+        # Determine if we need to update by comparing versions
+        if last_synced is None:
+            # First sync for this agent - always update
+            state_changed = True
+            internal_state_changed = True
+            conversation_manager_state_changed = True
+        else:
+            state_changed = current_state_version != last_synced.get("state_version")
+            internal_state_changed = current_interrupt_state_version != last_synced.get("interrupt_state_version")
+            conversation_manager_state_changed = (
+                current_conversation_manager_state != last_synced.get("conversation_manager_state")
+            )
+
+        if not state_changed and not internal_state_changed and not conversation_manager_state_changed:
+            logger.debug(
+                "agent_id=<%s> | session_id=<%s> | skipping sync, no changes detected",
+                agent.agent_id,
+                self.session_id,
+            )
+            return
+
+        logger.debug(
+            "agent_id=<%s> | session_id=<%s> | state_changed=<%s>, internal_state_changed=<%s>, "
+            "conversation_manager_state_changed=<%s> | syncing agent",
+            agent.agent_id,
+            self.session_id,
+            state_changed,
+            internal_state_changed,
+            conversation_manager_state_changed,
+        )
+
+        # Perform the update
         self.session_repository.update_agent(
             self.session_id,
             SessionAgent.from_agent(agent),
         )
 
+        # Update tracked versions after successful sync
+        self._last_synced_internal_state[agent.agent_id] = {
+            "state_version": current_state_version,
+            "interrupt_state_version": current_interrupt_state_version,
+            "conversation_manager_state": copy.deepcopy(current_conversation_manager_state),
+        }
+
     def initialize(self, agent: "Agent", **kwargs: Any) -> None:
         """Initialize an agent with a session.
 

src/strands/types/json_dict.py

@@ -15,6 +15,7 @@ class JSONSerializableDict:
     def __init__(self, initial_state: dict[str, Any] | None = None):
         """Initialize JSONSerializableDict."""
         self._data: dict[str, Any]
+        self._version: int = 0
         if initial_state:
             self._validate_json_serializable(initial_state)
             self._data = copy.deepcopy(initial_state)
@@ -34,6 +35,7 @@ def set(self, key: str, value: Any) -> None:
         self._validate_key(key)
         self._validate_json_serializable(value)
         self._data[key] = copy.deepcopy(value)
+        self._version += 1
 
     def get(self, key: str | None = None) -> Any:
         """Get a value or entire data.
@@ -57,6 +59,19 @@ def delete(self, key: str) -> None:
         """
         self._validate_key(key)
         self._data.pop(key, None)
+        self._version += 1
+
+    def _get_version(self) -> int:
+        """Get the current version number of the store.
+
+        The version is incremented each time set() or delete() is called.
+        Consumers can compare versions to detect changes without requiring
+        explicit dirty flag clearing.
+
+        Returns:
+            The current version number.
+        """
+        return self._version
 
     def _validate_key(self, key: str) -> None:
         """Validate that a key is valid.

tests/strands/session/test_repository_session_manager.py

@@ -595,3 +595,219 @@ def test_fix_broken_tool_use_does_not_affect_normal_conversations(session_manage
 
     # Should remain unchanged
     assert fixed_messages == messages
+
+
+# ============================================================================
+# Conditional Sync Tests
+# ============================================================================
+
+
+def test_sync_agent_skips_update_when_state_not_dirty_and_internal_state_unchanged(mock_repository):
+    """Test that sync_agent() skips update_agent() when state is not dirty and internal state unchanged."""
+    session_manager = RepositorySessionManager(session_id="test-session", session_repository=mock_repository)
+
+    # Create and initialize agent
+    agent = Agent(agent_id="test-agent", session_manager=session_manager)
+
+    # Track update_agent calls
+    update_agent_calls = []
+    original_update_agent = mock_repository.update_agent
+
+    def tracking_update_agent(session_id, session_agent):
+        update_agent_calls.append((session_id, session_agent))
+        return original_update_agent(session_id, session_agent)
+
+    mock_repository.update_agent = tracking_update_agent
+
+    # First sync should update (to establish baseline)
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 1
+
+    # Clear tracking
+    update_agent_calls.clear()
+
+    # Second sync without changes should skip update
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 0
+
+
+def test_sync_agent_calls_update_when_state_is_dirty(mock_repository):
+    """Test that sync_agent() calls update_agent() when agent.state is dirty."""
+    session_manager = RepositorySessionManager(session_id="test-session", session_repository=mock_repository)
+
+    # Create and initialize agent
+    agent = Agent(agent_id="test-agent", session_manager=session_manager)
+
+    # Track update_agent calls
+    update_agent_calls = []
+    original_update_agent = mock_repository.update_agent
+
+    def tracking_update_agent(session_id, session_agent):
+        update_agent_calls.append((session_id, session_agent))
+        return original_update_agent(session_id, session_agent)
+
+    mock_repository.update_agent = tracking_update_agent
+
+    # First sync to establish baseline
+    session_manager.sync_agent(agent)
+    update_agent_calls.clear()
+
+    # Modify state (makes it dirty)
+    agent.state.set("key", "value")
+
+    # Sync should call update_agent because state is dirty
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 1
+
+
+def test_sync_agent_calls_update_when_internal_state_changed(mock_repository):
+    """Test that sync_agent() calls update_agent() when internal state (interrupt_state) is dirty."""
+    session_manager = RepositorySessionManager(session_id="test-session", session_repository=mock_repository)
+
+    # Create and initialize agent
+    agent = Agent(agent_id="test-agent", session_manager=session_manager)
+
+    # Track update_agent calls
+    update_agent_calls = []
+    original_update_agent = mock_repository.update_agent
+
+    def tracking_update_agent(session_id, session_agent):
+        update_agent_calls.append((session_id, session_agent))
+        return original_update_agent(session_id, session_agent)
+
+    mock_repository.update_agent = tracking_update_agent
+
+    # First sync to establish baseline
+    session_manager.sync_agent(agent)
+    update_agent_calls.clear()
+
+    # Modify internal state (activate interrupt state which sets dirty flag)
+    agent._interrupt_state.activate()
+
+    # Sync should call update_agent because internal state is dirty
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 1
+
+
+def test_sync_agent_calls_update_when_conversation_manager_state_changed(mock_repository):
+    """Test that sync_agent() calls update_agent() when conversation manager state changed."""
+    session_manager = RepositorySessionManager(session_id="test-session", session_repository=mock_repository)
+
+    # Create and initialize agent
+    agent = Agent(agent_id="test-agent", session_manager=session_manager)
+
+    # Track update_agent calls
+    update_agent_calls = []
+    original_update_agent = mock_repository.update_agent
+
+    def tracking_update_agent(session_id, session_agent):
+        update_agent_calls.append((session_id, session_agent))
+        return original_update_agent(session_id, session_agent)
+
+    mock_repository.update_agent = tracking_update_agent
+
+    # First sync to establish baseline
+    session_manager.sync_agent(agent)
+    update_agent_calls.clear()
+
+    # Modify conversation manager state
+    agent.conversation_manager.removed_message_count = 5
+
+    # Sync should call update_agent because conversation manager state changed
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 1
+
+
+def test_sync_agent_tracks_version_after_successful_sync(mock_repository):
+    """Test that sync_agent() tracks version after successful sync."""
+    session_manager = RepositorySessionManager(session_id="test-session", session_repository=mock_repository)
+
+    # Create and initialize agent
+    agent = Agent(agent_id="test-agent", session_manager=session_manager)
+
+    # First sync to establish baseline
+    session_manager.sync_agent(agent)
+    initial_version = agent.state._get_version()
+
+    # Modify state (increments version)
+    agent.state.set("key", "value")
+    assert agent.state._get_version() == initial_version + 1
+
+    # Track update_agent calls
+    update_agent_calls = []
+    original_update_agent = mock_repository.update_agent
+
+    def tracking_update_agent(session_id, session_agent):
+        update_agent_calls.append((session_id, session_agent))
+        return original_update_agent(session_id, session_agent)
+
+    mock_repository.update_agent = tracking_update_agent
+
+    # Sync should update because version changed
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 1
+
+    # Second sync without changes should skip
+    update_agent_calls.clear()
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 0
+
+
+def test_sync_agent_retries_on_failure(mock_repository):
+    """Test that sync_agent() retries on next call if update_agent() fails."""
+    session_manager = RepositorySessionManager(session_id="test-session", session_repository=mock_repository)
+
+    # Create and initialize agent
+    agent = Agent(agent_id="test-agent", session_manager=session_manager)
+
+    # First sync to establish baseline
+    session_manager.sync_agent(agent)
+
+    # Modify state (increments version)
+    agent.state.set("key", "value")
+
+    # Make update_agent fail
+    def failing_update_agent(session_id, session_agent):
+        raise SessionException("Update failed")
+
+    mock_repository.update_agent = failing_update_agent
+
+    # Sync should fail
+    with pytest.raises(SessionException, match="Update failed"):
+        session_manager.sync_agent(agent)
+
+    # Restore working update_agent
+    update_agent_calls = []
+    original_update_agent = MockedSessionRepository.update_agent
+
+    def tracking_update_agent(self, session_id, session_agent):
+        update_agent_calls.append((session_id, session_agent))
+        return original_update_agent(self, session_id, session_agent)
+
+    mock_repository.update_agent = lambda sid, sa: tracking_update_agent(mock_repository, sid, sa)
+
+    # Retry should work because version wasn't updated on failure
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 1
+
+
+def test_sync_agent_first_sync_always_updates(mock_repository):
+    """Test that the first sync_agent() call always updates (no previous state to compare)."""
+    session_manager = RepositorySessionManager(session_id="test-session", session_repository=mock_repository)
+
+    # Create and initialize agent
+    agent = Agent(agent_id="test-agent", session_manager=session_manager)
+
+    # Track update_agent calls
+    update_agent_calls = []
+    original_update_agent = mock_repository.update_agent
+
+    def tracking_update_agent(session_id, session_agent):
+        update_agent_calls.append((session_id, session_agent))
+        return original_update_agent(session_id, session_agent)
+
+    mock_repository.update_agent = tracking_update_agent
+
+    # First sync should always update (no previous state)
+    session_manager.sync_agent(agent)
+    assert len(update_agent_calls) == 1