fix(cloud): preserve tenant scope for extension tasks

src/langbot/pkg/api/http/controller/groups/plugins.py

@@ -15,7 +15,6 @@ import posixpath
 import sqlalchemy
 
 from .....core import taskmgr
-from .....core.task_boundary import run_in_workspace_uow
 from .....entity.persistence import plugin as persistence_plugin
 from ...authz import Permission
 from ...context import ExecutionContext, RequestContext
@@ -311,11 +310,13 @@ class PluginsRouterGroup(group.RouterGroup):
     ):
         """Revalidate a captured task context immediately before Runtime I/O."""
 
-        await run_in_workspace_uow(
+        persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
-            self.ap,
+        tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
-            execution_context.workspace_uuid,
+        if callable(tenant_scope):
-            lambda: self.ap.plugin_connector.require_workspace_context(execution_context),
+            async with tenant_scope(execution_context.workspace_uuid):
-        )
+                await self.ap.plugin_connector.require_workspace_context(execution_context)
+                return await operation()
+        await self.ap.plugin_connector.require_workspace_context(execution_context)
         return await operation()
 
     async def _require_authenticated_plugin_runtime_context(

tests/unit_tests/api/test_plugin_runtime_route_fence.py

@@ -124,24 +124,37 @@ async def test_background_plugin_operation_refences_captured_generation(plugin_r
 
 
 @pytest.mark.asyncio
-async def test_background_plugin_operation_revalidates_inside_short_tenant_uow(plugin_router_cls):
+async def test_background_plugin_operation_revalidates_and_runs_inside_tenant_uow(plugin_router_cls):
     scopes = []
+    active_scope = None
+
+    transaction_active = False
 
     @asynccontextmanager
-    async def tenant_uow(workspace_uuid):
+    async def tenant_scope(workspace_uuid):
+        nonlocal active_scope
         scopes.append(workspace_uuid)
+        active_scope = workspace_uuid
+        try:
             yield
+        finally:
+            active_scope = None
 
     connector = SimpleNamespace(
         require_workspace_context=AsyncMock(side_effect=lambda context: context),
     )
-    operation = AsyncMock(return_value='done')
+
+    async def operation():
+        assert active_scope == CONTEXT.workspace_uuid
+        assert transaction_active is False
+        return 'done'
+
     router = object.__new__(plugin_router_cls)
     router.ap = SimpleNamespace(
         plugin_connector=connector,
         persistence_mgr=SimpleNamespace(
             mode=SimpleNamespace(value='cloud_runtime'),
-            tenant_uow=tenant_uow,
+            tenant_scope=tenant_scope,
         ),
     )
 
@@ -150,4 +163,3 @@ async def test_background_plugin_operation_revalidates_inside_short_tenant_uow(p
     assert result == 'done'
     assert scopes == [CONTEXT.workspace_uuid]
     connector.require_workspace_context.assert_awaited_once_with(CONTEXT)
-    operation.assert_awaited_once()