fix review comments

Author: seratch

src/agents/memory/sqlite_session.py

@@ -326,8 +326,7 @@ async def apply_history_mutations(self, args: SessionHistoryRewriteArgs) -> None
             return
 
         def _apply_history_mutations_sync() -> None:
-            conn = self._get_connection()
-            with self._lock if self._is_memory_db else threading.Lock():
+            with self._locked_connection() as conn:
                 cursor = conn.execute(
                     f"""
                     SELECT message_data FROM {self.messages_table}

tests/test_session.py

@@ -12,6 +12,19 @@
 from .test_responses import get_text_message
 
 
+class _RecordingLock:
+    def __init__(self, lock):
+        self._lock = lock
+        self.enter_count = 0
+
+    def __enter__(self):
+        self.enter_count += 1
+        return self._lock.__enter__()
+
+    def __exit__(self, exc_type, exc, tb):
+        return self._lock.__exit__(exc_type, exc, tb)
+
+
 # Helper functions for parametrized testing of different Runner methods
 def _run_sync_wrapper(agent, input_data, **kwargs):
     """Wrapper for run_sync that properly sets up an event loop."""
@@ -567,6 +580,49 @@ async def test_sqlite_session_file_lock_is_shared_across_instances():
         assert lock_path not in SQLiteSession._file_locks
 
 
+@pytest.mark.asyncio
+async def test_sqlite_session_apply_history_mutations_uses_file_lock():
+    """File-backed history rewrites should reuse the session lock."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        db_path = Path(temp_dir) / "test_rewrite_lock.db"
+        session = SQLiteSession("rewrite_lock_test", db_path)
+        function_call: TResponseInputItem = {
+            "type": "function_call",
+            "call_id": "call-1",
+            "id": "fc_1",
+            "name": "test_tool",
+            "arguments": '{"value":"before"}',
+        }
+        replacement: TResponseInputItem = {
+            "type": "function_call",
+            "call_id": "call-1",
+            "id": "fc_1",
+            "name": "test_tool",
+            "arguments": '{"value":"after"}',
+        }
+
+        await session.add_items([function_call])
+        recording_lock = _RecordingLock(session._lock)
+        session.__dict__["_lock"] = recording_lock
+
+        await session.apply_history_mutations(
+            {
+                "mutations": [
+                    {
+                        "type": "replace_function_call",
+                        "call_id": "call-1",
+                        "replacement": replacement,
+                    }
+                ]
+            }
+        )
+        assert recording_lock.enter_count == 1
+
+        retrieved = await session.get_items()
+        assert retrieved == [replacement]
+        session.close()
+
+
 @pytest.mark.asyncio
 async def test_session_add_items_exception_propagates_in_streamed():
     """Test that exceptions from session.add_items are properly propagated