mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 12:40:59 +00:00
Merge remote-tracking branch 'origin/master' into dev/4.11.x
# Conflicts: # src/langbot/pkg/pipeline/process/handlers/chat.py # tests/integration/persistence/test_workspace_migration.py
This commit is contained in:
@@ -74,6 +74,11 @@ class AuthorizationError(Exception):
|
||||
error_code = 'forbidden'
|
||||
|
||||
|
||||
class AuthenticationDeniedError(AuthorizationError):
|
||||
status_code = 401
|
||||
error_code = 'invalid_authentication'
|
||||
|
||||
|
||||
class WorkspaceRequiredError(AuthorizationError):
|
||||
status_code = 400
|
||||
error_code = 'workspace_required'
|
||||
|
||||
@@ -9,6 +9,7 @@ class PrincipalType(enum.StrEnum):
|
||||
|
||||
ACCOUNT = 'account'
|
||||
API_KEY = 'api_key'
|
||||
SUPPORT_ADMIN = 'support_admin'
|
||||
SYSTEM = 'system'
|
||||
PUBLIC_BOT = 'public_bot'
|
||||
|
||||
@@ -19,7 +20,9 @@ class PrincipalContext:
|
||||
|
||||
principal_type: PrincipalType
|
||||
account_uuid: str | None = None
|
||||
actor_account_uuid: str | None = None
|
||||
api_key_uuid: str | None = None
|
||||
support_session_id: str | None = None
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
|
||||
@@ -15,8 +15,17 @@ from ....workspace.collaboration import MembershipPermissionError, WorkspaceColl
|
||||
from ....workspace.errors import WorkspaceNotFoundError
|
||||
from ....cloud.entitlements import EntitlementUnavailableError
|
||||
from ....core.errors import TaskCapacityError
|
||||
from ..authz import AuthorizationError, Permission, permissions_for_role, require_permission
|
||||
from ..authz import (
|
||||
AuthenticationDeniedError,
|
||||
AuthorizationError,
|
||||
Permission,
|
||||
PermissionDeniedError,
|
||||
WorkspaceRequiredError,
|
||||
permissions_for_role,
|
||||
require_permission,
|
||||
)
|
||||
from ..context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
|
||||
from ....cloud.support_admin import SupportAdminSessionError
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from ....core.app import Application
|
||||
@@ -51,6 +60,17 @@ class AuthType(enum.Enum):
|
||||
USER_TOKEN_OR_API_KEY = 'user-token-or-api-key'
|
||||
|
||||
|
||||
_SUPPORT_ADMIN_DENIED_PERMISSIONS = frozenset(
|
||||
{
|
||||
Permission.OWNER_TRANSFER.value,
|
||||
Permission.MEMBER_VIEW.value,
|
||||
Permission.MEMBER_INVITE.value,
|
||||
Permission.MEMBER_UPDATE_ROLE.value,
|
||||
Permission.MEMBER_REMOVE.value,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class RouterGroup(abc.ABC):
|
||||
name: str
|
||||
|
||||
@@ -95,6 +115,10 @@ class RouterGroup(abc.ABC):
|
||||
return self.http_status(401, -1, 'No valid user token provided')
|
||||
|
||||
try:
|
||||
if self._is_support_admin_token(token):
|
||||
raise AuthenticationDeniedError(
|
||||
'Support admin tokens cannot be refreshed or used on account endpoints'
|
||||
)
|
||||
account, user_email = await self._authenticate_account(token)
|
||||
# Account-token routes deliberately stop before Workspace
|
||||
# selection. They may bootstrap a selector, but cannot
|
||||
@@ -111,8 +135,13 @@ class RouterGroup(abc.ABC):
|
||||
return self.http_status(401, -1, 'No valid user token provided')
|
||||
|
||||
try:
|
||||
account, user_email = await self._authenticate_account(token)
|
||||
request_context = await self._resolve_account_context(account, auth_type)
|
||||
request_context = await self._authenticate_support_admin(token, auth_type)
|
||||
if request_context is not None:
|
||||
self._require_support_admin_route_allowed(rule, f, permission)
|
||||
user_email = None
|
||||
else:
|
||||
account, user_email = await self._authenticate_account(token)
|
||||
request_context = await self._resolve_account_context(account, auth_type)
|
||||
if permission is not None:
|
||||
if request_context is None:
|
||||
raise AuthorizationError('Workspace authorization is unavailable')
|
||||
@@ -141,10 +170,20 @@ class RouterGroup(abc.ABC):
|
||||
return self._auth_error_response(e)
|
||||
|
||||
elif auth_type == AuthType.USER_TOKEN_OR_API_KEY:
|
||||
token = quart.request.headers.get('Authorization', '').replace('Bearer ', '')
|
||||
if token and self._is_support_admin_token(token):
|
||||
try:
|
||||
request_context = await self._authenticate_support_admin(token, auth_type)
|
||||
if request_context is None:
|
||||
raise AuthenticationDeniedError('Invalid support admin token')
|
||||
self._require_support_admin_route_allowed(rule, f, permission)
|
||||
if permission is not None:
|
||||
require_permission(request_context, permission)
|
||||
self._inject_handler_context(f, kwargs, None, request_context)
|
||||
except Exception as e:
|
||||
return self._auth_error_response(e)
|
||||
# Try API key first (check X-API-Key header)
|
||||
api_key = quart.request.headers.get('X-API-Key', '')
|
||||
|
||||
if api_key:
|
||||
elif api_key := quart.request.headers.get('X-API-Key', ''):
|
||||
# API key authentication
|
||||
try:
|
||||
request_context = await self._authenticate_api_key(api_key, auth_type)
|
||||
@@ -155,8 +194,6 @@ class RouterGroup(abc.ABC):
|
||||
return self._auth_error_response(e)
|
||||
else:
|
||||
# Try user token authentication (Authorization header)
|
||||
token = quart.request.headers.get('Authorization', '').replace('Bearer ', '')
|
||||
|
||||
if not token:
|
||||
return self.http_status(
|
||||
401, -1, 'No valid authentication provided (user token or API key required)'
|
||||
@@ -268,10 +305,83 @@ class RouterGroup(abc.ABC):
|
||||
raise ValueError('User not found')
|
||||
return account, account.user
|
||||
|
||||
def _is_support_admin_token(self, token: str) -> bool:
|
||||
service = getattr(self.ap, 'support_admin_session_service', None)
|
||||
detector = getattr(service, 'is_support_admin_token', None)
|
||||
return callable(detector) and detector(token) is True
|
||||
|
||||
async def _authenticate_support_admin(
|
||||
self,
|
||||
token: str,
|
||||
auth_type: AuthType,
|
||||
*,
|
||||
workspace_uuid: str | None = None,
|
||||
request_id: str | None = None,
|
||||
) -> RequestContext | None:
|
||||
service = getattr(self.ap, 'support_admin_session_service', None)
|
||||
detector = getattr(service, 'is_support_admin_token', None)
|
||||
if service is None or not callable(detector) or detector(token) is not True:
|
||||
return None
|
||||
|
||||
requested_workspace_uuid = (
|
||||
workspace_uuid if workspace_uuid is not None else quart.request.headers.get('X-Workspace-Id')
|
||||
)
|
||||
if not requested_workspace_uuid:
|
||||
raise WorkspaceRequiredError('Support admin token requires an explicit Workspace selector')
|
||||
try:
|
||||
identity = await service.authenticate_token(
|
||||
token,
|
||||
requested_workspace_uuid=requested_workspace_uuid,
|
||||
)
|
||||
except SupportAdminSessionError as exc:
|
||||
raise AuthenticationDeniedError(str(exc)) from exc
|
||||
|
||||
entitlement_revision = await self._resolve_entitlement_revision(
|
||||
identity.instance_uuid,
|
||||
identity.workspace_uuid,
|
||||
)
|
||||
request_context = RequestContext(
|
||||
instance_uuid=identity.instance_uuid,
|
||||
placement_generation=identity.placement_generation,
|
||||
request_id=request_id or self.request_id(),
|
||||
auth_type=auth_type.value,
|
||||
principal=PrincipalContext(
|
||||
principal_type=PrincipalType.SUPPORT_ADMIN,
|
||||
actor_account_uuid=identity.actor_account_uuid,
|
||||
support_session_id=identity.grant_jti_hash,
|
||||
),
|
||||
workspace=WorkspaceContext(
|
||||
workspace_uuid=identity.workspace_uuid,
|
||||
membership_uuid=None,
|
||||
role='owner',
|
||||
permissions=permissions_for_role('owner') - _SUPPORT_ADMIN_DENIED_PERMISSIONS,
|
||||
membership_revision=0,
|
||||
),
|
||||
entitlement_revision=entitlement_revision,
|
||||
)
|
||||
quart.g.request_context = request_context
|
||||
quart.g.workspace_membership = None
|
||||
return request_context
|
||||
|
||||
@staticmethod
|
||||
def _require_support_admin_route_allowed(
|
||||
rule: str,
|
||||
handler: RouteCallable,
|
||||
permission: Permission | str | None,
|
||||
) -> None:
|
||||
parameters = inspect.signature(handler).parameters
|
||||
if rule.startswith('/api/v1/user/') or 'account' in parameters or 'user_email' in parameters:
|
||||
raise AuthenticationDeniedError('Support admin tokens are not permitted on account endpoints')
|
||||
permission_value = permission.value if isinstance(permission, Permission) else permission
|
||||
if permission_value in _SUPPORT_ADMIN_DENIED_PERMISSIONS:
|
||||
raise PermissionDeniedError(permission_value)
|
||||
|
||||
async def _resolve_account_context(
|
||||
self,
|
||||
account: typing.Any,
|
||||
auth_type: AuthType,
|
||||
*,
|
||||
token: str | None = None,
|
||||
) -> RequestContext | None:
|
||||
collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None)
|
||||
account_uuid = getattr(account, 'uuid', None)
|
||||
|
||||
@@ -97,6 +97,16 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
if not token or not workspace_uuid:
|
||||
raise ValueError('Authentication is required')
|
||||
|
||||
support_context = await self._authenticate_support_admin(
|
||||
token,
|
||||
group.AuthType.USER_TOKEN,
|
||||
workspace_uuid=workspace_uuid,
|
||||
request_id=quart.websocket.headers.get('X-Request-Id') or str(uuid.uuid4()),
|
||||
)
|
||||
if support_context is not None:
|
||||
require_permission(support_context, Permission.RUNTIME_OPERATE)
|
||||
return support_context, token
|
||||
|
||||
account, _ = await self._authenticate_account(token)
|
||||
account_uuid = getattr(account, 'uuid', None)
|
||||
collaboration_service = getattr(self.ap, 'workspace_collaboration_service', None)
|
||||
@@ -131,6 +141,23 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
) -> RequestContext:
|
||||
"""Recheck revocable account, membership, permission, and placement state."""
|
||||
|
||||
if request_context.principal.principal_type == PrincipalType.SUPPORT_ADMIN:
|
||||
current_context = await self._authenticate_support_admin(
|
||||
token,
|
||||
group.AuthType.USER_TOKEN,
|
||||
workspace_uuid=request_context.workspace_uuid,
|
||||
request_id=request_context.request_id,
|
||||
)
|
||||
if current_context is None or current_context.principal != request_context.principal:
|
||||
raise ValueError('WebSocket support admin session changed')
|
||||
if (
|
||||
current_context.instance_uuid != request_context.instance_uuid
|
||||
or current_context.placement_generation != request_context.placement_generation
|
||||
):
|
||||
raise ValueError('WebSocket authorization changed')
|
||||
require_permission(current_context, Permission.RUNTIME_OPERATE)
|
||||
return current_context
|
||||
|
||||
account, _ = await self._authenticate_account(token)
|
||||
account_uuid = getattr(account, 'uuid', None)
|
||||
if account_uuid != request_context.account_uuid:
|
||||
@@ -211,6 +238,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
scope=WebSocketScope.from_context(request_context),
|
||||
pipeline_uuid=pipeline_uuid,
|
||||
session_type=session_type,
|
||||
trigger_principal=request_context.principal,
|
||||
metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')},
|
||||
send_queue_size=(
|
||||
self.ap.instance_config.data.get('system', {})
|
||||
@@ -391,7 +419,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
)
|
||||
elif message_type == 'message':
|
||||
try:
|
||||
await self._revalidate_websocket_authorization(request_context, token)
|
||||
request_context = await self._revalidate_websocket_authorization(request_context, token)
|
||||
except Exception:
|
||||
await connection.send_queue.put({'type': 'error', 'message': 'Unauthorized'})
|
||||
break
|
||||
|
||||
@@ -22,6 +22,7 @@ class _AdapterSessionScope:
|
||||
principal_type: str
|
||||
account_uuid: str | None
|
||||
api_key_uuid: str | None
|
||||
support_session_id: str | None
|
||||
|
||||
@classmethod
|
||||
def from_request_context(cls, request_context: RequestContext) -> '_AdapterSessionScope':
|
||||
@@ -33,6 +34,7 @@ class _AdapterSessionScope:
|
||||
principal_type=principal.principal_type.value,
|
||||
account_uuid=principal.account_uuid,
|
||||
api_key_uuid=principal.api_key_uuid,
|
||||
support_session_id=principal.support_session_id,
|
||||
)
|
||||
|
||||
def matches(self, request_context: RequestContext) -> bool:
|
||||
|
||||
@@ -318,6 +318,13 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
)
|
||||
return await operation()
|
||||
|
||||
async def _require_authenticated_plugin_runtime_context(
|
||||
self,
|
||||
request_context: RequestContext,
|
||||
) -> ExecutionContext:
|
||||
"""Fence an authenticated resource request to its injected Workspace."""
|
||||
return await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
|
||||
async def _require_public_plugin_runtime_context(self) -> ExecutionContext:
|
||||
"""Resolve public assets only for the OSS singleton Workspace.
|
||||
|
||||
@@ -372,7 +379,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
permission=Permission.RESOURCE_VIEW,
|
||||
)
|
||||
async def _(request_context: RequestContext) -> str:
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
plugins = await self.ap.plugin_connector.list_plugins()
|
||||
|
||||
return self.success(data={'plugins': redact_plugin_secrets(plugins)})
|
||||
@@ -385,7 +392,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
)
|
||||
async def _(request_context: RequestContext) -> str:
|
||||
"""Get plugin debug information including debug URL and key"""
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
debug_info = await self.ap.plugin_connector.get_debug_info()
|
||||
|
||||
# Get debug URL from config
|
||||
@@ -428,7 +435,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
permission=Permission.RESOURCE_VIEW,
|
||||
)
|
||||
async def _(author: str, plugin_name: str, request_context: RequestContext) -> str:
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
plugin = await self.ap.plugin_connector.get_plugin_info(author, plugin_name)
|
||||
if plugin is None:
|
||||
return self.http_status(404, -1, 'plugin not found')
|
||||
@@ -469,7 +476,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
permission=Permission.RESOURCE_VIEW,
|
||||
)
|
||||
async def _(author: str, plugin_name: str, request_context: RequestContext) -> quart.Response:
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
plugin = await self.ap.plugin_connector.get_plugin_info(author, plugin_name)
|
||||
if plugin is None:
|
||||
return self.http_status(404, -1, 'plugin not found')
|
||||
@@ -489,7 +496,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
permission=Permission.RESOURCE_MANAGE,
|
||||
)
|
||||
async def _(author: str, plugin_name: str, request_context: RequestContext) -> quart.Response:
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
plugin = await self.ap.plugin_connector.get_plugin_info(author, plugin_name)
|
||||
if plugin is None:
|
||||
return self.http_status(404, -1, 'plugin not found')
|
||||
@@ -506,7 +513,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
)
|
||||
except ValueError as exc:
|
||||
return self.http_status(400, -1, str(exc))
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
await self.ap.plugin_connector.set_plugin_config(author, plugin_name, config)
|
||||
return self.success(data={})
|
||||
|
||||
@@ -517,7 +524,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
permission=Permission.RESOURCE_VIEW,
|
||||
)
|
||||
async def _(author: str, plugin_name: str, request_context: RequestContext) -> quart.Response:
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
language = quart.request.args.get('language', 'en')
|
||||
readme = await self.ap.plugin_connector.get_plugin_readme(author, plugin_name, language=language)
|
||||
return self.success(data={'readme': readme})
|
||||
@@ -529,7 +536,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
permission=Permission.AUDIT_VIEW,
|
||||
)
|
||||
async def _(author: str, plugin_name: str, request_context: RequestContext) -> quart.Response:
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
try:
|
||||
limit = int(quart.request.args.get('limit', 200))
|
||||
except (TypeError, ValueError):
|
||||
@@ -538,6 +545,44 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
logs = await self.ap.plugin_connector.get_plugin_logs(author, plugin_name, limit=limit, level=level)
|
||||
return self.success(data={'logs': logs})
|
||||
|
||||
@self.route(
|
||||
'/<author>/<plugin_name>/authenticated-icon',
|
||||
methods=['GET'],
|
||||
auth_type=group.AuthType.USER_TOKEN_OR_API_KEY,
|
||||
permission=Permission.RESOURCE_VIEW,
|
||||
)
|
||||
async def _(
|
||||
author: str,
|
||||
plugin_name: str,
|
||||
request_context: RequestContext,
|
||||
) -> quart.Response:
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
icon_data = await self.ap.plugin_connector.get_plugin_icon(author, plugin_name)
|
||||
icon_bytes = await asyncio.to_thread(base64.b64decode, icon_data['plugin_icon_base64'])
|
||||
return quart.Response(icon_bytes, mimetype=icon_data['mime_type'])
|
||||
|
||||
@self.route(
|
||||
'/<author>/<plugin_name>/authenticated-assets/<path:filepath>',
|
||||
methods=['GET'],
|
||||
auth_type=group.AuthType.USER_TOKEN_OR_API_KEY,
|
||||
permission=Permission.RESOURCE_VIEW,
|
||||
)
|
||||
async def _(
|
||||
author: str,
|
||||
plugin_name: str,
|
||||
filepath: str,
|
||||
request_context: RequestContext,
|
||||
) -> quart.Response:
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
asset_path = _normalize_plugin_asset_path(filepath)
|
||||
if asset_path is None:
|
||||
return quart.Response('Asset not found', status=404)
|
||||
asset_data = await self.ap.plugin_connector.get_plugin_assets(author, plugin_name, asset_path)
|
||||
if not asset_data.get('asset_base64'):
|
||||
return quart.Response('Asset not found', status=404)
|
||||
asset_bytes = await asyncio.to_thread(base64.b64decode, asset_data['asset_base64'])
|
||||
return quart.Response(asset_bytes, mimetype=asset_data['mime_type'])
|
||||
|
||||
@self.route(
|
||||
'/<author>/<plugin_name>/icon',
|
||||
methods=['GET'],
|
||||
@@ -596,7 +641,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
)
|
||||
async def _(author: str, plugin_name: str, request_context: RequestContext) -> str:
|
||||
"""Forward a page API request to the plugin."""
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
data = await quart.request.json
|
||||
if not isinstance(data, dict):
|
||||
return self.http_status(400, -1, 'invalid request body')
|
||||
@@ -625,7 +670,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
)
|
||||
async def _(request_context: RequestContext) -> str:
|
||||
"""Get releases from a GitHub repository URL"""
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
data = await quart.request.json
|
||||
repo_url = data.get('repo_url', '')
|
||||
|
||||
@@ -705,7 +750,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
)
|
||||
async def _(request_context: RequestContext) -> str:
|
||||
"""Get assets from a specific GitHub release"""
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
data = await quart.request.json
|
||||
owner = data.get('owner', '')
|
||||
repo = data.get('repo', '')
|
||||
@@ -901,7 +946,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
permission=Permission.RESOURCE_MANAGE,
|
||||
)
|
||||
async def _(request_context: RequestContext) -> str:
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
file = (await quart.request.files).get('file')
|
||||
if file is None:
|
||||
return self.http_status(400, -1, 'file is required')
|
||||
@@ -942,7 +987,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
)
|
||||
async def _(request_context: RequestContext) -> str:
|
||||
"""Upload a file for plugin configuration"""
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
file = (await quart.request.files).get('file')
|
||||
if file is None:
|
||||
return self.http_status(400, -1, 'file is required')
|
||||
@@ -974,7 +1019,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
)
|
||||
async def _(file_key: str, request_context: RequestContext) -> str:
|
||||
"""Delete a plugin configuration file"""
|
||||
await self.ap.plugin_connector.require_workspace_context(request_context)
|
||||
await self._require_authenticated_plugin_runtime_context(request_context)
|
||||
if not self.ap.storage_mgr.is_scoped_object_key(file_key, expected_owner_type='plugin_config'):
|
||||
return self.http_status(400, -1, 'invalid file key')
|
||||
|
||||
|
||||
@@ -396,6 +396,19 @@ class UserRouterGroup(group.RouterGroup):
|
||||
launch_assertion,
|
||||
expected_workspace_uuid=workspace_uuid,
|
||||
)
|
||||
if launch.get('launch_mode') == 'support_admin':
|
||||
token = launch.get('support_admin_token')
|
||||
if not token:
|
||||
raise SpaceLaunchError('Support admin launch session was not issued')
|
||||
return self.success(
|
||||
data={
|
||||
'token': token,
|
||||
'workspace_uuid': launch['workspace_uuid'],
|
||||
'principal_type': 'support_admin',
|
||||
'actor_account_uuid': launch['actor_account_uuid'],
|
||||
}
|
||||
)
|
||||
|
||||
account = await self.ap.user_service.get_user_by_uuid(launch['account_uuid'])
|
||||
if account is None:
|
||||
raise SpaceLaunchError('Launch Account is not projected into Core')
|
||||
|
||||
@@ -5,7 +5,7 @@ import typing
|
||||
import quart
|
||||
|
||||
from ...authz import Permission, permissions_for_role
|
||||
from ...context import RequestContext
|
||||
from ...context import PrincipalType, RequestContext
|
||||
from ...service.user import AccountExistsLoginRequiredError, ControlPlaneDirectoryRequiredError
|
||||
from .....entity.persistence.workspace import Workspace, WorkspaceInvitation, WorkspaceMembership
|
||||
from .....entity.persistence.workspace import WorkspaceSource
|
||||
@@ -120,9 +120,6 @@ class WorkspacesRouterGroup(group.RouterGroup):
|
||||
@self.route('/current', methods=['GET'], permission=Permission.WORKSPACE_VIEW)
|
||||
async def _(request_context: RequestContext) -> typing.Any:
|
||||
membership = quart.g.workspace_membership
|
||||
account = await self.ap.user_service.get_user_by_uuid(request_context.account_uuid)
|
||||
if account is None:
|
||||
return self.http_status(401, 'invalid_authentication', 'Account not found')
|
||||
workspace = await self.ap.workspace_service.get_workspace(request_context.workspace_uuid)
|
||||
plan_name: str | None = None
|
||||
resolver = getattr(self.ap, 'entitlement_resolver', None)
|
||||
@@ -132,6 +129,28 @@ class WorkspacesRouterGroup(group.RouterGroup):
|
||||
minimum_revision=request_context.entitlement_revision,
|
||||
)
|
||||
plan_name = entitlement.plan_name
|
||||
if request_context.principal.principal_type == PrincipalType.SUPPORT_ADMIN:
|
||||
return self.success(
|
||||
data={
|
||||
'workspace': _workspace_payload(workspace),
|
||||
'membership': {
|
||||
'uuid': None,
|
||||
'workspace_uuid': request_context.workspace_uuid,
|
||||
'account_uuid': None,
|
||||
'email': None,
|
||||
'role': 'owner',
|
||||
'status': 'active',
|
||||
'joined_at': None,
|
||||
'created_at': None,
|
||||
},
|
||||
'permissions': sorted(request_context.workspace.permissions),
|
||||
'placement_generation': request_context.placement_generation,
|
||||
'plan_name': plan_name,
|
||||
}
|
||||
)
|
||||
account = await self.ap.user_service.get_user_by_uuid(request_context.account_uuid)
|
||||
if account is None:
|
||||
return self.http_status(401, 'invalid_authentication', 'Account not found')
|
||||
return self.success(
|
||||
data={
|
||||
'workspace': _workspace_payload(workspace),
|
||||
|
||||
@@ -400,7 +400,10 @@ class UserService:
|
||||
|
||||
return await self.generate_jwt_token(user_obj)
|
||||
|
||||
async def generate_jwt_token(self, account: user.User | str) -> str:
|
||||
async def generate_jwt_token(
|
||||
self,
|
||||
account: user.User | str,
|
||||
) -> str:
|
||||
jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
|
||||
jwt_expire = self.ap.instance_config.data['system']['jwt']['expire']
|
||||
|
||||
@@ -413,7 +416,7 @@ class UserService:
|
||||
# Lightweight unit-test and bootstrap callers may not have persistence wired.
|
||||
account_obj = None
|
||||
|
||||
payload = {
|
||||
payload: dict[str, typing.Any] = {
|
||||
'user': user_email,
|
||||
'iss': self._jwt_identity()[0],
|
||||
'aud': self._jwt_identity()[1],
|
||||
|
||||
@@ -15,12 +15,15 @@ from cryptography.exceptions import InvalidSignature
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
|
||||
|
||||
from .support_admin import SupportAdminReplayError, SupportAdminSessionError, hash_grant_jti
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from ..core.app import Application
|
||||
|
||||
|
||||
CONTROL_PLANE_TYP = 'langbot-control-plane+jwt'
|
||||
LAUNCH_KIND = 'workspace.launch'
|
||||
SUPPORT_ADMIN_LAUNCH_KIND = 'workspace.support_admin_launch'
|
||||
EXPECTED_ISSUER = 'langbot-space'
|
||||
EXPECTED_AUDIENCE = 'langbot-cloud-runtime'
|
||||
_CONSUMED_JTI_MAX_ENTRIES = 4096
|
||||
@@ -123,15 +126,60 @@ class SpaceLaunchService:
|
||||
payload = claims.get('payload')
|
||||
if not isinstance(payload, dict):
|
||||
raise SpaceLaunchError('Launch assertion payload must be a JSON object')
|
||||
account_uuid = _required_string(payload, 'account_uuid')
|
||||
kind = _required_string(claims, 'kind')
|
||||
workspace_uuid = _required_string(payload, 'workspace_uuid')
|
||||
if expected_workspace_uuid is not None and workspace_uuid != expected_workspace_uuid:
|
||||
raise SpaceLaunchError('Launch assertion targets another Workspace')
|
||||
await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1))
|
||||
return {
|
||||
if kind == SUPPORT_ADMIN_LAUNCH_KIND:
|
||||
if 'account_uuid' in payload:
|
||||
raise SpaceLaunchError('Admin launch assertion must not identify a customer Account')
|
||||
if payload.get('launch_mode') != 'support_admin' or payload.get('principal_type') != 'support_admin':
|
||||
raise SpaceLaunchError('Admin launch principal must be support_admin')
|
||||
actor_account_uuid = _required_string(payload, 'actor_account_uuid')
|
||||
if _required_string(payload, 'effective_role') != 'owner':
|
||||
raise SpaceLaunchError('Admin launch effective role must be owner')
|
||||
issued_at = _required_int(claims, 'iat')
|
||||
expires_at = _required_int(claims, 'exp', minimum=1)
|
||||
if expires_at - issued_at > 90:
|
||||
raise SpaceLaunchError('Admin launch assertion lifetime exceeds 90 seconds')
|
||||
grant_jti_hash = hash_grant_jti(_required_string(claims, 'jti'))
|
||||
result = {
|
||||
'workspace_uuid': workspace_uuid,
|
||||
'launch_mode': 'support_admin',
|
||||
'actor_account_uuid': actor_account_uuid,
|
||||
'effective_role': 'owner',
|
||||
'grant_jti_hash': grant_jti_hash,
|
||||
}
|
||||
support_service = getattr(self.ap, 'support_admin_session_service', None)
|
||||
if support_service is None or not callable(getattr(support_service, 'consume_launch_grant', None)):
|
||||
raise SpaceLaunchError('Durable support admin session service is unavailable')
|
||||
try:
|
||||
support_session = await support_service.consume_launch_grant(
|
||||
grant_jti_hash=grant_jti_hash,
|
||||
workspace_uuid=workspace_uuid,
|
||||
actor_account_uuid=actor_account_uuid,
|
||||
)
|
||||
except SupportAdminReplayError as exc:
|
||||
raise SpaceLaunchError('Launch assertion has already been consumed') from exc
|
||||
except SupportAdminSessionError as exc:
|
||||
raise SpaceLaunchError(str(exc)) from exc
|
||||
result['support_admin_token'] = support_session.token
|
||||
self.ap.logger.info(
|
||||
'cloud_support_admin_launch_consumed actor_account_uuid=%s workspace_uuid=%s',
|
||||
result['actor_account_uuid'],
|
||||
workspace_uuid,
|
||||
)
|
||||
return result
|
||||
|
||||
if payload.get('launch_mode') is not None:
|
||||
raise SpaceLaunchError('Launch assertion mode is unsupported')
|
||||
account_uuid = _required_string(payload, 'account_uuid')
|
||||
result = {
|
||||
'account_uuid': account_uuid,
|
||||
'workspace_uuid': workspace_uuid,
|
||||
}
|
||||
await self._consume_jti(_required_string(claims, 'jti'), _required_int(claims, 'exp', minimum=1))
|
||||
return result
|
||||
|
||||
def _verify_assertion(self, token: str) -> dict[str, typing.Any]:
|
||||
if not getattr(getattr(self.ap, 'deployment', None), 'multi_workspace_enabled', False):
|
||||
@@ -169,8 +217,9 @@ class SpaceLaunchService:
|
||||
raise SpaceLaunchError('Launch assertion subject targets another instance')
|
||||
if _required_string(claims, 'instance_uuid') != instance_uuid:
|
||||
raise SpaceLaunchError('Launch assertion instance UUID does not match this Core')
|
||||
if _required_string(claims, 'kind') != LAUNCH_KIND:
|
||||
raise SpaceLaunchError('Launch assertion kind is not workspace.launch')
|
||||
kind = _required_string(claims, 'kind')
|
||||
if kind not in {LAUNCH_KIND, SUPPORT_ADMIN_LAUNCH_KIND}:
|
||||
raise SpaceLaunchError('Launch assertion kind is not supported')
|
||||
|
||||
issued_at = _required_int(claims, 'iat')
|
||||
not_before = _required_int(claims, 'nbf')
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import datetime
|
||||
import hashlib
|
||||
import re
|
||||
import time
|
||||
import typing
|
||||
|
||||
import jwt
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from ..entity.persistence.support_admin import SupportAdminTemporarySession
|
||||
from ..workspace.errors import WorkspaceError
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from ..core.app import Application
|
||||
|
||||
|
||||
SUPPORT_ADMIN_TOKEN_TYP = 'langbot-support-admin+jwt'
|
||||
SUPPORT_ADMIN_TOKEN_KIND = 'support_admin.session'
|
||||
SUPPORT_ADMIN_EFFECTIVE_ROLE = 'owner'
|
||||
SUPPORT_ADMIN_MAX_TOKEN_SECONDS = 300
|
||||
_SHA256_HEX = re.compile(r'^[0-9a-f]{64}$')
|
||||
|
||||
|
||||
class SupportAdminSessionError(ValueError):
|
||||
"""Raised when a support-admin session or token is not admissible."""
|
||||
|
||||
|
||||
class SupportAdminReplayError(SupportAdminSessionError):
|
||||
"""Raised when a launch grant JTI has already been consumed."""
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class IssuedSupportAdminSession:
|
||||
token: str
|
||||
grant_jti_hash: str
|
||||
workspace_uuid: str
|
||||
actor_account_uuid: str
|
||||
issued_at: datetime.datetime
|
||||
expires_at: datetime.datetime
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class SupportAdminSessionIdentity:
|
||||
grant_jti_hash: str
|
||||
workspace_uuid: str
|
||||
actor_account_uuid: str
|
||||
instance_uuid: str
|
||||
placement_generation: int
|
||||
|
||||
|
||||
def hash_grant_jti(jti: str) -> str:
|
||||
return hashlib.sha256(jti.encode('utf-8')).hexdigest()
|
||||
|
||||
|
||||
class SupportAdminSessionService:
|
||||
"""Issue and validate temporary Workspace-scoped support-admin sessions."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ap: Application,
|
||||
*,
|
||||
wall_time: typing.Callable[[], float] = time.time,
|
||||
) -> None:
|
||||
self.ap = ap
|
||||
self._wall_time = wall_time
|
||||
|
||||
async def consume_launch_grant(
|
||||
self,
|
||||
*,
|
||||
grant_jti_hash: str,
|
||||
workspace_uuid: str,
|
||||
actor_account_uuid: str,
|
||||
) -> IssuedSupportAdminSession:
|
||||
self._validate_grant_hash(grant_jti_hash)
|
||||
if not workspace_uuid or not actor_account_uuid:
|
||||
raise SupportAdminSessionError('Support admin session requires an actor and Workspace')
|
||||
|
||||
issued_at = self._utcnow()
|
||||
expires_at = issued_at + datetime.timedelta(seconds=SUPPORT_ADMIN_MAX_TOKEN_SECONDS)
|
||||
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||
if not callable(tenant_uow):
|
||||
raise SupportAdminSessionError('Support admin sessions require tenant persistence')
|
||||
|
||||
try:
|
||||
async with tenant_uow(workspace_uuid) as uow:
|
||||
await self.ap.workspace_service.get_execution_binding(workspace_uuid, session=uow.session)
|
||||
uow.session.add(
|
||||
SupportAdminTemporarySession(
|
||||
grant_jti_hash=grant_jti_hash,
|
||||
workspace_uuid=workspace_uuid,
|
||||
actor_account_uuid=actor_account_uuid,
|
||||
issued_at=issued_at,
|
||||
expires_at=expires_at,
|
||||
)
|
||||
)
|
||||
await uow.session.flush()
|
||||
except IntegrityError as exc:
|
||||
raise SupportAdminReplayError('Launch assertion has already been consumed') from exc
|
||||
except WorkspaceError as exc:
|
||||
raise SupportAdminSessionError('Workspace is unavailable for support access') from exc
|
||||
|
||||
return IssuedSupportAdminSession(
|
||||
token=self._encode_token(
|
||||
grant_jti_hash=grant_jti_hash,
|
||||
workspace_uuid=workspace_uuid,
|
||||
actor_account_uuid=actor_account_uuid,
|
||||
issued_at=issued_at,
|
||||
expires_at=expires_at,
|
||||
),
|
||||
grant_jti_hash=grant_jti_hash,
|
||||
workspace_uuid=workspace_uuid,
|
||||
actor_account_uuid=actor_account_uuid,
|
||||
issued_at=issued_at,
|
||||
expires_at=expires_at,
|
||||
)
|
||||
|
||||
def is_support_admin_token(self, token: str) -> bool:
|
||||
"""Return True only for compact JWTs marked as support-admin tokens."""
|
||||
|
||||
if not isinstance(token, str) or token.count('.') != 2:
|
||||
return False
|
||||
try:
|
||||
header = jwt.get_unverified_header(token)
|
||||
except jwt.PyJWTError:
|
||||
return False
|
||||
if header.get('typ') == SUPPORT_ADMIN_TOKEN_TYP:
|
||||
return True
|
||||
try:
|
||||
payload = jwt.decode(token, options={'verify_signature': False})
|
||||
except jwt.PyJWTError:
|
||||
return False
|
||||
return payload.get('kind') == SUPPORT_ADMIN_TOKEN_KIND
|
||||
|
||||
async def authenticate_token(
|
||||
self,
|
||||
token: str,
|
||||
*,
|
||||
requested_workspace_uuid: str | None,
|
||||
) -> SupportAdminSessionIdentity:
|
||||
if not self.is_support_admin_token(token):
|
||||
raise SupportAdminSessionError('Not a support admin token')
|
||||
workspace_uuid = (requested_workspace_uuid or '').strip()
|
||||
if not workspace_uuid:
|
||||
raise SupportAdminSessionError('Support admin token requires an explicit Workspace selector')
|
||||
|
||||
jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
|
||||
try:
|
||||
payload = jwt.decode(
|
||||
token,
|
||||
jwt_secret,
|
||||
algorithms=['HS256'],
|
||||
issuer='langbot-core',
|
||||
audience=self._audience(workspace_uuid),
|
||||
options={'require': ['exp', 'iat', 'nbf', 'iss', 'aud']},
|
||||
)
|
||||
except jwt.PyJWTError as exc:
|
||||
raise SupportAdminSessionError('Invalid support admin token') from exc
|
||||
self._validate_payload(payload, workspace_uuid)
|
||||
grant_jti_hash = payload['grant_jti_hash']
|
||||
actor_account_uuid = payload['actor_account_uuid']
|
||||
|
||||
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||
if not callable(tenant_uow):
|
||||
raise SupportAdminSessionError('Support admin sessions require tenant persistence')
|
||||
|
||||
now = self._utcnow()
|
||||
async with tenant_uow(workspace_uuid) as uow:
|
||||
session = await uow.session.get(SupportAdminTemporarySession, grant_jti_hash)
|
||||
if (
|
||||
session is None
|
||||
or session.workspace_uuid != workspace_uuid
|
||||
or session.actor_account_uuid != actor_account_uuid
|
||||
or session.revoked_at is not None
|
||||
or session.expires_at <= now
|
||||
):
|
||||
raise SupportAdminSessionError('Support admin session is inactive')
|
||||
binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid, session=uow.session)
|
||||
session.last_used_at = now
|
||||
await uow.session.flush()
|
||||
|
||||
return SupportAdminSessionIdentity(
|
||||
grant_jti_hash=grant_jti_hash,
|
||||
workspace_uuid=workspace_uuid,
|
||||
actor_account_uuid=actor_account_uuid,
|
||||
instance_uuid=binding.instance_uuid,
|
||||
placement_generation=binding.placement_generation,
|
||||
)
|
||||
|
||||
async def revoke_session(self, grant_jti_hash: str, workspace_uuid: str) -> None:
|
||||
self._validate_grant_hash(grant_jti_hash)
|
||||
now = self._utcnow()
|
||||
async with self.ap.persistence_mgr.tenant_uow(workspace_uuid) as uow:
|
||||
row = await uow.session.get(SupportAdminTemporarySession, grant_jti_hash)
|
||||
if row is not None and row.revoked_at is None:
|
||||
row.revoked_at = now
|
||||
|
||||
def _encode_token(
|
||||
self,
|
||||
*,
|
||||
grant_jti_hash: str,
|
||||
workspace_uuid: str,
|
||||
actor_account_uuid: str,
|
||||
issued_at: datetime.datetime,
|
||||
expires_at: datetime.datetime,
|
||||
) -> str:
|
||||
jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
|
||||
payload: dict[str, typing.Any] = {
|
||||
'kind': SUPPORT_ADMIN_TOKEN_KIND,
|
||||
'iss': 'langbot-core',
|
||||
'aud': self._audience(workspace_uuid),
|
||||
'sub': f'support-admin:{actor_account_uuid}',
|
||||
'iat': issued_at,
|
||||
'nbf': issued_at,
|
||||
'exp': expires_at,
|
||||
'actor_account_uuid': actor_account_uuid,
|
||||
'workspace_uuid': workspace_uuid,
|
||||
'effective_role': SUPPORT_ADMIN_EFFECTIVE_ROLE,
|
||||
'grant_jti_hash': grant_jti_hash,
|
||||
}
|
||||
return jwt.encode(payload, jwt_secret, algorithm='HS256', headers={'typ': SUPPORT_ADMIN_TOKEN_TYP})
|
||||
|
||||
def _validate_payload(self, payload: dict[str, typing.Any], workspace_uuid: str) -> None:
|
||||
if payload.get('kind') != SUPPORT_ADMIN_TOKEN_KIND:
|
||||
raise SupportAdminSessionError('Invalid support admin token kind')
|
||||
if payload.get('workspace_uuid') != workspace_uuid:
|
||||
raise SupportAdminSessionError('Support admin session is scoped to another Workspace')
|
||||
if payload.get('effective_role') != SUPPORT_ADMIN_EFFECTIVE_ROLE:
|
||||
raise SupportAdminSessionError('Invalid support admin token role')
|
||||
actor_account_uuid = payload.get('actor_account_uuid')
|
||||
if not isinstance(actor_account_uuid, str) or not actor_account_uuid.strip():
|
||||
raise SupportAdminSessionError('Invalid support admin actor')
|
||||
grant_jti_hash = payload.get('grant_jti_hash')
|
||||
if not isinstance(grant_jti_hash, str) or not _SHA256_HEX.match(grant_jti_hash):
|
||||
raise SupportAdminSessionError('Invalid support admin grant')
|
||||
|
||||
def _audience(self, workspace_uuid: str) -> str:
|
||||
return f'langbot-support-admin:{self.ap.workspace_service.instance_uuid}:{workspace_uuid}'
|
||||
|
||||
@staticmethod
|
||||
def _validate_grant_hash(grant_jti_hash: str) -> None:
|
||||
if not _SHA256_HEX.match(grant_jti_hash):
|
||||
raise SupportAdminSessionError('Invalid support admin grant')
|
||||
|
||||
def _utcnow(self) -> datetime.datetime:
|
||||
return datetime.datetime.fromtimestamp(self._wall_time(), datetime.UTC).replace(tzinfo=None)
|
||||
@@ -53,6 +53,7 @@ from ..workspace import collaboration as workspace_collaboration_module
|
||||
from ..workspace import invitation_delivery as invitation_delivery_module
|
||||
from ..cloud import bootstrap as cloud_bootstrap_module
|
||||
from ..cloud import launch as cloud_launch_module
|
||||
from ..cloud import support_admin as cloud_support_admin_module
|
||||
from ..cloud import directory_projection as cloud_directory_projection_module
|
||||
from ..cloud import entitlements as cloud_entitlements_module
|
||||
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType
|
||||
@@ -141,6 +142,8 @@ class Application:
|
||||
|
||||
space_launch_service: cloud_launch_module.SpaceLaunchService = None
|
||||
|
||||
support_admin_session_service: cloud_support_admin_module.SupportAdminSessionService = None
|
||||
|
||||
deployment: cloud_bootstrap_module.OpenSourceDeployment | cloud_bootstrap_module.VerifiedCloudDeployment = None
|
||||
|
||||
deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None
|
||||
|
||||
@@ -44,6 +44,7 @@ from ...workspace import collaboration as workspace_collaboration_module
|
||||
from ...workspace import invitation_delivery as invitation_delivery_module
|
||||
from ...cloud import bootstrap as cloud_bootstrap
|
||||
from ...cloud import launch as cloud_launch_module
|
||||
from ...cloud import support_admin as cloud_support_admin_module
|
||||
from ...cloud.directory import directory_projection_limits_from_config
|
||||
from ...cloud.directory_projection import DirectoryProjectionService
|
||||
from ...cloud.entitlements import EntitlementResolver
|
||||
@@ -185,6 +186,7 @@ class BuildAppStage(stage.BootingStage):
|
||||
workspace_service_inst,
|
||||
)
|
||||
ap.invitation_delivery_service = invitation_delivery_module.InvitationDeliveryService(ap)
|
||||
ap.support_admin_session_service = cloud_support_admin_module.SupportAdminSessionService(ap)
|
||||
ap.space_launch_service = cloud_launch_module.SpaceLaunchService(ap)
|
||||
|
||||
user_service_inst = user_service.UserService(ap)
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy
|
||||
|
||||
from .base import Base
|
||||
|
||||
|
||||
class SupportAdminTemporarySession(Base):
|
||||
"""Temporary support-admin Workspace access session."""
|
||||
|
||||
__tablename__ = 'support_admin_temporary_sessions'
|
||||
|
||||
grant_jti_hash = sqlalchemy.Column(sqlalchemy.String(64), primary_key=True)
|
||||
workspace_uuid = sqlalchemy.Column(
|
||||
sqlalchemy.String(36),
|
||||
sqlalchemy.ForeignKey('workspaces.uuid', ondelete='CASCADE'),
|
||||
nullable=False,
|
||||
)
|
||||
actor_account_uuid = sqlalchemy.Column(sqlalchemy.String(36), nullable=False)
|
||||
issued_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
|
||||
expires_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=False)
|
||||
revoked_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
|
||||
last_used_at = sqlalchemy.Column(sqlalchemy.DateTime, nullable=True)
|
||||
|
||||
__table_args__ = (
|
||||
sqlalchemy.Index(
|
||||
'ix_support_admin_sessions_workspace_expiry',
|
||||
'workspace_uuid',
|
||||
'expires_at',
|
||||
),
|
||||
sqlalchemy.CheckConstraint(
|
||||
'length(grant_jti_hash) = 64',
|
||||
name='ck_support_admin_sessions_grant_jti_hash',
|
||||
),
|
||||
)
|
||||
@@ -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,88 @@
|
||||
"""add temporary support-admin sessions
|
||||
|
||||
Revision ID: 0016_support_admin_sessions
|
||||
Revises: 0015_cloud_core_collab
|
||||
Create Date: 2026-07-31
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
|
||||
revision = '0016_support_admin_sessions'
|
||||
down_revision = '0015_cloud_core_collab'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
_TABLE_NAME = 'support_admin_temporary_sessions'
|
||||
_POLICY_NAME = 'langbot_workspace_isolation'
|
||||
_TENANT_SETTING = 'langbot.workspace_uuid'
|
||||
|
||||
|
||||
def _setting(name: str) -> str:
|
||||
return f"NULLIF(current_setting('{name}', true), '')"
|
||||
|
||||
|
||||
def _quote(conn: sa.Connection, identifier: str) -> str:
|
||||
return conn.dialect.identifier_preparer.quote(identifier)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
existing_tables = set(sa.inspect(conn).get_table_names())
|
||||
if _TABLE_NAME not in existing_tables:
|
||||
op.create_table(
|
||||
_TABLE_NAME,
|
||||
sa.Column('grant_jti_hash', sa.String(64), nullable=False),
|
||||
sa.Column(
|
||||
'workspace_uuid',
|
||||
sa.String(36),
|
||||
sa.ForeignKey('workspaces.uuid', ondelete='CASCADE'),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column('actor_account_uuid', sa.String(36), nullable=False),
|
||||
sa.Column('issued_at', sa.DateTime(), nullable=False),
|
||||
sa.Column('expires_at', sa.DateTime(), nullable=False),
|
||||
sa.Column('revoked_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('last_used_at', sa.DateTime(), nullable=True),
|
||||
sa.CheckConstraint(
|
||||
'length(grant_jti_hash) = 64',
|
||||
name='ck_support_admin_sessions_grant_jti_hash',
|
||||
),
|
||||
sa.PrimaryKeyConstraint('grant_jti_hash'),
|
||||
)
|
||||
op.create_index(
|
||||
'ix_support_admin_sessions_workspace_expiry',
|
||||
_TABLE_NAME,
|
||||
['workspace_uuid', 'expires_at'],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
if conn.dialect.name != 'postgresql':
|
||||
return
|
||||
|
||||
table = _quote(conn, _TABLE_NAME)
|
||||
policy = _quote(conn, _POLICY_NAME)
|
||||
expression = f'workspace_uuid::text = {_setting(_TENANT_SETTING)}'
|
||||
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
|
||||
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
|
||||
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
|
||||
op.execute(
|
||||
sa.text(
|
||||
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
|
||||
f'USING ({expression}) WITH CHECK ({expression})'
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
if conn.dialect.name == 'postgresql':
|
||||
table = _quote(conn, _TABLE_NAME)
|
||||
policy = _quote(conn, _POLICY_NAME)
|
||||
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
|
||||
op.drop_index('ix_support_admin_sessions_workspace_expiry', table_name=_TABLE_NAME)
|
||||
op.drop_table(_TABLE_NAME)
|
||||
@@ -0,0 +1,167 @@
|
||||
"""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
|
||||
+4
-4
@@ -1,7 +1,7 @@
|
||||
"""repair an ownerless local Workspace after tenancy migration
|
||||
|
||||
Revision ID: 0017_local_owner_repair
|
||||
Revises: 0016_agent_workspace
|
||||
Revision ID: 0018_local_owner_repair
|
||||
Revises: 0016_agent_workspace, 0017_oss_workspace_identity
|
||||
Create Date: 2026-07-31
|
||||
"""
|
||||
|
||||
@@ -14,8 +14,8 @@ import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
|
||||
revision = '0017_local_owner_repair'
|
||||
down_revision = '0016_agent_workspace'
|
||||
revision = '0018_local_owner_repair'
|
||||
down_revision = ('0016_agent_workspace', '0017_oss_workspace_identity')
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
@@ -53,6 +53,7 @@ _ALEMBIC_TENANT_TABLES = {
|
||||
'workspace_memberships',
|
||||
'workspace_invitations',
|
||||
'workspace_execution_states',
|
||||
'support_admin_temporary_sessions',
|
||||
'workspace_metadata',
|
||||
'api_keys',
|
||||
'agents',
|
||||
|
||||
@@ -43,6 +43,7 @@ TENANT_TABLE_COLUMNS: dict[str, str] = {
|
||||
'workspace_memberships': 'workspace_uuid',
|
||||
'workspace_invitations': 'workspace_uuid',
|
||||
'workspace_execution_states': 'workspace_uuid',
|
||||
'support_admin_temporary_sessions': 'workspace_uuid',
|
||||
'workspace_metadata': 'workspace_uuid',
|
||||
'api_keys': 'workspace_uuid',
|
||||
'agents': 'workspace_uuid',
|
||||
|
||||
@@ -16,6 +16,7 @@ from ....agent.runner.config_resolver import RunnerConfigResolver
|
||||
from ....agent.runner import config_schema
|
||||
from ....utils import constants, runner as runner_utils
|
||||
from ....telemetry import features as telemetry_features
|
||||
from ....telemetry.identity import workspace_identity
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
@@ -303,7 +304,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,
|
||||
|
||||
@@ -179,10 +179,13 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
|
||||
)
|
||||
file_format = 'image/jpeg'
|
||||
|
||||
# NOTE: Telegram's file.file_path is a full URL of the form
|
||||
# https://api.telegram.org/file/bot<TOKEN>/<path> which embeds the
|
||||
# bot token. Unlike the public CDN URLs used by other adapters, it
|
||||
# cannot be exposed safely, so only base64 is stored here.
|
||||
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
|
||||
message_components.append(
|
||||
platform_message.Image(
|
||||
url=file.file_path,
|
||||
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
|
||||
)
|
||||
)
|
||||
|
||||
@@ -9,7 +9,7 @@ from datetime import datetime
|
||||
|
||||
import pydantic
|
||||
|
||||
from ...api.http.context import ExecutionContext
|
||||
from ...api.http.context import ExecutionContext, PrincipalContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_SESSION_FILTER_UNSET = object()
|
||||
@@ -95,6 +95,9 @@ class WebSocketConnection(pydantic.BaseModel):
|
||||
metadata: dict = pydantic.Field(default_factory=dict)
|
||||
"""连接元数据(可存储额外信息)"""
|
||||
|
||||
trigger_principal: PrincipalContext | None = None
|
||||
"""Authenticated principal that opened this dashboard connection."""
|
||||
|
||||
@property
|
||||
def scope(self) -> WebSocketScope:
|
||||
return WebSocketScope(
|
||||
@@ -112,6 +115,7 @@ class WebSocketConnection(pydantic.BaseModel):
|
||||
workspace_uuid=self.workspace_uuid,
|
||||
placement_generation=self.placement_generation,
|
||||
pipeline_uuid=self.pipeline_uuid,
|
||||
trigger_principal=self.trigger_principal,
|
||||
)
|
||||
|
||||
|
||||
@@ -138,6 +142,7 @@ class WebSocketConnectionManager:
|
||||
pipeline_uuid: str,
|
||||
session_type: str,
|
||||
metadata: dict | None = None,
|
||||
trigger_principal: PrincipalContext | None = None,
|
||||
session_id: str | None = None,
|
||||
send_queue_size: int = _DEFAULT_SEND_QUEUE_SIZE,
|
||||
max_connections: int = 1024,
|
||||
@@ -174,6 +179,7 @@ class WebSocketConnectionManager:
|
||||
session_id=session_id,
|
||||
websocket=websocket,
|
||||
metadata=metadata or {},
|
||||
trigger_principal=trigger_principal,
|
||||
send_queue=asyncio.Queue(maxsize=send_queue_size),
|
||||
)
|
||||
|
||||
|
||||
@@ -27,6 +27,19 @@ 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]
|
||||
execution_generation: int
|
||||
|
||||
|
||||
async def _count(
|
||||
ap: core_app.Application,
|
||||
table,
|
||||
@@ -52,8 +65,64 @@ async def _count(
|
||||
return -1
|
||||
|
||||
|
||||
async def build_heartbeat_payload(ap: core_app.Application) -> dict:
|
||||
"""Collect the anonymous instance profile snapshot."""
|
||||
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 []
|
||||
|
||||
counts: dict[str, WorkspaceResourceSnapshot] = {
|
||||
binding.workspace_uuid: {
|
||||
'workspace_uuid': binding.workspace_uuid,
|
||||
'bot_count': 0,
|
||||
'pipeline_count': 0,
|
||||
'knowledge_base_count': 0,
|
||||
'plugin_count': 0,
|
||||
'mcp_server_count': 0,
|
||||
'extension_count': 0,
|
||||
'skill_count': 0,
|
||||
'adapters': [],
|
||||
'execution_generation': binding.placement_generation,
|
||||
}
|
||||
for binding in bindings
|
||||
}
|
||||
|
||||
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
|
||||
for key in getattr(ap.rag_mgr, 'knowledge_bases', {}):
|
||||
if len(key) >= 1 and key[0] in counts:
|
||||
counts[key[0]]['knowledge_base_count'] += 1
|
||||
for key in getattr(ap.tool_mgr.mcp_tool_loader, '_sessions', {}):
|
||||
if len(key) >= 2 and key[1] in counts:
|
||||
counts[key[1]]['mcp_server_count'] += 1
|
||||
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 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,
|
||||
*,
|
||||
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
|
||||
@@ -136,11 +205,14 @@ async def build_heartbeat_payload(ap: core_app.Application) -> dict:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
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,
|
||||
@@ -148,14 +220,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