From 6654473514b765efe7d98b060baaed8fe0b7b5be Mon Sep 17 00:00:00 2001 From: Sarah Wooders Date: Thu, 30 Oct 2025 15:18:04 -0700 Subject: [PATCH] fix: handle block race conditions (#5819) --- letta/orm/sqlalchemy_base.py | 7 +++++++ tests/managers/test_block_manager.py | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/letta/orm/sqlalchemy_base.py b/letta/orm/sqlalchemy_base.py index 0b0b70f1..58692dc6 100644 --- a/letta/orm/sqlalchemy_base.py +++ b/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) diff --git a/tests/managers/test_block_manager.py b/tests/managers/test_block_manager.py index de8f01e7..c3432847 100644 --- a/tests/managers/test_block_manager.py +++ b/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()