fix(plugin): keep pre-tenancy plugin storage readable (#2446)

* fix(plugin): adopt legacy scoped storage rows

* fix(plugin): delete adopted legacy storage rows

* test(plugin): cover legacy storage deletion

* fix(plugin): avoid mutating legacy storage on reads

* fix(plugin): make legacy storage adoption atomic

* fix(plugin): resolve concurrent legacy adoption

* fix(plugin): upsert concurrent storage writes

* fix(plugin): retry reads after legacy adoption

* fix(plugin): deduplicate migrated storage keys

---------

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>

src/langbot/pkg/plugin/handler.py

@@ -11,6 +11,8 @@ import traceback
 from dataclasses import dataclass
 
 import sqlalchemy
+import sqlalchemy.dialects.postgresql
+import sqlalchemy.dialects.sqlite
 
 from langbot_plugin.runtime.io import handler
 from langbot_plugin.runtime.io.connection import Connection
@@ -431,6 +433,19 @@ class RuntimeConnectionHandler(handler.Handler):
             return f'{identity.plugin_author}/{identity.plugin_name}'
         raise ValueError(f'Unsupported binary storage owner_type {owner_type!r}')
 
+    @staticmethod
+    def _legacy_binary_storage_key(
+        action_context: ActionContext,
+        *,
+        owner_type: str,
+        owner: str,
+        key: str,
+    ) -> str:
+        """Return the pre-tenancy key shape for a row already scoped to this Workspace."""
+
+        legacy_owner = action_context.workspace_uuid if owner_type == 'workspace' else owner
+        return f'{owner_type}:{legacy_owner}:{key}'
+
     @classmethod
     def _binary_storage_key(
         cls,
@@ -896,17 +911,70 @@ class RuntimeConnectionHandler(handler.Handler):
                 .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid)
                 .where(persistence_bstorage.BinaryStorage.unique_key == unique_key)
             )
+            storage = result.first()
+            if storage is None:
+                legacy_key = self._legacy_binary_storage_key(
+                    action_context,
+                    owner_type=owner_type,
+                    owner=owner,
+                    key=key,
+                )
+                result = await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.select(persistence_bstorage.BinaryStorage)
+                    .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid)
+                    .where(persistence_bstorage.BinaryStorage.unique_key == legacy_key)
+                    .where(persistence_bstorage.BinaryStorage.key == key)
+                    .where(persistence_bstorage.BinaryStorage.owner_type == owner_type)
+                    .where(persistence_bstorage.BinaryStorage.owner == owner)
+                )
+                storage = result.first()
+                if storage is not None:
+                    update_result = await self.ap.persistence_mgr.execute_async(
+                        sqlalchemy.update(persistence_bstorage.BinaryStorage)
+                        .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid)
+                        .where(persistence_bstorage.BinaryStorage.unique_key == legacy_key)
+                        .where(persistence_bstorage.BinaryStorage.key == key)
+                        .where(persistence_bstorage.BinaryStorage.owner_type == owner_type)
+                        .where(persistence_bstorage.BinaryStorage.owner == owner)
+                        .values(unique_key=unique_key, value=value)
+                    )
+                    if update_result.rowcount:
+                        return handler.ActionResponse.success(data={})
+                    canonical_update = await self.ap.persistence_mgr.execute_async(
+                        sqlalchemy.update(persistence_bstorage.BinaryStorage)
+                        .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid)
+                        .where(persistence_bstorage.BinaryStorage.unique_key == unique_key)
+                        .where(persistence_bstorage.BinaryStorage.key == key)
+                        .where(persistence_bstorage.BinaryStorage.owner_type == owner_type)
+                        .where(persistence_bstorage.BinaryStorage.owner == owner)
+                        .values(value=value)
+                    )
+                    if canonical_update.rowcount:
+                        return handler.ActionResponse.success(data={})
+                    storage = None
 
-            if result.first() is not None:
+            if storage is not None:
                 await self.ap.persistence_mgr.execute_async(
                     sqlalchemy.update(persistence_bstorage.BinaryStorage)
                     .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid)
                     .where(persistence_bstorage.BinaryStorage.unique_key == unique_key)
+                    .where(persistence_bstorage.BinaryStorage.key == key)
+                    .where(persistence_bstorage.BinaryStorage.owner_type == owner_type)
+                    .where(persistence_bstorage.BinaryStorage.owner == owner)
                     .values(value=value)
                 )
-            else:
+                return handler.ActionResponse.success(data={})
+
+            dialect_name = self.ap.persistence_mgr.get_db_engine().dialect.name
+            insert = {
+                'postgresql': sqlalchemy.dialects.postgresql.insert,
+                'sqlite': sqlalchemy.dialects.sqlite.insert,
+            }.get(dialect_name)
+            if insert is None:
+                return handler.ActionResponse.error(message=f'Unsupported storage database dialect: {dialect_name}')
             await self.ap.persistence_mgr.execute_async(
-                    sqlalchemy.insert(persistence_bstorage.BinaryStorage).values(
+                insert(persistence_bstorage.BinaryStorage)
+                .values(
                     workspace_uuid=action_context.workspace_uuid,
                     unique_key=unique_key,
                     key=key,
@@ -914,6 +982,10 @@ class RuntimeConnectionHandler(handler.Handler):
                     owner=owner,
                     value=value,
                 )
+                .on_conflict_do_update(
+                    index_elements=['workspace_uuid', 'unique_key'],
+                    set_={'value': value},
+                )
             )
 
             return handler.ActionResponse.success(
@@ -946,6 +1018,29 @@ class RuntimeConnectionHandler(handler.Handler):
             )
 
             storage = result.first()
+            if storage is None:
+                legacy_key = self._legacy_binary_storage_key(
+                    action_context,
+                    owner_type=owner_type,
+                    owner=owner,
+                    key=key,
+                )
+                result = await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.select(persistence_bstorage.BinaryStorage)
+                    .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid)
+                    .where(persistence_bstorage.BinaryStorage.unique_key == legacy_key)
+                    .where(persistence_bstorage.BinaryStorage.key == key)
+                    .where(persistence_bstorage.BinaryStorage.owner_type == owner_type)
+                    .where(persistence_bstorage.BinaryStorage.owner == owner)
+                )
+                storage = result.first()
+                if storage is None:
+                    retry_result = await self.ap.persistence_mgr.execute_async(
+                        sqlalchemy.select(persistence_bstorage.BinaryStorage)
+                        .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid)
+                        .where(persistence_bstorage.BinaryStorage.unique_key == unique_key)
+                    )
+                    storage = retry_result.first()
             if storage is None:
                 return handler.ActionResponse.error(
                     message=f'Storage with key {key} not found',
@@ -981,10 +1076,19 @@ class RuntimeConnectionHandler(handler.Handler):
                     message=str(e),
                 )
 
+            legacy_key = self._legacy_binary_storage_key(
+                action_context,
+                owner_type=owner_type,
+                owner=owner,
+                key=key,
+            )
             await self.ap.persistence_mgr.execute_async(
                 sqlalchemy.delete(persistence_bstorage.BinaryStorage)
                 .where(persistence_bstorage.BinaryStorage.workspace_uuid == action_context.workspace_uuid)
-                .where(persistence_bstorage.BinaryStorage.unique_key == unique_key)
+                .where(persistence_bstorage.BinaryStorage.unique_key.in_((unique_key, legacy_key)))
+                .where(persistence_bstorage.BinaryStorage.key == key)
+                .where(persistence_bstorage.BinaryStorage.owner_type == owner_type)
+                .where(persistence_bstorage.BinaryStorage.owner == owner)
             )
 
             return handler.ActionResponse.success(
@@ -1012,7 +1116,7 @@ class RuntimeConnectionHandler(handler.Handler):
 
             return handler.ActionResponse.success(
                 data={
-                    'keys': result.scalars().all(),
+                    'keys': list(dict.fromkeys(result.scalars().all())),
                 },
             )
 

tests/unit_tests/plugin/test_handler_actions.py

@@ -234,6 +234,7 @@ class TestSetBinaryStorage:
             },
         }
         mock_app.persistence_mgr = Mock()
+        mock_app.persistence_mgr.get_db_engine.return_value = SimpleNamespace(dialect=SimpleNamespace(name='sqlite'))
         mock_app.persistence_mgr.execute_async = AsyncMock(return_value=make_result())
         mock_app.logger = Mock()
         return mock_app
@@ -270,8 +271,8 @@ class TestSetBinaryStorage:
         )
 
         assert response.code == 0
-        assert app.persistence_mgr.execute_async.await_count == 2
+        assert app.persistence_mgr.execute_async.await_count == 3
-        insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
+        insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0])
         assert insert_params['workspace_uuid'] == 'workspace-a'
         assert insert_params['unique_key'] == canonical_binary_key(
             'plugin',
@@ -301,6 +302,69 @@ class TestSetBinaryStorage:
         assert expected_key in update_params.values()
         assert update_params['value'] == b'new'
 
+    @pytest.mark.asyncio
+    async def test_adopts_legacy_storage_before_updating(self, app):
+        """A migrated pre-tenancy row is updated in place rather than duplicated."""
+        runtime_handler = make_handler(app)
+        legacy_storage = SimpleNamespace(unique_key='plugin:test-author/test-plugin:test-key')
+        adopted = SimpleNamespace(rowcount=1)
+        app.persistence_mgr.execute_async.side_effect = [
+            make_result(),
+            make_result(legacy_storage),
+            adopted,
+        ]
+
+        response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](self.payload(b'new'))
+
+        assert response.code == 0
+        assert app.persistence_mgr.execute_async.await_count == 3
+        adoption_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0])
+        expected_key = canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key')
+        assert expected_key in adoption_params.values()
+        assert adoption_params['value'] == b'new'
+
+    @pytest.mark.asyncio
+    async def test_legacy_adoption_race_updates_winning_canonical_row(self, app):
+        runtime_handler = make_handler(app)
+        legacy_storage = SimpleNamespace(unique_key='plugin:test-author/test-plugin:test-key')
+        lost_race = SimpleNamespace(rowcount=0)
+        canonical_winner = SimpleNamespace(rowcount=1)
+        app.persistence_mgr.execute_async.side_effect = [
+            make_result(),
+            make_result(legacy_storage),
+            lost_race,
+            canonical_winner,
+        ]
+
+        response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](self.payload(b'new'))
+
+        assert response.code == 0
+        assert app.persistence_mgr.execute_async.await_count == 4
+        winner_update = compiled_params(app.persistence_mgr.execute_async.await_args_list[3].args[0])
+        assert canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key') in winner_update.values()
+        assert winner_update['value'] == b'new'
+
+    @pytest.mark.asyncio
+    async def test_legacy_adoption_lost_to_delete_inserts_new_value(self, app):
+        runtime_handler = make_handler(app)
+        legacy_storage = SimpleNamespace(unique_key='plugin:test-author/test-plugin:test-key')
+        lost_race = SimpleNamespace(rowcount=0)
+        app.persistence_mgr.execute_async.side_effect = [
+            make_result(),
+            make_result(legacy_storage),
+            lost_race,
+            SimpleNamespace(rowcount=0),
+            make_result(),
+        ]
+
+        response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value](self.payload(b'new'))
+
+        assert response.code == 0
+        assert app.persistence_mgr.execute_async.await_count == 5
+        insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[4].args[0])
+        assert insert_params['unique_key'] == canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key')
+        assert insert_params['value'] == b'new'
+
     @pytest.mark.asyncio
     async def test_invalid_max_value_bytes_falls_back_to_default_limit(self, app):
         """Invalid max_value_bytes uses the 10MB default limit."""
@@ -525,6 +589,46 @@ class TestGetBinaryStorage:
             in statement_params.values()
         )
 
+    @pytest.mark.asyncio
+    async def test_reads_legacy_storage_without_mutating_key(self, app):
+        runtime_handler = make_handler(app)
+        legacy_storage = SimpleNamespace(
+            unique_key='plugin:test-author/test-plugin:test-key',
+            value=b'legacy bytes',
+        )
+        app.persistence_mgr.execute_async.side_effect = [
+            make_result(),
+            make_result(legacy_storage),
+        ]
+
+        response = await runtime_handler.actions[RuntimeToLangBotAction.GET_BINARY_STORAGE.value](
+            {'key': 'test-key', 'owner_type': 'plugin', 'owner': 'ignored'}
+        )
+
+        assert response.code == 0
+        assert base64.b64decode(response.data['value_base64']) == b'legacy bytes'
+        assert app.persistence_mgr.execute_async.await_count == 2
+
+    @pytest.mark.asyncio
+    async def test_retries_canonical_after_concurrent_legacy_adoption(self, app):
+        runtime_handler = make_handler(app)
+        canonical_storage = SimpleNamespace(value=b'adopted bytes')
+        app.persistence_mgr.execute_async.side_effect = [
+            make_result(),
+            make_result(),
+            make_result(canonical_storage),
+        ]
+
+        response = await runtime_handler.actions[RuntimeToLangBotAction.GET_BINARY_STORAGE.value](
+            {'key': 'test-key', 'owner_type': 'plugin', 'owner': 'ignored'}
+        )
+
+        assert response.code == 0
+        assert base64.b64decode(response.data['value_base64']) == b'adopted bytes'
+        assert app.persistence_mgr.execute_async.await_count == 3
+        retry_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[2].args[0])
+        assert canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key') in retry_params.values()
+
     @pytest.mark.asyncio
     async def test_returns_error_when_not_found(self, app):
         """Missing binary storage rows return an error response."""
@@ -567,21 +671,47 @@ class TestDeleteAndListBinaryStorage:
 
         assert response.code == 0
         statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0])
-        assert 'workspace-a' in statement_params.values()
+        flat_values = [
+            item for value in statement_params.values() for item in (value if isinstance(value, list) else [value])
+        ]
+        assert 'workspace-a' in flat_values
         assert (
             canonical_binary_key(
                 'plugin',
                 'test-author/test-plugin',
                 'test-key',
             )
-            in statement_params.values()
+            in flat_values
         )
-        assert 'forged-owner' not in statement_params.values()
+        assert 'forged-owner' not in flat_values
+
+    @pytest.mark.asyncio
+    async def test_delete_removes_canonical_and_legacy_scoped_keys(self, app):
+        runtime_handler = make_handler(app)
+
+        response = await runtime_handler.actions[RuntimeToLangBotAction.DELETE_BINARY_STORAGE.value](
+            {
+                'key': 'test-key',
+                'owner_type': 'plugin',
+                'owner': 'forged-owner',
+            }
+        )
+
+        assert response.code == 0
+        statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0])
+        values = [
+            item for value in statement_params.values() for item in (value if isinstance(value, list) else [value])
+        ]
+        assert 'workspace-a' in values
+        assert canonical_binary_key('plugin', 'test-author/test-plugin', 'test-key') in values
+        assert 'plugin:test-author/test-plugin:test-key' in values
+        assert 'test-author/test-plugin' in values
+        assert 'forged-owner' not in values
 
     @pytest.mark.asyncio
     async def test_list_keys_uses_trusted_plugin_owner(self, app):
         result = Mock()
-        result.scalars.return_value.all.return_value = ['first', 'second']
+        result.scalars.return_value.all.return_value = ['first', 'second', 'first']
         app.persistence_mgr.execute_async.return_value = result
         runtime_handler = make_handler(app)