From 78068db9c866b8a2d208d24c52602466e4e62073 Mon Sep 17 00:00:00 2001 From: Hyu Date: Fri, 7 Aug 2026 11:32:38 +0800 Subject: [PATCH] fix(cloud): preserve tenant scope for extension tasks (#2408) * fix(cloud): preserve tenant scope for extension tasks * ci: retrigger extension scope checks --------- Co-authored-by: Chan --- .../pkg/api/http/controller/groups/plugins.py | 13 +++++----- .../api/test_plugin_runtime_route_fence.py | 24 ++++++++++++++----- 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/src/langbot/pkg/api/http/controller/groups/plugins.py b/src/langbot/pkg/api/http/controller/groups/plugins.py index e96f8d213..069fcb027 100644 --- a/src/langbot/pkg/api/http/controller/groups/plugins.py +++ b/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( - self.ap, - execution_context.workspace_uuid, - lambda: self.ap.plugin_connector.require_workspace_context(execution_context), - ) + persistence_mgr = getattr(self.ap, 'persistence_mgr', None) + tenant_scope = getattr(persistence_mgr, 'tenant_scope', None) + if callable(tenant_scope): + 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( diff --git a/tests/unit_tests/api/test_plugin_runtime_route_fence.py b/tests/unit_tests/api/test_plugin_runtime_route_fence.py index 3594a7374..509684fe8 100644 --- a/tests/unit_tests/api/test_plugin_runtime_route_fence.py +++ b/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) - yield + 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()