fix: handle block race conditions (#5819)

letta/orm/sqlalchemy_base.py

@@ -9,6 +9,7 @@ from sqlalchemy import Sequence, String, and_, delete, func, or_, select
 from sqlalchemy.exc import DBAPIError, IntegrityError, TimeoutError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Mapped, Session, mapped_column
+from sqlalchemy.orm.exc import StaleDataError
 from sqlalchemy.orm.interfaces import ORMOption
 
 from letta.log import get_logger
@@ -625,6 +626,12 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             if not no_refresh:
                 await db_session.refresh(self)
             return self
+        except StaleDataError as e:
+            # This can occur when using optimistic locking (version_id_col) and:
+            # 1. The row doesn't exist (0 rows matched)
+            # 2. The version has changed (concurrent update)
+            # We convert this to NoResultFound to return a proper 404 error
+            raise NoResultFound(f"{self.__class__.__name__} with id '{self.id}' not found or was updated by another transaction") from e
         except (DBAPIError, IntegrityError) as e:
             self._handle_dbapi_error(e)
 

tests/managers/test_block_manager.py

@@ -483,6 +483,19 @@ async def test_update_block_limit_does_not_reset(server: SyncServer, default_use
     assert updated_block.value == new_content
 
 
+@pytest.mark.asyncio
+async def test_update_nonexistent_block(server: SyncServer, default_user):
+    """Test that updating a non-existent block raises NoResultFound (which maps to 404)."""
+    block_manager = BlockManager()
+
+    # Try to update a block that doesn't exist
+    nonexistent_block_id = "block-7d73d0a7-6e86-4db7-b53a-411c11ed958a"
+    update_data = BlockUpdate(value="Updated Content")
+
+    with pytest.raises(NoResultFound):
+        await block_manager.update_block_async(block_id=nonexistent_block_id, block_update=update_data, actor=default_user)
+
+
 @pytest.mark.asyncio
 async def test_delete_block(server: SyncServer, default_user):
     block_manager = BlockManager()