From f59343fd5be8c92e99f7f6ba4a5a2bf299887d8b Mon Sep 17 00:00:00 2001 From: Hyu Date: Wed, 5 Aug 2026 23:26:23 +0800 Subject: [PATCH] fix: provision Cloud models after Workspace activation (#2403) Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com> --- src/langbot/pkg/cloud/directory_projection.py | 10 +++ src/langbot/pkg/cloud/model_catalog.py | 12 ++- .../cloud/test_directory_projection.py | 33 +++++++ tests/unit_tests/cloud/test_model_catalog.py | 88 +++++++++++++++++++ 4 files changed, 142 insertions(+), 1 deletion(-) diff --git a/src/langbot/pkg/cloud/directory_projection.py b/src/langbot/pkg/cloud/directory_projection.py index 3d0f82e44..46d06f153 100644 --- a/src/langbot/pkg/cloud/directory_projection.py +++ b/src/langbot/pkg/cloud/directory_projection.py @@ -358,6 +358,7 @@ class DirectoryProjectionService: await self._reconcile_entitlement_snapshot_set(snapshot) self._publish_runtime_execution_projection(snapshot.workspaces) + self._request_model_catalog_sync() self._record_batch_cardinality( active_workspaces=active_workspace_count, workspaces=workspace_count, @@ -466,6 +467,7 @@ class DirectoryProjectionService: returned.values(), affected_workspace_uuids=requested, ) + self._request_model_catalog_sync() self._record_batch_cardinality( active_workspaces=active_workspace_count, workspaces=workspace_count, @@ -475,6 +477,14 @@ class DirectoryProjectionService: self._record_success() self._consumer_cursor = batch.cursor + def _request_model_catalog_sync(self) -> None: + """Wake model provisioning after a committed directory change.""" + + service = getattr(self.ap, 'cloud_model_catalog_service', None) + request_sync = getattr(service, 'request_sync', None) + if callable(request_sync): + request_sync() + def _publish_runtime_execution_projection( self, workspaces: Iterable[DirectoryWorkspace], diff --git a/src/langbot/pkg/cloud/model_catalog.py b/src/langbot/pkg/cloud/model_catalog.py index 056159e0a..1d7808202 100644 --- a/src/langbot/pkg/cloud/model_catalog.py +++ b/src/langbot/pkg/cloud/model_catalog.py @@ -151,6 +151,7 @@ class CloudModelCatalogSyncService: # following database reconciliation is a no-op. self._runtime_reload_pending = False self._workspace_credits: dict[str, int | None] = {} + self._sync_requested = asyncio.Event() def get_workspace_credits(self, workspace_uuid: str) -> int | None: """Return the latest signed owner-credit projection for a Workspace.""" @@ -159,9 +160,18 @@ class CloudModelCatalogSyncService: async def initialize(self) -> None: await self.sync_once(reload_runtime=False) + def request_sync(self) -> None: + """Wake the catalog loop after a directory Workspace change.""" + + self._sync_requested.set() + async def run(self) -> None: while True: - await asyncio.sleep(self.sync_interval_seconds) + try: + await asyncio.wait_for(self._sync_requested.wait(), timeout=self.sync_interval_seconds) + except TimeoutError: + pass + self._sync_requested.clear() try: await self.sync_once(reload_runtime=True) except asyncio.CancelledError: diff --git a/tests/unit_tests/cloud/test_directory_projection.py b/tests/unit_tests/cloud/test_directory_projection.py index ec11b9f74..0256ae595 100644 --- a/tests/unit_tests/cloud/test_directory_projection.py +++ b/tests/unit_tests/cloud/test_directory_projection.py @@ -181,6 +181,39 @@ def _delta( ) +async def test_directory_delta_requests_model_catalog_sync_after_commit(projection_context): + application, _session_factory = projection_context + request_sync = Mock() + application.cloud_model_catalog_service = SimpleNamespace(request_sync=request_sync) + event = DirectoryEvent( + cursor=2, + uuid='20000000-0000-4000-8000-000000000002', + aggregate_uuid=WORKSPACE_UUID, + event_type='directory.changed', + revision=2, + payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2}, + created_at=datetime.datetime(2026, 7, 24, 12, 30, tzinfo=datetime.UTC), + ) + batch = DirectoryEventBatch( + instance_uuid=INSTANCE_UUID, + after_cursor=1, + cursor=2, + high_water_cursor=2, + events=[event], + ) + service = DirectoryProjectionService( + application, + _Provider([_snapshot(1)], [batch], [_delta(workspaces=[_workspace(revision=2)])]), + INSTANCE_UUID, + ) + await service.initialize() + request_sync.reset_mock() + + await service.sync_once() + + request_sync.assert_called_once_with() + + async def test_initial_snapshot_projects_core_owned_rows(projection_context): application, session_factory = projection_context reconcile_execution_projection = Mock() diff --git a/tests/unit_tests/cloud/test_model_catalog.py b/tests/unit_tests/cloud/test_model_catalog.py index 5d1de001f..d8d9e625d 100644 --- a/tests/unit_tests/cloud/test_model_catalog.py +++ b/tests/unit_tests/cloud/test_model_catalog.py @@ -273,6 +273,94 @@ async def test_snapshot_must_cover_every_active_workspace() -> None: await service.sync_once() +async def test_periodic_sync_discovers_workspace_created_after_startup_cache_release(tmp_path) -> None: + engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "model-catalog-new-workspace.db"}') + manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME) + manager.db = SimpleNamespace(get_engine=lambda: engine) + startup_bindings = [ + SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_A, placement_generation=1) + ] + live_bindings = [ + *startup_bindings, + SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_B, placement_generation=1), + ] + + class _WorkspaceService: + startup_released = False + + async def list_active_execution_bindings(self): + return list(live_bindings if self.startup_released else startup_bindings) + + def release_startup_execution_bindings(self): + self.startup_released = True + + workspace_service = _WorkspaceService() + app = SimpleNamespace( + persistence_mgr=manager, + workspace_service=workspace_service, + model_mgr=SimpleNamespace(load_models_from_db=_AsyncCounter()), + logger=logging.getLogger(__name__), + ) + service = CloudModelCatalogSyncService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID) + + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + await connection.execute( + sqlalchemy.insert(Workspace), + [ + { + 'uuid': WORKSPACE_A, + 'instance_uuid': INSTANCE_UUID, + 'name': 'A', + 'slug': 'a', + 'source': 'cloud_projection', + }, + { + 'uuid': WORKSPACE_B, + 'instance_uuid': INSTANCE_UUID, + 'name': 'B', + 'slug': 'b', + 'source': 'cloud_projection', + }, + ], + ) + + await service.initialize() + workspace_service.release_startup_execution_bindings() + await service.sync_once() + + async with engine.connect() as connection: + provider_b = await connection.scalar( + sqlalchemy.select(ModelProvider).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_B)) + ) + assert provider_b is not None + finally: + await engine.dispose() + + +async def test_catalog_run_wakes_immediately_when_directory_changes() -> None: + sync_started = asyncio.Event() + + class _WakeService(CloudModelCatalogSyncService): + async def sync_once(self, *, reload_runtime: bool = True): + del reload_runtime + sync_started.set() + return {'workspaces': 0, 'created': 0, 'updated': 0, 'deleted': 0} + + app = SimpleNamespace(logger=logging.getLogger(__name__)) + service = _WakeService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID, sync_interval_seconds=3600) + task = asyncio.create_task(service.run()) + try: + await asyncio.sleep(0) + service.request_sync() + await asyncio.wait_for(sync_started.wait(), timeout=0.2) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + async def _async_value(value): return value