mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
feat: use workspace identity for telemetry (#2371)
Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user