feat(cloud): provision workspace model catalog (#2376)

* feat(cloud): provision workspace model catalog

* ci(cloud): pin model catalog adapter source

* fix: make cloud model catalog sync recoverable

* ci: pin cloud adapter source for release

---------

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
Hyu
2026-08-01 18:16:02 +08:00
committed by GitHub
parent 161ea9b3eb
commit d64278ab3f
15 changed files with 979 additions and 15 deletions
@@ -0,0 +1,112 @@
"""Cloud Runtime write protection for the managed LangBot Models catalog."""
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from langbot.pkg.api.http.service import model as model_service_module
from langbot.pkg.api.http.service.model import (
EmbeddingModelsService,
LLMModelsService,
RerankModelsService,
_assert_cloud_managed_provider_mutable,
)
from langbot.pkg.cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER
WORKSPACE = 'workspace-a'
PROVIDER = 'managed-provider'
MODEL = 'managed-model'
@pytest.mark.asyncio
async def test_managed_provider_guard_is_cloud_only(monkeypatch) -> None:
async def managed_provider(_ap, _context, provider_uuid):
assert provider_uuid == PROVIDER
return {'uuid': PROVIDER, 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER}
monkeypatch.setattr(model_service_module, '_require_workspace_provider', managed_provider)
application = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='cloud_runtime')))
with pytest.raises(ValueError, match='managed by Cloud'):
await _assert_cloud_managed_provider_mutable(
application,
WORKSPACE,
PROVIDER,
)
application.persistence_mgr.mode.value = 'normal'
await _assert_cloud_managed_provider_mutable(
application,
WORKSPACE,
PROVIDER,
)
@pytest.mark.parametrize(
('service_type', 'create_method', 'model_data'),
[
(LLMModelsService, 'create_llm_model', {'provider_uuid': PROVIDER, 'name': 'chat', 'abilities': []}),
(EmbeddingModelsService, 'create_embedding_model', {'provider_uuid': PROVIDER, 'name': 'embedding'}),
(RerankModelsService, 'create_rerank_model', {'provider_uuid': PROVIDER, 'name': 'rerank'}),
],
)
@pytest.mark.asyncio
async def test_all_model_types_reject_creation_under_managed_provider(
monkeypatch,
service_type,
create_method: str,
model_data: dict,
) -> None:
guard = AsyncMock(side_effect=ValueError('LangBot Models is managed by Cloud and cannot be modified'))
monkeypatch.setattr(model_service_module, '_assert_cloud_managed_provider_mutable', guard)
application = SimpleNamespace(
persistence_mgr=SimpleNamespace(),
provider_service=SimpleNamespace(
get_provider=AsyncMock(return_value={'uuid': PROVIDER, 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER})
),
model_mgr=None,
)
service = service_type(application)
with pytest.raises(ValueError, match='managed by Cloud'):
await getattr(service, create_method)(WORKSPACE, model_data)
guard.assert_awaited_once()
@pytest.mark.parametrize(
('service_type', 'get_method', 'write_method', 'payload'),
[
(LLMModelsService, 'get_llm_model', 'update_llm_model', {'name': 'changed'}),
(LLMModelsService, 'get_llm_model', 'delete_llm_model', None),
(EmbeddingModelsService, 'get_embedding_model', 'update_embedding_model', {'name': 'changed'}),
(EmbeddingModelsService, 'get_embedding_model', 'delete_embedding_model', None),
(RerankModelsService, 'get_rerank_model', 'update_rerank_model', {'name': 'changed'}),
(RerankModelsService, 'get_rerank_model', 'delete_rerank_model', None),
],
)
@pytest.mark.asyncio
async def test_all_model_types_reject_update_and_delete_for_managed_provider(
monkeypatch,
service_type,
get_method: str,
write_method: str,
payload: dict | None,
) -> None:
guard = AsyncMock(side_effect=ValueError('LangBot Models is managed by Cloud and cannot be modified'))
monkeypatch.setattr(model_service_module, '_assert_cloud_managed_provider_mutable', guard)
application = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='cloud_runtime')))
service = service_type(application)
monkeypatch.setattr(
service,
get_method,
AsyncMock(return_value={'uuid': MODEL, 'provider_uuid': PROVIDER, 'extra_args': {}}),
)
args = (WORKSPACE, MODEL) if payload is None else (WORKSPACE, MODEL, payload)
with pytest.raises(ValueError, match='managed by Cloud'):
await getattr(service, write_method)(*args)
guard.assert_awaited_once()
@@ -25,6 +25,7 @@ from langbot.pkg.workspace.errors import WorkspaceNotFoundError
pytestmark = pytest.mark.asyncio
WORKSPACE_UUID = 'workspace-a'
SYSTEM_REQUESTER = 'space-chat-completions'
def _create_mock_provider(
@@ -1005,3 +1006,56 @@ class TestProviderSecretRoundtrip:
)
ap.persistence_mgr.execute_async.assert_not_awaited()
class TestCloudManagedProviderProtection:
@staticmethod
def _service() -> ModelProviderService:
ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(
mode=SimpleNamespace(value='cloud_runtime'),
execute_async=AsyncMock(),
),
model_mgr=SimpleNamespace(),
)
return ModelProviderService(ap)
async def test_cloud_rejects_user_created_system_requester(self):
service = self._service()
with pytest.raises(ValueError, match='reserved'):
await service.create_provider(
WORKSPACE_UUID,
{
'name': 'Fake LangBot Models',
'requester': SYSTEM_REQUESTER,
'base_url': 'https://example.invalid/v1',
'api_keys': ['fake'],
},
)
with pytest.raises(ValueError, match='reserved'):
await service.find_or_create_provider(
WORKSPACE_UUID,
SYSTEM_REQUESTER,
'https://api.langbot.cloud/v1',
['fake'],
)
service.ap.persistence_mgr.execute_async.assert_not_awaited()
async def test_cloud_rejects_update_and_delete_of_managed_provider(self):
service = self._service()
service.get_provider = AsyncMock(
return_value={'uuid': 'system-provider', 'requester': SYSTEM_REQUESTER}
)
with pytest.raises(ValueError, match='managed by Cloud'):
await service.update_provider(WORKSPACE_UUID, 'system-provider', {'name': 'Renamed'})
with pytest.raises(ValueError, match='managed by Cloud'):
await service.delete_provider(WORKSPACE_UUID, 'system-provider')
service.ap.persistence_mgr.execute_async.assert_not_awaited()
async def test_oss_does_not_reserve_space_requester(self):
ap = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='oss_compat')))
service = ModelProviderService(ap)
assert service._system_requester_is_reserved(SYSTEM_REQUESTER) is False
+5
View File
@@ -66,6 +66,10 @@ class _Provider:
def __init__(self):
self.manifest_provider = _Manifest()
async def fetch_model_catalog(self, instance_uuid: str):
del instance_uuid
raise AssertionError('not used by bootstrap contract tests')
def bootstrap(self, *, instance_uuid: str, instance_config: dict):
del instance_config
return VerifiedCloudDeployment(
@@ -79,6 +83,7 @@ class _Provider:
entitlement_provider=_Entitlements(),
directory_provider=_Directory(),
manifest_provider=self.manifest_provider,
model_catalog_provider=self,
verification_key_id='root-2026',
)
@@ -0,0 +1,363 @@
from __future__ import annotations
import asyncio
import logging
from datetime import UTC, datetime
from types import SimpleNamespace
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.cloud.model_catalog import (
CloudModelCatalogSnapshot,
CloudModelCatalogSyncService,
system_model_uuid,
system_provider_uuid,
)
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.model import EmbeddingModel, LLMModel, ModelProvider
from langbot.pkg.entity.persistence.workspace import Workspace
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
pytestmark = pytest.mark.asyncio
INSTANCE_UUID = 'instance-model-catalog'
WORKSPACE_A = '00000000-0000-4000-8000-000000000001'
WORKSPACE_B = '00000000-0000-4000-8000-000000000002'
OWNER_A = '10000000-0000-4000-8000-000000000001'
OWNER_B = '10000000-0000-4000-8000-000000000002'
class _CatalogProvider:
def __init__(self, snapshot: CloudModelCatalogSnapshot) -> None:
self.snapshot = snapshot
async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot:
assert instance_uuid == INSTANCE_UUID
return self.snapshot
def _snapshot(
*,
key_a: str | None = 'owner-a-key',
model_id: str = 'gpt-test',
include_embedding: bool = True,
) -> CloudModelCatalogSnapshot:
models = [
{
'uuid': 'upstream-chat',
'model_id': model_id,
'category': 'chat',
'llm_abilities': ['chat', 'vision'],
'is_featured': True,
'featured_order': 7,
}
]
if include_embedding:
models.append(
{
'uuid': 'upstream-embedding',
'model_id': 'embedding-test',
'category': 'embedding',
}
)
return CloudModelCatalogSnapshot.model_validate(
{
'instance_uuid': INSTANCE_UUID,
'generated_at': datetime.now(UTC),
'base_url': 'https://api.langbot.cloud/v1/',
'models': models,
'workspaces': [
{
'workspace_uuid': WORKSPACE_A,
'owner_account_uuid': OWNER_A,
'api_key': key_a,
},
{
'workspace_uuid': WORKSPACE_B,
'owner_account_uuid': OWNER_B,
'api_key': 'owner-b-key',
},
],
}
)
async def test_catalog_reconciles_every_workspace_idempotently_and_tracks_owner_and_downlisting(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "model-catalog.db"}')
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
bindings = [
SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_A, placement_generation=1),
SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_B, placement_generation=1),
]
workspace_service = SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings))
reload_counter = _AsyncCounter()
runtime_reload = SimpleNamespace(load_models_from_db=reload_counter)
app = SimpleNamespace(
persistence_mgr=manager,
workspace_service=workspace_service,
model_mgr=runtime_reload,
logger=logging.getLogger(__name__),
)
provider = _CatalogProvider(_snapshot())
service = CloudModelCatalogSyncService(app, provider, 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 connection.execute(
sqlalchemy.insert(ModelProvider).values(
uuid='custom-provider',
workspace_uuid=WORKSPACE_A,
name='Custom',
requester='openai-chat-completions',
base_url='https://custom.example/v1',
api_keys=['custom-key'],
)
)
await connection.execute(
sqlalchemy.insert(LLMModel).values(
uuid='custom-model',
workspace_uuid=WORKSPACE_A,
name='custom-model',
provider_uuid='custom-provider',
abilities=['chat'],
extra_args={},
prefered_ranking=0,
)
)
first = await service.sync_once()
assert first == {'workspaces': 2, 'created': 6, 'updated': 0, 'deleted': 0}
assert reload_counter.calls == 1
async with engine.connect() as connection:
providers = (
await connection.execute(
sqlalchemy.select(
ModelProvider.uuid,
ModelProvider.workspace_uuid,
ModelProvider.api_keys,
).where(ModelProvider.requester == 'space-chat-completions')
)
).all()
assert {item.workspace_uuid for item in providers} == {WORKSPACE_A, WORKSPACE_B}
assert {item.uuid for item in providers} == {
system_provider_uuid(WORKSPACE_A),
system_provider_uuid(WORKSPACE_B),
}
assert {item.workspace_uuid: item.api_keys for item in providers} == {
WORKSPACE_A: ['owner-a-key'],
WORKSPACE_B: ['owner-b-key'],
}
assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(LLMModel)) == 3
assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(EmbeddingModel)) == 2
second = await service.sync_once()
assert second == {'workspaces': 2, 'created': 0, 'updated': 0, 'deleted': 0}
assert reload_counter.calls == 1
provider.snapshot = _snapshot(
key_a='new-owner-key',
model_id='gpt-renamed',
include_embedding=False,
)
third = await service.sync_once()
assert third == {'workspaces': 2, 'created': 0, 'updated': 3, 'deleted': 2}
assert reload_counter.calls == 2
async with engine.connect() as connection:
provider_a_keys = await connection.scalar(
sqlalchemy.select(ModelProvider.api_keys).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_A))
)
assert provider_a_keys == ['new-owner-key']
system_model_names = (
(
await connection.execute(
sqlalchemy.select(LLMModel.name).where(
LLMModel.provider_uuid.in_(
[system_provider_uuid(WORKSPACE_A), system_provider_uuid(WORKSPACE_B)]
)
)
)
)
.scalars()
.all()
)
assert set(system_model_names) == {'gpt-renamed'}
assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(EmbeddingModel)) == 0
assert (
await connection.scalar(
sqlalchemy.select(sqlalchemy.func.count())
.select_from(ModelProvider)
.where(ModelProvider.uuid == 'custom-provider')
)
== 1
)
assert (
await connection.scalar(
sqlalchemy.select(sqlalchemy.func.count())
.select_from(LLMModel)
.where(LLMModel.uuid == 'custom-model')
)
== 1
)
provider.snapshot = _snapshot(key_a=None, model_id='gpt-renamed', include_embedding=False)
fourth = await service.sync_once()
assert fourth == {'workspaces': 2, 'created': 0, 'updated': 1, 'deleted': 0}
assert reload_counter.calls == 3
async with engine.connect() as connection:
provider_a_keys = await connection.scalar(
sqlalchemy.select(ModelProvider.api_keys).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_A))
)
assert provider_a_keys == []
finally:
await engine.dispose()
def test_workspace_scoped_ids_are_stable_and_secrets_are_redacted() -> None:
assert system_provider_uuid(WORKSPACE_A) == system_provider_uuid(WORKSPACE_A)
assert system_provider_uuid(WORKSPACE_A) != system_provider_uuid(WORKSPACE_B)
assert system_model_uuid(WORKSPACE_A, 'chat', 'upstream') != system_model_uuid(WORKSPACE_B, 'chat', 'upstream')
snapshot = _snapshot()
assert 'owner-a-key' not in repr(snapshot)
async def test_snapshot_must_cover_every_active_workspace() -> None:
snapshot = _snapshot().model_copy(update={'workspaces': _snapshot().workspaces[:1]})
app = SimpleNamespace(
workspace_service=SimpleNamespace(
list_active_execution_bindings=lambda: _async_value(
[SimpleNamespace(workspace_uuid=WORKSPACE_A), SimpleNamespace(workspace_uuid=WORKSPACE_B)]
)
),
logger=logging.getLogger(__name__),
)
service = CloudModelCatalogSyncService(app, _CatalogProvider(snapshot), INSTANCE_UUID)
with pytest.raises(ValueError, match='missing billing projections for 1 active Workspaces'):
await service.sync_once()
async def _async_value(value):
return value
class _AsyncCounter:
def __init__(self) -> None:
self.calls = 0
async def __call__(self) -> None:
self.calls += 1
async def test_partial_workspace_failure_reloads_already_committed_changes() -> None:
bindings = [
SimpleNamespace(workspace_uuid=WORKSPACE_A),
SimpleNamespace(workspace_uuid=WORKSPACE_B),
]
reload_counter = _AsyncCounter()
app = SimpleNamespace(
workspace_service=SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings)),
model_mgr=SimpleNamespace(load_models_from_db=reload_counter),
logger=logging.getLogger(__name__),
)
service = CloudModelCatalogSyncService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID)
calls = 0
async def sync_workspace(*_args):
nonlocal calls
calls += 1
if calls == 1:
return {'created': 1, 'updated': 0, 'deleted': 0}
raise RuntimeError('second Workspace failed')
service._sync_workspace = sync_workspace # type: ignore[method-assign]
with pytest.raises(RuntimeError, match='second Workspace failed'):
await service.sync_once()
assert reload_counter.calls == 1
async def test_failed_runtime_reload_is_retried_after_noop_sync() -> None:
bindings = [SimpleNamespace(workspace_uuid=WORKSPACE_A)]
class _FlakyReload:
def __init__(self) -> None:
self.calls = 0
async def __call__(self) -> None:
self.calls += 1
if self.calls == 1:
raise RuntimeError('reload failed')
runtime_reload = _FlakyReload()
app = SimpleNamespace(
workspace_service=SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings)),
model_mgr=SimpleNamespace(load_models_from_db=runtime_reload),
logger=logging.getLogger(__name__),
)
service = CloudModelCatalogSyncService(app, _CatalogProvider(_snapshot()), INSTANCE_UUID)
calls = 0
async def sync_workspace(*_args):
nonlocal calls
calls += 1
if calls == 1:
return {'created': 1, 'updated': 0, 'deleted': 0}
return {'created': 0, 'updated': 0, 'deleted': 0}
service._sync_workspace = sync_workspace # type: ignore[method-assign]
with pytest.raises(RuntimeError, match='reload failed'):
await service.sync_once()
summary = await service.sync_once()
assert summary == {'workspaces': 1, 'created': 0, 'updated': 0, 'deleted': 0}
assert runtime_reload.calls == 2
async def test_background_sync_log_redacts_exception_message(caplog) -> None:
secret = 'owner-secret-api-key'
attempted = asyncio.Event()
class _FailingProvider:
async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot:
del instance_uuid
attempted.set()
raise RuntimeError(f'database parameters include {secret}')
app = SimpleNamespace(logger=logging.getLogger(__name__))
service = CloudModelCatalogSyncService(app, _FailingProvider(), INSTANCE_UUID)
service.sync_interval_seconds = 0.001
task = asyncio.create_task(service.run())
try:
await asyncio.wait_for(attempted.wait(), timeout=1)
await asyncio.sleep(0.01)
finally:
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert secret not in caplog.text
assert 'Cloud model catalog synchronization failed (RuntimeError)' in caplog.text