mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-07 03:46:38 +00:00
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 <dadachann@users.noreply.github.com>
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user