mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-31 16:46:07 +00:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 22de4b14ee | |||
| 404e3466d9 |
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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}')
|
||||
|
||||
@@ -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}
|
||||
@@ -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")}'
|
||||
|
||||
@@ -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))
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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"}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user