From 22de4b14eec4b529f6ab710677dde517c57eeff1 Mon Sep 17 00:00:00 2001 From: dadachann <185672915+dadachann@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:33:16 +0000 Subject: [PATCH] feat: use workspace identity for telemetry --- .../versions/0009_workspace_tenancy_kernel.py | 13 +- .../versions/0017_oss_workspace_identity.py | 177 ++++++++++++++++++ .../pkg/pipeline/process/handlers/chat.py | 4 +- src/langbot/pkg/telemetry/heartbeat.py | 71 +++++-- src/langbot/pkg/telemetry/identity.py | 16 ++ src/langbot/pkg/telemetry/telemetry.py | 20 +- src/langbot/pkg/workspace/identity.py | 20 ++ src/langbot/pkg/workspace/service.py | 3 +- .../persistence/test_workspace_migration.py | 90 +++++++++ tests/unit_tests/telemetry/test_heartbeat.py | 61 +++--- tests/unit_tests/telemetry/test_telemetry.py | 27 +++ .../telemetry/test_workspace_identity.py | 31 +++ .../workspace/test_workspace_service.py | 2 + 13 files changed, 493 insertions(+), 42 deletions(-) create mode 100644 src/langbot/pkg/persistence/alembic/versions/0017_oss_workspace_identity.py create mode 100644 src/langbot/pkg/telemetry/identity.py create mode 100644 src/langbot/pkg/workspace/identity.py create mode 100644 tests/unit_tests/telemetry/test_workspace_identity.py diff --git a/src/langbot/pkg/persistence/alembic/versions/0009_workspace_tenancy_kernel.py b/src/langbot/pkg/persistence/alembic/versions/0009_workspace_tenancy_kernel.py index d0a6c233e..1da7d0286 100644 --- a/src/langbot/pkg/persistence/alembic/versions/0009_workspace_tenancy_kernel.py +++ b/src/langbot/pkg/persistence/alembic/versions/0009_workspace_tenancy_kernel.py @@ -18,6 +18,17 @@ down_revision = '0008_mcp_resource_prefs' branch_labels = None depends_on = None +_WORKSPACE_IDENTITY_NAMESPACE = uuid.UUID('8ea04f29-8528-4cc3-bb28-30a838c89d76') + + +def _workspace_uuid_from_instance_id(instance_id: str) -> str: + value = instance_id.strip() + candidate = value[len('instance_') :] if value.startswith('instance_') else value + try: + return str(uuid.UUID(candidate)) + except ValueError: + return str(uuid.uuid5(_WORKSPACE_IDENTITY_NAMESPACE, value)) + def _table_names(conn: sa.Connection) -> set[str]: return set(sa.inspect(conn).get_table_names()) @@ -403,7 +414,7 @@ def _bootstrap_default_workspace(conn: sa.Connection) -> None: .values(created_by_account_uuid=owner_account_uuid) ) else: - workspace_uuid = str(uuid.uuid4()) + workspace_uuid = _workspace_uuid_from_instance_id(instance_uuid) conn.execute( workspaces.insert().values( uuid=workspace_uuid, diff --git a/src/langbot/pkg/persistence/alembic/versions/0017_oss_workspace_identity.py b/src/langbot/pkg/persistence/alembic/versions/0017_oss_workspace_identity.py new file mode 100644 index 000000000..5574c96c1 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0017_oss_workspace_identity.py @@ -0,0 +1,177 @@ +"""align the OSS Workspace UUID with the persisted instance identity + +Revision ID: 0017_oss_workspace_identity +Revises: 0016_support_admin_sessions +Create Date: 2026-07-31 +""" + +from __future__ import annotations + +import uuid + +import sqlalchemy as sa +from alembic import op + + +revision = '0017_oss_workspace_identity' +down_revision = '0016_support_admin_sessions' +branch_labels = None +depends_on = None + +_WORKSPACE_IDENTITY_NAMESPACE = uuid.UUID('8ea04f29-8528-4cc3-bb28-30a838c89d76') +_OSS_WORKSPACE_METADATA_KEY = 'oss_workspace_uuid' + + +def _workspace_uuid_from_instance_id(instance_id: str) -> str: + value = instance_id.strip() + candidate = value[len('instance_') :] if value.startswith('instance_') else value + try: + return str(uuid.UUID(candidate)) + except ValueError: + return str(uuid.uuid5(_WORKSPACE_IDENTITY_NAMESPACE, value)) + + +def _quote(conn: sa.Connection, identifier: str) -> str: + return conn.dialect.identifier_preparer.quote(identifier) + + +def _defer_foreign_keys(conn: sa.Connection, inspector: sa.Inspector, table_names: list[str]) -> None: + """Allow the transaction to re-key a connected tenant graph atomically.""" + + if conn.dialect.name == 'sqlite': + conn.execute(sa.text('PRAGMA defer_foreign_keys = ON')) + return + if conn.dialect.name != 'postgresql': + raise RuntimeError(f'Unsupported Workspace identity migration dialect: {conn.dialect.name}') + + for table_name in table_names: + for foreign_key in inspector.get_foreign_keys(table_name): + constraint_name = foreign_key.get('name') + if not constraint_name: + continue + conn.execute( + sa.text( + f'ALTER TABLE {_quote(conn, table_name)} ' + f'ALTER CONSTRAINT {_quote(conn, constraint_name)} DEFERRABLE INITIALLY DEFERRED' + ) + ) + + +def _suspend_postgres_rls( + conn: sa.Connection, + table_names: list[str], +) -> dict[str, tuple[bool, bool]]: + if conn.dialect.name != 'postgresql': + return {} + + states: dict[str, tuple[bool, bool]] = {} + for table_name in table_names: + row = conn.execute( + sa.text( + 'SELECT relrowsecurity, relforcerowsecurity ' + 'FROM pg_class WHERE oid = to_regclass(:table_name)' + ), + {'table_name': table_name}, + ).one() + enabled, forced = bool(row.relrowsecurity), bool(row.relforcerowsecurity) + states[table_name] = (enabled, forced) + table = _quote(conn, table_name) + if forced: + conn.execute(sa.text(f'ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY')) + if enabled: + conn.execute(sa.text(f'ALTER TABLE {table} DISABLE ROW LEVEL SECURITY')) + return states + + +def _restore_postgres_rls(conn: sa.Connection, states: dict[str, tuple[bool, bool]]) -> None: + for table_name, (enabled, forced) in states.items(): + table = _quote(conn, table_name) + if enabled: + conn.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY')) + if forced: + conn.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY')) + + +def upgrade() -> None: + conn = op.get_bind() + inspector = sa.inspect(conn) + table_names = inspector.get_table_names() + if 'workspaces' not in table_names: + return + + metadata = sa.MetaData() + workspaces = sa.Table('workspaces', metadata, autoload_with=conn) + local_rows = conn.execute(sa.select(workspaces).where(workspaces.c.source == 'local')).mappings().all() + if not local_rows: + return + if len(local_rows) != 1: + raise RuntimeError('Cannot align OSS Workspace identity: expected exactly one local Workspace') + + old_row = dict(local_rows[0]) + old_uuid = old_row['uuid'] + canonical_uuid = _workspace_uuid_from_instance_id(old_row['instance_uuid']) + if old_uuid == canonical_uuid: + return + if conn.execute(sa.select(workspaces.c.uuid).where(workspaces.c.uuid == canonical_uuid)).scalar_one_or_none(): + raise RuntimeError(f'Cannot align OSS Workspace identity: target {canonical_uuid!r} already exists') + + tenant_tables = [ + table_name + for table_name in table_names + if table_name == 'workspaces' + or 'workspace_uuid' in {column['name'] for column in inspector.get_columns(table_name)} + ] + rls_states = _suspend_postgres_rls(conn, tenant_tables) + try: + _defer_foreign_keys(conn, inspector, table_names) + + # Release local source/slug uniqueness while the canonical parent exists + # alongside the old parent for the duration of this transaction. + temporary_slug = f'__workspace_rekey__{old_uuid}' + conn.execute( + workspaces.update() + .where(workspaces.c.uuid == old_uuid) + .values(source='cloud_projection', slug=temporary_slug) + ) + new_row = dict(old_row) + new_row['uuid'] = canonical_uuid + conn.execute(workspaces.insert().values(**new_row)) + + for table_name in tenant_tables: + if table_name == 'workspaces': + continue + table = sa.Table(table_name, metadata, autoload_with=conn, extend_existing=True) + conn.execute( + table.update() + .where(table.c.workspace_uuid == old_uuid) + .values(workspace_uuid=canonical_uuid) + ) + + if 'metadata' in table_names: + conn.execute( + sa.text( + 'UPDATE metadata SET value = :canonical_uuid ' + 'WHERE key = :key AND value = :old_uuid' + ), + { + 'canonical_uuid': canonical_uuid, + 'key': _OSS_WORKSPACE_METADATA_KEY, + 'old_uuid': old_uuid, + }, + ) + conn.execute(workspaces.delete().where(workspaces.c.uuid == old_uuid)) + if conn.dialect.name == 'postgresql': + # Fire deferred FK triggers before ALTER TABLE restores RLS; PostgreSQL + # rejects ALTER TABLE while a relation has pending trigger events. + conn.execute(sa.text('SET CONSTRAINTS ALL IMMEDIATE')) + except Exception: + # Alembic owns the transaction. Rollback restores the transactional RLS DDL. + raise + else: + _restore_postgres_rls(conn, rls_states) + + +def downgrade() -> None: + # The previous random UUID is intentionally not recoverable. Keeping the + # canonical identity preserves every FK and is safe for older application code. + pass diff --git a/src/langbot/pkg/pipeline/process/handlers/chat.py b/src/langbot/pkg/pipeline/process/handlers/chat.py index ef6dd874d..d3dd4e79d 100644 --- a/src/langbot/pkg/pipeline/process/handlers/chat.py +++ b/src/langbot/pkg/pipeline/process/handlers/chat.py @@ -15,6 +15,7 @@ from ....provider import runner as runner_module import langbot_plugin.api.entities.events as events from ....utils import importutil, constants, runner as runner_utils from ....telemetry import features as telemetry_features +from ....telemetry.identity import workspace_identity from ....provider import runners import langbot_plugin.api.entities.builtin.provider.session as provider_session import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query @@ -265,7 +266,8 @@ class ChatMessageHandler(handler.MessageHandler): 'duration_ms': duration_ms, 'model_name': model_name, 'version': constants.semantic_version, - 'instance_id': constants.instance_id, + **workspace_identity(get_query_execution_context(query)), + 'runtime_instance_id': constants.instance_id, 'edition': constants.edition, 'pipeline_plugins': pipeline_plugins, 'features': features, diff --git a/src/langbot/pkg/telemetry/heartbeat.py b/src/langbot/pkg/telemetry/heartbeat.py index c535120ca..1296d9145 100644 --- a/src/langbot/pkg/telemetry/heartbeat.py +++ b/src/langbot/pkg/telemetry/heartbeat.py @@ -27,6 +27,18 @@ if typing.TYPE_CHECKING: HEARTBEAT_INTERVAL_SECONDS = 24 * 3600 +class WorkspaceResourceSnapshot(typing.TypedDict): + workspace_uuid: str + bot_count: int + pipeline_count: int + knowledge_base_count: int + plugin_count: int + mcp_server_count: int + extension_count: int + skill_count: int + adapters: list[str] + + async def _count( ap: core_app.Application, table, @@ -52,14 +64,13 @@ async def _count( return -1 -async def _cloud_workspace_resource_counts(ap: core_app.Application) -> list[dict]: +async def _cloud_workspace_resource_counts(ap: core_app.Application, bindings) -> list[WorkspaceResourceSnapshot]: """Summarize already-loaded Cloud registries without per-tenant SQL.""" persistence_mgr = ap.persistence_mgr if getattr(getattr(persistence_mgr, 'mode', None), 'value', None) != 'cloud_runtime': return [] - bindings = await ap.workspace_service.list_active_execution_bindings() - counts = { + counts: dict[str, WorkspaceResourceSnapshot] = { binding.workspace_uuid: { 'workspace_uuid': binding.workspace_uuid, 'bot_count': 0, @@ -68,13 +79,19 @@ async def _cloud_workspace_resource_counts(ap: core_app.Application) -> list[dic 'plugin_count': 0, 'mcp_server_count': 0, 'extension_count': 0, + 'skill_count': 0, + 'adapters': [], } for binding in bindings } - for key in getattr(ap.platform_mgr, '_bots_by_key', {}): + adapter_sets: dict[str, set[str]] = {workspace_uuid: set() for workspace_uuid in counts} + for key, bot in getattr(ap.platform_mgr, '_bots_by_key', {}).items(): if len(key) >= 2 and key[1] in counts: counts[key[1]]['bot_count'] += 1 + adapter = getattr(bot, 'adapter', None) + if adapter is not None and getattr(bot, 'enable', False): + adapter_sets[key[1]].add(adapter.__class__.__name__) for key in getattr(ap.pipeline_mgr, '_pipelines_by_key', {}): if len(key) >= 2 and key[1] in counts: counts[key[1]]['pipeline_count'] += 1 @@ -87,14 +104,23 @@ async def _cloud_workspace_resource_counts(ap: core_app.Application) -> list[dic for workspace_uuid, installations in getattr(ap.plugin_connector, '_workspace_installations', {}).items(): if workspace_uuid in counts: counts[workspace_uuid]['plugin_count'] = len(installations) + for key, skills in getattr(ap.skill_mgr, '_skills_by_scope', {}).items(): + if len(key) >= 2 and key[1] in counts: + counts[key[1]]['skill_count'] += len(skills) - for resource in counts.values(): + for workspace_uuid, resource in counts.items(): resource['extension_count'] = resource['plugin_count'] + resource['mcp_server_count'] + resource['adapters'] = sorted(adapter_sets[workspace_uuid]) return list(counts.values()) -async def build_heartbeat_payload(ap: core_app.Application) -> dict: - """Collect the anonymous instance profile snapshot.""" +async def build_heartbeat_payload( + ap: core_app.Application, + *, + workspace_uuid: str, + workspace_resource: WorkspaceResourceSnapshot | None = None, +) -> dict: + """Collect one anonymous Workspace profile snapshot.""" from ..entity.persistence import bot as persistence_bot from ..entity.persistence import mcp as persistence_mcp from ..entity.persistence import pipeline as persistence_pipeline @@ -177,15 +203,14 @@ async def build_heartbeat_payload(ap: core_app.Application) -> dict: except Exception: pass - workspace_resources = await _cloud_workspace_resource_counts(ap) - if workspace_resources: - features['workspace_resources'] = workspace_resources + if workspace_resource is not None: + features.update({key: value for key, value in workspace_resource.items() if key != 'workspace_uuid'}) return { 'event_type': 'instance_heartbeat', 'query_id': '', 'version': constants.semantic_version, - 'instance_id': constants.instance_id, + 'workspace_uuid': workspace_uuid, 'instance_create_ts': constants.instance_create_ts, 'edition': constants.edition, 'features': features, @@ -193,14 +218,34 @@ async def build_heartbeat_payload(ap: core_app.Application) -> dict: } +async def build_heartbeat_payloads(ap: core_app.Application) -> list[dict]: + """Build one heartbeat per active Workspace.""" + bindings = await ap.workspace_service.list_active_execution_bindings() + workspace_uuids = sorted({binding.workspace_uuid for binding in bindings}) + resources = { + resource['workspace_uuid']: resource for resource in await _cloud_workspace_resource_counts(ap, bindings) + } + return [ + await build_heartbeat_payload( + ap, + workspace_uuid=workspace_uuid, + workspace_resource=resources.get(workspace_uuid), + ) + for workspace_uuid in workspace_uuids + ] + + async def heartbeat_loop(ap: core_app.Application) -> None: """Send one heartbeat shortly after startup, then daily.""" # Small delay so managers (platform, skills, plugins) finish loading first await asyncio.sleep(30) while True: try: - payload = await build_heartbeat_payload(ap) - await ap.telemetry.start_send_task(payload) + for payload in await build_heartbeat_payloads(ap): + # Heartbeats are a daily bounded batch, not best-effort query events. + # Await each send so the TelemetryManager's 8-task queue cannot drop + # Workspaces after the first batch. + await ap.telemetry.send(payload) except Exception as e: try: ap.logger.debug(f'Telemetry heartbeat failed: {e}') diff --git a/src/langbot/pkg/telemetry/identity.py b/src/langbot/pkg/telemetry/identity.py new file mode 100644 index 000000000..8385ce190 --- /dev/null +++ b/src/langbot/pkg/telemetry/identity.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +import typing + + +class WorkspaceExecutionContext(typing.Protocol): + @property + def workspace_uuid(self) -> str: ... + + +def workspace_identity(execution_context: WorkspaceExecutionContext) -> dict[str, str]: + """Build the canonical telemetry identity for one Workspace execution.""" + workspace_uuid = execution_context.workspace_uuid.strip() + if not workspace_uuid: + raise ValueError("Telemetry execution Workspace UUID is empty") + return {"workspace_uuid": workspace_uuid} diff --git a/src/langbot/pkg/telemetry/telemetry.py b/src/langbot/pkg/telemetry/telemetry.py index 6c3b8ba29..28bcf1a0a 100644 --- a/src/langbot/pkg/telemetry/telemetry.py +++ b/src/langbot/pkg/telemetry/telemetry.py @@ -2,7 +2,11 @@ from __future__ import annotations import asyncio import contextlib +import os +import typing + import httpx + from ..core import app as core_app from ..utils import httpclient @@ -21,7 +25,7 @@ class TelemetryManager: def __init__(self, ap: core_app.Application): self.ap = ap - self.telemetry_config = {} + self.telemetry_config: dict[str, typing.Any] = {} self.send_tasks: list[asyncio.Task] = [] self._client: httpx.AsyncClient | None = None @@ -131,7 +135,16 @@ class TelemetryManager: async with self._client_context() as client: try: # Use asyncio.wait_for to ensure we always bound the total time - resp = await asyncio.wait_for(client.post(url, json=sanitized), timeout=10 + 1) + telemetry_token = os.getenv('LANGBOT_TELEMETRY_INGEST_TOKEN', '').strip() + if telemetry_token: + request = client.post( + url, + json=sanitized, + headers={'X-LangBot-Telemetry-Token': telemetry_token}, + ) + else: + request = client.post(url, json=sanitized) + resp = await asyncio.wait_for(request, timeout=10 + 1) if resp.status_code >= 400: body = await httpclient.response_text(resp, max_chars=200) @@ -143,7 +156,8 @@ class TelemetryManager: app_err = False try: j = await httpclient.parse_json_response(resp) - if isinstance(j, dict) and j.get('code') is not None and int(j.get('code')) >= 400: + app_code = j.get('code') if isinstance(j, dict) else None + if app_code is not None and int(app_code) >= 400: app_err = True self.ap.logger.warning( f'Telemetry post to {url} returned application error code {j.get("code")} - {j.get("msg")}' diff --git a/src/langbot/pkg/workspace/identity.py b/src/langbot/pkg/workspace/identity.py new file mode 100644 index 000000000..8500be798 --- /dev/null +++ b/src/langbot/pkg/workspace/identity.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +import uuid + + +_INSTANCE_PREFIX = "instance_" +_WORKSPACE_IDENTITY_NAMESPACE = uuid.UUID("8ea04f29-8528-4cc3-bb28-30a838c89d76") + + +def workspace_uuid_from_instance_id(instance_id: str) -> str: + """Return the stable OSS Workspace UUID for a persisted instance identity.""" + value = instance_id.strip() + if not value: + raise ValueError("LangBot instance identity is empty") + + candidate = value[len(_INSTANCE_PREFIX) :] if value.startswith(_INSTANCE_PREFIX) else value + try: + return str(uuid.UUID(candidate)) + except ValueError: + return str(uuid.uuid5(_WORKSPACE_IDENTITY_NAMESPACE, value)) diff --git a/src/langbot/pkg/workspace/service.py b/src/langbot/pkg/workspace/service.py index bab30d4f2..b4eb0a845 100644 --- a/src/langbot/pkg/workspace/service.py +++ b/src/langbot/pkg/workspace/service.py @@ -30,6 +30,7 @@ from .errors import ( WorkspaceOwnerAlreadyExistsError, ) from .entities import WorkspaceExecutionBinding +from .identity import workspace_uuid_from_instance_id from .policy import CloudWorkspacePolicy, SingleWorkspacePolicy from .repository import WorkspaceRepository @@ -497,7 +498,7 @@ class WorkspaceService: created_by_account_uuid: str | None = None, ) -> Workspace: return Workspace( - uuid=str(uuid.uuid4()), + uuid=workspace_uuid_from_instance_id(self.instance_uuid), instance_uuid=self.instance_uuid, name=name, slug=slug, diff --git a/tests/integration/persistence/test_workspace_migration.py b/tests/integration/persistence/test_workspace_migration.py index 4fc56eba5..70c46f89a 100644 --- a/tests/integration/persistence/test_workspace_migration.py +++ b/tests/integration/persistence/test_workspace_migration.py @@ -22,6 +22,7 @@ from langbot.pkg.persistence.alembic_runner import ( from langbot.pkg.utils import constants from langbot.pkg.utils import importutil from langbot.pkg.workspace.collaboration import normalize_email +from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id pytestmark = [pytest.mark.integration, pytest.mark.asyncio] @@ -109,6 +110,7 @@ async def test_legacy_instance_gets_stable_accounts_and_default_workspace(legacy .mappings() .one() ) + assert workspace['uuid'] == workspace_uuid_from_instance_id('instance_migration_test') assert workspace['instance_uuid'] == 'instance_migration_test' assert workspace['slug'] == 'default' assert workspace['status'] == 'active' @@ -149,6 +151,53 @@ async def test_workspace_upgrade_is_idempotent_and_preserves_identifiers(legacy_ assert workspace_uuid_after == workspace_uuid_before +async def test_existing_oss_workspace_is_rekeyed_to_instance_identity(tmp_path): + engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-rekey.db"}') + instance_id = 'instance_a711d9e4-0953-443f-a0e9-7dd50193a79f' + old_workspace_uuid = '11111111-1111-4111-8111-111111111111' + canonical_uuid = workspace_uuid_from_instance_id(instance_id) + schema = sa.MetaData() + sa.Table( + 'metadata', + schema, + sa.Column('key', sa.String(255), primary_key=True), + sa.Column('value', sa.String(255)), + ) + sa.Table( + 'workspaces', + schema, + sa.Column('uuid', sa.String(36), primary_key=True), + sa.Column('instance_uuid', sa.String(255), nullable=False), + sa.Column('slug', sa.String(255), nullable=False), + sa.Column('source', sa.String(32), nullable=False), + ) + sa.Table( + 'tenant_rows', + schema, + sa.Column('id', sa.Integer, primary_key=True), + sa.Column('workspace_uuid', sa.String(36), sa.ForeignKey('workspaces.uuid'), nullable=False), + ) + async with engine.begin() as conn: + await conn.run_sync(schema.create_all) + await conn.execute(sa.text("INSERT INTO metadata (key, value) VALUES ('instance_uuid', :value)"), {'value': instance_id}) + await conn.execute( + sa.text("INSERT INTO workspaces (uuid, instance_uuid, slug, source) VALUES (:uuid, :instance, 'default', 'local')"), + {'uuid': old_workspace_uuid, 'instance': instance_id}, + ) + await conn.execute( + sa.text("INSERT INTO tenant_rows (id, workspace_uuid) VALUES (1, :uuid)"), + {'uuid': old_workspace_uuid}, + ) + await run_alembic_stamp(engine, '0016_support_admin_sessions') + + await run_alembic_upgrade(engine, 'head') + + async with engine.connect() as conn: + assert (await conn.execute(sa.text("SELECT uuid FROM workspaces"))).scalar_one() == canonical_uuid + assert (await conn.execute(sa.text("SELECT workspace_uuid FROM tenant_rows"))).scalar_one() == canonical_uuid + await engine.dispose() + + async def test_workspace_kernel_upgrade_downgrade_upgrade_round_trip(tmp_path): engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-round-trip.db"}') try: @@ -362,6 +411,47 @@ async def test_persistence_startup_defers_workspace_tables_until_account_upgrade await engine.dispose() +async def test_oss_workspace_identity_rekeys_fk_graph_and_metadata(tmp_path): + engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-rekey.db"}') + try: + await _create_legacy_schema(engine) + await run_alembic_stamp(engine, '0008_mcp_resource_prefs') + await run_alembic_upgrade(engine, '0016_support_admin_sessions') + + async with engine.begin() as conn: + old_uuid = await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'")) + instance_uuid = await conn.scalar(sa.text("SELECT instance_uuid FROM workspaces WHERE source = 'local'")) + assert old_uuid + assert instance_uuid + await conn.execute( + sa.text( + "INSERT INTO workspace_metadata (workspace_uuid, key, value) " + "VALUES (:workspace_uuid, 'migration_probe', 'present')" + ), + {'workspace_uuid': old_uuid}, + ) + await conn.execute( + sa.text( + "INSERT INTO metadata (key, value) VALUES ('oss_workspace_uuid', :workspace_uuid) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value" + ), + {'workspace_uuid': old_uuid}, + ) + + await run_alembic_upgrade(engine, 'head') + expected_uuid = workspace_uuid_from_instance_id(instance_uuid) + async with engine.connect() as conn: + assert await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'")) == expected_uuid + assert await conn.scalar( + sa.text("SELECT workspace_uuid FROM workspace_metadata WHERE key = 'migration_probe'") + ) == expected_uuid + assert await conn.scalar( + sa.text("SELECT value FROM metadata WHERE key = 'oss_workspace_uuid'") + ) == expected_uuid + finally: + await engine.dispose() + + async def test_persistence_startup_rejects_instance_uuid_drift(tmp_path, monkeypatch): engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "instance-drift.db"}') try: diff --git a/tests/unit_tests/telemetry/test_heartbeat.py b/tests/unit_tests/telemetry/test_heartbeat.py index 24c25b4ee..7a1b76484 100644 --- a/tests/unit_tests/telemetry/test_heartbeat.py +++ b/tests/unit_tests/telemetry/test_heartbeat.py @@ -60,10 +60,12 @@ class TestBuildHeartbeatPayload: async def test_payload_shape(self): heartbeat = get_heartbeat_module() ap = make_app() - payload = await heartbeat.build_heartbeat_payload(ap) + payload = await heartbeat.build_heartbeat_payload(ap, workspace_uuid='workspace-a') assert payload['event_type'] == 'instance_heartbeat' assert payload['query_id'] == '' + assert payload['workspace_uuid'] == 'workspace-a' + assert 'instance_id' not in payload assert 'instance_create_ts' in payload assert 'timestamp' in payload f = payload['features'] @@ -86,7 +88,7 @@ class TestBuildHeartbeatPayload: @pytest.mark.asyncio async def test_payload_is_json_serializable(self): heartbeat = get_heartbeat_module() - payload = await heartbeat.build_heartbeat_payload(make_app()) + payload = await heartbeat.build_heartbeat_payload(make_app(), workspace_uuid='workspace-a') json.dumps(payload) @pytest.mark.asyncio @@ -94,7 +96,7 @@ class TestBuildHeartbeatPayload: heartbeat = get_heartbeat_module() ap = make_app() ap.persistence_mgr.execute_async = AsyncMock(side_effect=RuntimeError('db down')) - payload = await heartbeat.build_heartbeat_payload(ap) + payload = await heartbeat.build_heartbeat_payload(ap, workspace_uuid='workspace-a') assert payload['features']['pipeline_count'] == -1 @pytest.mark.asyncio @@ -111,8 +113,10 @@ class TestBuildHeartbeatPayload: ('instance-a', 'workspace-a', 'pipeline-b'): object(), }, ) + adapter_a = type('WorkspaceAAdapter', (), {})() + adapter_b = type('WorkspaceBAdapter', (), {})() ap.platform_mgr._bots_by_key = { - ('instance-a', 'workspace-a', 'bot-a'): object(), + ('instance-a', 'workspace-a', 'bot-a'): SimpleNamespace(enable=True, adapter=adapter_a), } ap.tool_mgr = SimpleNamespace( mcp_tool_loader=SimpleNamespace( @@ -129,35 +133,46 @@ class TestBuildHeartbeatPayload: ap.plugin_connector._workspace_installations = { 'workspace-a': {'plugin-a', 'plugin-b'}, } + ap.skill_mgr._skills_by_scope = { + ('instance-a', 'workspace-a', 1): {'skill-a': {}, 'skill-b': {}}, + ('instance-a', 'workspace-b', 1): {'skill-c': {}}, + } ap.workspace_service.list_active_execution_bindings = AsyncMock( - return_value=[SimpleNamespace(workspace_uuid='workspace-a')], + return_value=[ + SimpleNamespace(workspace_uuid='workspace-a'), + SimpleNamespace(workspace_uuid='workspace-b'), + ], + ) + ap.platform_mgr._bots_by_key[('instance-a', 'workspace-b', 'bot-b')] = SimpleNamespace( + enable=True, adapter=adapter_b ) - payload = await heartbeat.build_heartbeat_payload(ap) + payloads = await heartbeat.build_heartbeat_payloads(ap) - features = payload['features'] - assert features['pipeline_count'] == 2 - assert features['mcp_server_count'] == 3 - assert features['knowledge_base_count'] == 1 - assert features['bot_count'] == 1 - assert features['workspace_resources'] == [ - { - 'workspace_uuid': 'workspace-a', - 'bot_count': 1, - 'pipeline_count': 2, - 'knowledge_base_count': 1, - 'plugin_count': 2, - 'mcp_server_count': 3, - 'extension_count': 5, - } - ] + assert [payload['workspace_uuid'] for payload in payloads] == ['workspace-a', 'workspace-b'] + assert all('instance_id' not in payload for payload in payloads) + by_workspace = {payload['workspace_uuid']: payload['features'] for payload in payloads} + assert by_workspace['workspace-a']['pipeline_count'] == 2 + assert by_workspace['workspace-a']['mcp_server_count'] == 3 + assert by_workspace['workspace-a']['knowledge_base_count'] == 1 + assert by_workspace['workspace-a']['bot_count'] == 1 + assert by_workspace['workspace-a']['plugin_count'] == 2 + assert by_workspace['workspace-a']['extension_count'] == 5 + assert by_workspace['workspace-a']['skill_count'] == 2 + assert by_workspace['workspace-a']['adapters'] == ['WorkspaceAAdapter'] + assert by_workspace['workspace-b']['bot_count'] == 1 + assert by_workspace['workspace-b']['pipeline_count'] == 0 + assert by_workspace['workspace-b']['skill_count'] == 1 + assert by_workspace['workspace-b']['adapters'] == ['WorkspaceBAdapter'] + assert 'workspace_resources' not in by_workspace['workspace-a'] ap.persistence_mgr.execute_async.assert_not_awaited() + ap.workspace_service.list_active_execution_bindings.assert_awaited_once() @pytest.mark.asyncio async def test_no_user_content_fields(self): """The heartbeat must never carry message content / credentials keys.""" heartbeat = get_heartbeat_module() - payload = await heartbeat.build_heartbeat_payload(make_app()) + payload = await heartbeat.build_heartbeat_payload(make_app(), workspace_uuid='workspace-a') flat = json.dumps(payload).lower() for forbidden in ('api_key', 'password', 'token', 'message_content'): assert forbidden not in flat diff --git a/tests/unit_tests/telemetry/test_telemetry.py b/tests/unit_tests/telemetry/test_telemetry.py index b15a989ee..2242a1aa4 100644 --- a/tests/unit_tests/telemetry/test_telemetry.py +++ b/tests/unit_tests/telemetry/test_telemetry.py @@ -569,6 +569,33 @@ class TestHTTPScenarios: await manager.send({'query_id': 'test'}) +class TestTelemetryManagedRuntimeAuthentication: + @pytest.mark.asyncio + async def test_send_includes_managed_runtime_token_header(self): + telemetry = get_telemetry_module() + mock_app = Mock() + mock_app.logger = Mock() + manager = telemetry.TelemetryManager(mock_app) + manager.telemetry_config = {'url': 'https://example.com'} + captured = {} + + async def mock_post(url, json, headers): + captured['headers'] = headers + return Mock(status_code=200, text='', json=Mock(return_value={'code': 0})) + + mock_client = Mock() + mock_client.post = mock_post + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + with ( + patch.dict('os.environ', {'LANGBOT_TELEMETRY_INGEST_TOKEN': 'managed-runtime-secret'}), + patch.object(httpx, 'AsyncClient', return_value=mock_client), + ): + await manager.send({'event_type': 'instance_heartbeat'}) + + assert captured['headers'] == {'X-LangBot-Telemetry-Token': 'managed-runtime-secret'} + + class TestStartSendTask: """Tests for start_send_task() method.""" diff --git a/tests/unit_tests/telemetry/test_workspace_identity.py b/tests/unit_tests/telemetry/test_workspace_identity.py new file mode 100644 index 000000000..efbd68a11 --- /dev/null +++ b/tests/unit_tests/telemetry/test_workspace_identity.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import uuid +from types import SimpleNamespace + + +def test_standard_oss_instance_id_aligns_to_embedded_uuid(): + from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id + + instance_uuid = "a711d9e4-0953-443f-a0e9-7dd50193a79f" + + assert workspace_uuid_from_instance_id(instance_uuid) == instance_uuid + assert workspace_uuid_from_instance_id(f"instance_{instance_uuid}") == instance_uuid + + +def test_custom_legacy_instance_id_maps_to_stable_valid_uuid(): + from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id + + first = workspace_uuid_from_instance_id("instance_migration_test") + second = workspace_uuid_from_instance_id("instance_migration_test") + + assert first == second + assert str(uuid.UUID(first)) == first + + +def test_query_telemetry_identity_uses_execution_workspace_only(): + from langbot.pkg.telemetry.identity import workspace_identity + + identity = workspace_identity(SimpleNamespace(workspace_uuid="workspace-a", instance_uuid="instance-a")) + + assert identity == {"workspace_uuid": "workspace-a"} diff --git a/tests/unit_tests/workspace/test_workspace_service.py b/tests/unit_tests/workspace/test_workspace_service.py index 8657998a3..f24f2ca44 100644 --- a/tests/unit_tests/workspace/test_workspace_service.py +++ b/tests/unit_tests/workspace/test_workspace_service.py @@ -26,6 +26,7 @@ from langbot.pkg.workspace import ( WorkspaceOwnerAlreadyExistsError, WorkspaceService, ) +from langbot.pkg.workspace.identity import workspace_uuid_from_instance_id from langbot.pkg.workspace.policy import CloudWorkspacePolicy @@ -113,6 +114,7 @@ async def test_ensure_singleton_workspace_is_idempotent(workspace_test_context): first = await service.ensure_singleton_workspace() second = await service.ensure_singleton_workspace() + assert first.uuid == workspace_uuid_from_instance_id('instance_service_test') assert second.uuid == first.uuid async with session_factory() as session: assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(Workspace)) == 1