feat(tenancy): harden shared cloud runtime boundaries

tests/unit_tests/pipeline/test_aggregator.py

@@ -13,8 +13,11 @@ from __future__ import annotations
 
 import pytest
 import asyncio
+import contextvars
+from contextlib import asynccontextmanager
 from unittest.mock import Mock, AsyncMock
 from importlib import import_module
+from types import SimpleNamespace
 
 from tests.factories import (
     FakeApp,
@@ -894,17 +897,48 @@ class TestMessageAggregatorWorkspaceIsolation:
         app = make_aggregator_app()
         enable_aggregation(app)
         agg = get_aggregator_module().MessageAggregator(app)
-        delayed_flush = AsyncMock()
+        request_value = contextvars.ContextVar('aggregator_request_value', default=None)
+        token = request_value.set('request-scope')
+        observed = []
+
+        async def delayed_flush(*args):
+            observed.append((request_value.get(), args[2]))
+
         monkeypatch.setattr(agg, '_delayed_flush', delayed_flush)
         context = execution_context(pipeline_uuid='test-pipeline')
 
-        await agg.add_message(**scoped_message_kwargs(context))
-        await asyncio.sleep(0)
+        try:
+            await agg.add_message(**scoped_message_kwargs(context))
+            await asyncio.sleep(0)
+        finally:
+            request_value.reset(token)
 
-        delayed_flush.assert_awaited_once()
-        assert delayed_flush.await_args.args[2] is context
+        assert observed == [(None, context)]
         await agg.flush_all()
 
+    @pytest.mark.asyncio
+    async def test_delayed_flush_opens_explicit_workspace_uow(self, monkeypatch):
+        app = make_aggregator_app()
+        app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
+        scopes = []
+
+        @asynccontextmanager
+        async def tenant_uow(workspace_uuid):
+            scopes.append(workspace_uuid)
+            yield
+
+        app.persistence_mgr.tenant_uow = tenant_uow
+        agg = get_aggregator_module().MessageAggregator(app)
+        flush = AsyncMock()
+        monkeypatch.setattr(agg, '_flush_buffer', flush)
+        context = execution_context('workspace-a', pipeline_uuid='test-pipeline')
+        key = aggregation_key(context, pipeline_uuid='test-pipeline')
+
+        await agg._delayed_flush(key, 0, context)
+
+        assert scopes == ['workspace-a']
+        flush.assert_awaited_once_with(key, context)
+
     @pytest.mark.asyncio
     async def test_flush_rejects_context_from_another_workspace(self):
         app = make_aggregator_app()

tests/unit_tests/pipeline/test_controller_tenancy.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
+import asyncio
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, Mock
 
 import pytest
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import create_async_engine
 
+from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
+from langbot.pkg.persistence.tenant_uow import PersistenceScopeKind
 from langbot.pkg.pipeline.controller import Controller
 from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError
 
@@ -44,6 +49,79 @@ async def test_controller_drops_stale_query_before_pipeline_lookup(
     query_pool.condition.notify_all.assert_called_once_with()
 
 
+@pytest.mark.asyncio
+async def test_cloud_controller_releases_database_connection_during_pipeline_wait(
+    tmp_path,
+    mock_app,
+    sample_query,
+):
+    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "pipeline-short-scope.db"}')
+    table = sa.Table('pipeline_scope_probe', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
+    manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
+    manager.db = SimpleNamespace(get_engine=lambda: engine)
+    checked_out = 0
+
+    def on_checkout(*_args):
+        nonlocal checked_out
+        checked_out += 1
+
+    def on_checkin(*_args):
+        nonlocal checked_out
+        checked_out -= 1
+
+    sa.event.listen(engine.sync_engine, 'checkout', on_checkout)
+    sa.event.listen(engine.sync_engine, 'checkin', on_checkin)
+    try:
+        async with engine.begin() as conn:
+            await conn.run_sync(table.metadata.create_all)
+
+        _prepare_scheduler(mock_app)
+        mock_app.persistence_mgr = manager
+        pipeline_waiting = asyncio.Event()
+        release_pipeline = asyncio.Event()
+
+        async def get_binding(*_args, **_kwargs):
+            assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
+            assert manager.current_session() is None
+            await manager.execute_async(sa.select(table.c.id))
+            assert manager.current_session() is None
+            return SimpleNamespace(
+                instance_uuid='test-instance',
+                workspace_uuid='test-workspace',
+                placement_generation=1,
+            )
+
+        async def run_pipeline(_query):
+            await manager.execute_async(sa.select(table.c.id))
+            assert manager.current_session() is None
+            pipeline_waiting.set()
+            await release_pipeline.wait()
+            assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
+            assert manager.current_session() is None
+
+        runtime_pipeline = SimpleNamespace(run=AsyncMock(side_effect=run_pipeline))
+
+        async def get_pipeline(*_args, **_kwargs):
+            await manager.execute_async(sa.select(table.c.id))
+            assert manager.current_session() is None
+            return runtime_pipeline
+
+        mock_app.workspace_service.get_execution_binding = AsyncMock(side_effect=get_binding)
+        mock_app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock(side_effect=get_pipeline)
+        controller = Controller(mock_app)
+
+        task = asyncio.create_task(controller._process_query(sample_query))
+        await asyncio.wait_for(pipeline_waiting.wait(), timeout=2)
+        assert checked_out == 0
+        assert not task.done()
+        release_pipeline.set()
+        await asyncio.wait_for(task, timeout=2)
+        assert checked_out == 0
+        runtime_pipeline.run.assert_awaited_once_with(sample_query)
+    finally:
+        await engine.dispose()
+
+
 @pytest.mark.asyncio
 async def test_controller_revalidates_generation_before_running_pipeline(
     mock_app,