feat(tenancy): connect cloud workspace control plane

This commit is contained in:
Junyan Qin
2026-07-24 19:11:33 +08:00
parent d7cdd206c2
commit 98f45aa88e
40 changed files with 4159 additions and 452 deletions
+60 -2
View File
@@ -109,12 +109,15 @@ class UserService:
return await asyncio.to_thread(argon2.PasswordHasher().hash, password)
def _require_local_directory(self) -> None:
workspace_service = getattr(self.ap, 'workspace_service', None)
if workspace_service is not None and workspace_service.policy.multi_workspace_enabled:
if self._uses_control_plane_directory():
raise ControlPlaneDirectoryRequiredError(
'Cloud Accounts and directory changes are managed by the SaaS control plane'
)
def _uses_control_plane_directory(self) -> bool:
workspace_service = getattr(self.ap, 'workspace_service', None)
return bool(workspace_service is not None and workspace_service.policy.multi_workspace_enabled)
async def _verify_password(self, hashed_password: str, password: str) -> None:
async with self._password_hash_lock:
await asyncio.to_thread(argon2.PasswordHasher().verify, hashed_password, password)
@@ -427,6 +430,15 @@ class UserService:
expires_in: int = 0,
) -> user.User:
"""Create or update a Space user account (only if system not initialized or user exists)"""
if self._uses_control_plane_directory():
return await self._update_projected_space_user(
space_account_uuid=space_account_uuid,
email=email,
access_token=access_token,
refresh_token=refresh_token,
api_key=api_key,
expires_in=expires_in,
)
self._require_local_directory()
expires_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in) if expires_in > 0 else None
@@ -507,6 +519,52 @@ class UserService:
await self._update_space_provider_for_account(created_user, api_key)
return created_user
async def _update_projected_space_user(
self,
*,
space_account_uuid: str,
email: str,
access_token: str,
refresh_token: str,
api_key: str,
expires_in: int,
) -> user.User:
"""Attach OAuth credentials to an already projected Cloud Account."""
normalized_email = normalize_email(email)
expires_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in) if expires_in > 0 else None
async with self._create_user_lock:
projected = await self.get_user_by_space_account_uuid(space_account_uuid)
if (
projected is None
or projected.uuid != space_account_uuid
or projected.normalized_email != normalized_email
or projected.source != user.AccountSource.CLOUD_PROJECTION.value
or projected.account_type != 'space'
):
raise ControlPlaneDirectoryRequiredError('Space Account is not present in the verified Cloud directory')
self._require_active_account(projected)
await self._identity_execute(
sqlalchemy.update(user.User)
.where(
user.User.uuid == projected.uuid,
user.User.space_account_uuid == space_account_uuid,
user.User.source == user.AccountSource.CLOUD_PROJECTION.value,
)
.values(
space_access_token=access_token,
space_refresh_token=refresh_token,
space_api_key=api_key,
space_access_token_expires_at=expires_at,
),
f'space:{space_account_uuid}',
)
refreshed = await self.get_user_by_space_account_uuid(space_account_uuid)
if refreshed is None:
raise ControlPlaneDirectoryRequiredError('Space Account disappeared from the verified Cloud directory')
self._require_active_account(refreshed)
return refreshed
async def authenticate_space_user(
self, access_token: str, refresh_token: str, expires_in: int = 0
) -> typing.Tuple[str, user.User]:
+1
View File
@@ -351,6 +351,7 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
make_connection_failed_callback=on_connect_failed,
additional_headers=self.get_control_headers(),
)
self._ctrl = ctrl
self._ctrl_task = asyncio.create_task(
ctrl.run(self._make_connection_callback(transport_name, connected, connect_error, self._generation))
)
+50 -3
View File
@@ -145,8 +145,38 @@ class BoxService:
self._available = False
self._connector_error = str(exc)
if self._cloud_managed:
await self._abort_failed_cloud_initialization()
raise
async def _abort_failed_cloud_initialization(self) -> None:
"""Close a connected Cloud transport before propagating readiness failure.
Connector initialization starts control and heartbeat tasks before Core
performs the stricter Cloud readiness challenge. If that challenge
fails, startup must remain fail-closed without leaving those tasks free
to schedule reconnect work while the application loop is unwinding.
"""
self._closing = True
self._available = False
reconnect_task = self._reconnect_task
self._reconnect_task = None
self._reconnecting = False
if reconnect_task is not None and reconnect_task is not asyncio.current_task():
reconnect_task.cancel()
await asyncio.gather(reconnect_task, return_exceptions=True)
connector = self._runtime_connector
if connector is None:
return
connector.runtime_disconnect_callback = None
try:
await connector.aclose()
except Exception:
# Cleanup failure must not replace the readiness error which caused
# Cloud startup to fail closed.
self.ap.logger.exception('Failed to close Box runtime after Cloud readiness validation failed')
async def _on_runtime_disconnect(self, connector: BoxRuntimeConnector) -> None:
"""Called by the connector when the Box runtime connection drops.
@@ -156,13 +186,28 @@ class BoxService:
"""
if not self._enabled or self._closing:
return
try:
loop = asyncio.get_running_loop()
except RuntimeError:
return
if loop.is_closed():
return
if self._reconnect_task is not None and not self._reconnect_task.done():
return # Another reconnect loop is already running
self._reconnecting = True
self._available = False
self._connector_error = 'Disconnected from Box runtime'
self.ap.logger.warning('Box runtime disconnected, sandbox features temporarily disabled.')
self._reconnect_task = asyncio.create_task(self._reconnect_loop(connector))
reconnect = self._reconnect_loop(connector)
try:
self._reconnect_task = loop.create_task(reconnect)
except RuntimeError:
# The loop may begin closing between get_running_loop() and task
# creation. Explicitly close the coroutine so shutdown emits no
# "coroutine was never awaited" warning.
reconnect.close()
self._reconnecting = False
self._reconnect_task = None
async def _reconnect_loop(self, connector: BoxRuntimeConnector) -> None:
"""Retry reconnection with exponential backoff (3s → 60s max)."""
@@ -173,9 +218,11 @@ class BoxService:
self.ap.logger.info(f'Attempting to reconnect to Box runtime in {delay}s...')
await asyncio.sleep(delay)
try:
connector.dispose()
await connector.initialize()
await connector.reconnect()
self._ensure_default_workspace()
await self._verify_cloud_runtime()
if not self._cloud_managed:
await self._purge_attachment_dirs()
self._available = True
self._connector_error = ''
skill_mgr = getattr(self.ap, 'skill_mgr', None)
+24
View File
@@ -2,10 +2,23 @@
from .bootstrap import (
CloudBootstrapError,
CloudManifestProvider,
CloudManifestRefreshService,
OpenSourceDeployment,
VerifiedCloudDeployment,
resolve_deployment,
)
from .directory import (
DirectoryDelta,
DirectoryEvent,
DirectoryEventBatch,
DirectoryMember,
DirectoryProjectionProvider,
DirectoryProjectionUnavailableError,
DirectorySnapshot,
DirectoryWorkspace,
)
from .directory_projection import DirectoryProjectionService
from .entitlements import (
EntitlementProvider,
EntitlementResolver,
@@ -16,6 +29,17 @@ from .entitlements import (
__all__ = [
'CloudBootstrapError',
'CloudManifestProvider',
'CloudManifestRefreshService',
'DirectoryDelta',
'DirectoryEvent',
'DirectoryEventBatch',
'DirectoryMember',
'DirectoryProjectionProvider',
'DirectoryProjectionService',
'DirectoryProjectionUnavailableError',
'DirectorySnapshot',
'DirectoryWorkspace',
'EntitlementProvider',
'EntitlementResolver',
'EntitlementSnapshot',
+83 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import dataclasses
import importlib.metadata
import inspect
@@ -7,9 +8,10 @@ import os
import threading
import time
from collections.abc import Awaitable, Callable
from typing import Any, Protocol
from typing import Any, Protocol, runtime_checkable
from ..workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
from .directory import DirectoryProjectionProvider
from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider
@@ -26,6 +28,17 @@ class CloudRuntimeUnavailableError(CloudBootstrapError):
"""The verified Cloud receipt no longer admits runtime work."""
@runtime_checkable
class CloudManifestProvider(Protocol):
"""Closed adapter responsible for renewing the signed deployment receipt."""
async def refresh_manifest(self) -> VerifiedCloudDeployment:
"""Fetch, verify, and return the newest deployment receipt."""
async def aclose(self) -> None:
"""Release control-plane transport resources."""
@dataclasses.dataclass(frozen=True, slots=True)
class OpenSourceDeployment:
"""Default deployment selected when no closed bootstrap is installed."""
@@ -35,6 +48,8 @@ class OpenSourceDeployment:
entitlement_provider: OpenSourceEntitlementProvider = dataclasses.field(
default_factory=OpenSourceEntitlementProvider
)
directory_provider: None = None
manifest_provider: None = None
persistence_mode: str = 'oss_compat'
required_vector_backend: str | None = None
@@ -63,6 +78,8 @@ class VerifiedCloudDeployment:
capabilities: frozenset[str]
tenant_isolation_version: int
entitlement_provider: EntitlementProvider
directory_provider: DirectoryProjectionProvider
manifest_provider: CloudManifestProvider
verification_key_id: str
mode: str = dataclasses.field(default='cloud', init=False)
workspace_policy: CloudWorkspacePolicy = dataclasses.field(default_factory=CloudWorkspacePolicy, init=False)
@@ -89,6 +106,10 @@ class VerifiedCloudDeployment:
raise CloudBootstrapError('Verified Cloud Manifest does not grant multi_workspace_v2')
if not isinstance(self.entitlement_provider, EntitlementProvider):
raise CloudBootstrapError('Verified Cloud bootstrap did not provide an entitlement adapter')
if not isinstance(self.directory_provider, DirectoryProjectionProvider):
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a directory adapter')
if not isinstance(self.manifest_provider, CloudManifestProvider):
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a Manifest renewal adapter')
def validate_instance_config(self, config: dict[str, Any]) -> None:
if config.get('database', {}).get('use') != 'postgresql':
@@ -262,6 +283,67 @@ class DeploymentAdmissionGuard:
return deployment
class CloudManifestRefreshService:
"""Renew a short-lived verified Manifest before runtime admission expires."""
def __init__(
self,
admission: DeploymentAdmissionGuard,
provider: CloudManifestProvider,
logger: Any,
*,
wall_time: Callable[[], float] = time.time,
refresh_margin_seconds: int = 180,
maximum_sleep_seconds: int = 300,
) -> None:
if not isinstance(provider, CloudManifestProvider):
raise TypeError('Cloud Manifest refresh requires a CloudManifestProvider')
if refresh_margin_seconds < 120:
raise ValueError('Cloud Manifest refresh margin must be at least 120 seconds')
if maximum_sleep_seconds <= 0:
raise ValueError('Cloud Manifest refresh maximum sleep must be positive')
self.admission = admission
self.provider = provider
self.logger = logger
self._wall_time = wall_time
self.refresh_margin_seconds = refresh_margin_seconds
self.maximum_sleep_seconds = maximum_sleep_seconds
def next_refresh_delay(self) -> float:
deployment = self.admission.deployment
if not isinstance(deployment, VerifiedCloudDeployment):
return float(self.maximum_sleep_seconds)
remaining = deployment.expires_at - self._wall_time()
return max(
5.0,
min(
float(self.maximum_sleep_seconds),
remaining - self.refresh_margin_seconds,
),
)
async def refresh_once(self) -> VerifiedCloudDeployment:
candidate = await self.provider.refresh_manifest()
if not isinstance(candidate, VerifiedCloudDeployment):
raise CloudBootstrapError('Cloud Manifest provider returned an invalid deployment receipt')
self.admission.replace(candidate)
return candidate
async def run(self) -> None:
retry_delay = 5.0
while True:
try:
await asyncio.sleep(self.next_refresh_delay())
await self.refresh_once()
retry_delay = 5.0
except asyncio.CancelledError:
raise
except Exception:
self.logger.exception('Cloud Manifest refresh failed')
await asyncio.sleep(retry_delay)
retry_delay = min(retry_delay * 2, 30.0)
async def _invoke_provider(
loaded: Any,
*,
+215
View File
@@ -0,0 +1,215 @@
from __future__ import annotations
import datetime
from collections.abc import Sequence
from typing import Any, Protocol, runtime_checkable
import pydantic
class DirectoryProjectionUnavailableError(RuntimeError):
"""Raised when the verified Cloud directory cannot safely admit work."""
class DirectoryMember(pydantic.BaseModel):
"""One account membership published by the SaaS control plane."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
membership_uuid: str = pydantic.Field(min_length=1, max_length=36)
account_uuid: str = pydantic.Field(min_length=1, max_length=36)
normalized_email: str = pydantic.Field(min_length=1, max_length=320)
display_name: str = pydantic.Field(min_length=1, max_length=255)
account_status: str = pydantic.Field(pattern=r'^(active|blocked|disabled|deleted)$')
role: str = pydantic.Field(pattern=r'^(owner|admin|member|developer|operator|viewer)$')
membership_status: str = pydantic.Field(pattern=r'^(active|invited|disabled|removed)$')
projection_revision: int = pydantic.Field(ge=1)
joined_at: datetime.datetime | None = None
@pydantic.field_validator('normalized_email')
@classmethod
def _normalize_email(cls, value: str) -> str:
normalized = value.strip().casefold()
if normalized != value:
raise ValueError('Directory email must already be normalized')
return normalized
class DirectoryWorkspace(pydantic.BaseModel):
"""One Workspace and its authoritative membership projection."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
uuid: str = pydantic.Field(min_length=1, max_length=36)
name: str = pydantic.Field(min_length=1, max_length=255)
slug: str = pydantic.Field(min_length=1, max_length=255)
type: str = pydantic.Field(pattern=r'^(personal|team)$')
status: str = pydantic.Field(pattern=r'^(provisioning|active|suspended|archived|deleted)$')
created_by_account_uuid: str = pydantic.Field(min_length=1, max_length=36)
projection_revision: int = pydantic.Field(ge=1)
execution_generation: int = pydantic.Field(ge=1)
members: tuple[DirectoryMember, ...] = ()
@pydantic.field_validator('members', mode='before')
@classmethod
def _copy_members(cls, value: Sequence[DirectoryMember] | None) -> tuple[DirectoryMember, ...]:
return tuple(value or ())
@pydantic.model_validator(mode='after')
def _validate_members(self) -> DirectoryWorkspace:
membership_uuids = [member.membership_uuid for member in self.members]
account_uuids = [member.account_uuid for member in self.members]
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory Workspace contains duplicate membership UUIDs')
if len(account_uuids) != len(set(account_uuids)):
raise ValueError('Directory Workspace contains duplicate account UUIDs')
if self.created_by_account_uuid not in set(account_uuids):
raise ValueError('Directory Workspace must include its creator')
return self
class DirectorySnapshot(pydantic.BaseModel):
"""Full signed directory state at one monotonic outbox cursor."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
instance_uuid: str = pydantic.Field(min_length=1, max_length=255)
cursor: int = pydantic.Field(ge=0)
generated_at: datetime.datetime
workspaces: tuple[DirectoryWorkspace, ...] = ()
@pydantic.field_validator('workspaces', mode='before')
@classmethod
def _copy_workspaces(cls, value: Sequence[DirectoryWorkspace] | None) -> tuple[DirectoryWorkspace, ...]:
return tuple(value or ())
@pydantic.model_validator(mode='after')
def _validate_workspaces(self) -> DirectorySnapshot:
workspace_uuids = [workspace.uuid for workspace in self.workspaces]
slugs = [workspace.slug for workspace in self.workspaces]
membership_uuids = [member.membership_uuid for workspace in self.workspaces for member in workspace.members]
if len(workspace_uuids) != len(set(workspace_uuids)):
raise ValueError('Directory snapshot contains duplicate Workspace UUIDs')
if len(slugs) != len(set(slugs)):
raise ValueError('Directory snapshot contains duplicate Workspace slugs')
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory snapshot contains duplicate membership UUIDs')
return self
class DirectoryDelta(pydantic.BaseModel):
"""Signed authoritative state for an explicitly requested Workspace set."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
instance_uuid: str = pydantic.Field(min_length=1, max_length=255)
requested_workspace_uuids: tuple[str, ...]
generated_at: datetime.datetime
workspaces: tuple[DirectoryWorkspace, ...] = ()
@pydantic.field_validator('requested_workspace_uuids', mode='before')
@classmethod
def _copy_requested_workspace_uuids(cls, value: Sequence[str]) -> tuple[str, ...]:
return tuple(value)
@pydantic.field_validator('workspaces', mode='before')
@classmethod
def _copy_workspaces(cls, value: Sequence[DirectoryWorkspace] | None) -> tuple[DirectoryWorkspace, ...]:
return tuple(value or ())
@pydantic.model_validator(mode='after')
def _validate_workspaces(self) -> DirectoryDelta:
requested = self.requested_workspace_uuids
if not requested or len(requested) > 100:
raise ValueError('Directory delta must request between 1 and 100 Workspaces')
if any(not workspace_uuid or len(workspace_uuid) > 36 for workspace_uuid in requested):
raise ValueError('Directory delta contains an invalid requested Workspace UUID')
if len(requested) != len(set(requested)):
raise ValueError('Directory delta contains duplicate requested Workspace UUIDs')
workspace_uuids = [workspace.uuid for workspace in self.workspaces]
slugs = [workspace.slug for workspace in self.workspaces]
membership_uuids = [member.membership_uuid for workspace in self.workspaces for member in workspace.members]
if len(workspace_uuids) != len(set(workspace_uuids)):
raise ValueError('Directory delta contains duplicate Workspace UUIDs')
if not set(workspace_uuids).issubset(set(requested)):
raise ValueError('Directory delta returned an unrequested Workspace')
if len(slugs) != len(set(slugs)):
raise ValueError('Directory delta contains duplicate Workspace slugs')
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory delta contains duplicate membership UUIDs')
return self
class DirectoryEvent(pydantic.BaseModel):
"""One signed control-plane outbox notification."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
cursor: int = pydantic.Field(ge=1)
uuid: str = pydantic.Field(min_length=1, max_length=36)
aggregate_uuid: str = pydantic.Field(min_length=1, max_length=36)
event_type: str = pydantic.Field(min_length=1, max_length=128)
revision: int = pydantic.Field(ge=1)
payload: dict[str, Any] = pydantic.Field(default_factory=dict)
created_at: datetime.datetime
class DirectoryEventBatch(pydantic.BaseModel):
"""Signed events returned after a caller-supplied directory cursor."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
instance_uuid: str = pydantic.Field(min_length=1, max_length=255)
after_cursor: int = pydantic.Field(ge=0)
cursor: int = pydantic.Field(ge=0)
high_water_cursor: int = pydantic.Field(ge=0)
events: tuple[DirectoryEvent, ...] = ()
@pydantic.field_validator('events', mode='before')
@classmethod
def _copy_events(cls, value: Sequence[DirectoryEvent] | None) -> tuple[DirectoryEvent, ...]:
return tuple(value or ())
@pydantic.model_validator(mode='after')
def _validate_events(self) -> DirectoryEventBatch:
if self.cursor < self.after_cursor:
raise ValueError('Directory event cursor rolled back')
if self.high_water_cursor < self.cursor:
raise ValueError('Directory event high-water mark rolled back')
event_cursors = [event.cursor for event in self.events]
event_uuids = [event.uuid for event in self.events]
if event_cursors != sorted(event_cursors) or len(event_cursors) != len(set(event_cursors)):
raise ValueError('Directory events must have strictly increasing cursors')
if len(event_uuids) != len(set(event_uuids)):
raise ValueError('Directory event batch contains duplicate UUIDs')
if any(cursor <= self.after_cursor or cursor > self.cursor for cursor in event_cursors):
raise ValueError('Directory event falls outside the requested cursor window')
if not self.events and (self.cursor != self.after_cursor or self.high_water_cursor != self.after_cursor):
raise ValueError('Empty Directory event batch cannot advance or trail the high-water mark')
if self.events and self.cursor != self.events[-1].cursor:
raise ValueError('Directory event batch cursor must equal its final event cursor')
return self
@runtime_checkable
class DirectoryProjectionProvider(Protocol):
"""Closed adapter that returns signature-verified control-plane data."""
async def fetch_snapshot(self, instance_uuid: str) -> DirectorySnapshot:
"""Fetch and verify an authoritative full snapshot."""
async def fetch_events(
self,
instance_uuid: str,
after_cursor: int,
limit: int,
) -> DirectoryEventBatch:
"""Fetch and verify directory events after one process-local cursor."""
async def fetch_workspaces(
self,
instance_uuid: str,
workspace_uuids: tuple[str, ...],
) -> DirectoryDelta:
"""Fetch and verify authoritative state for an explicit Workspace set."""
@@ -0,0 +1,881 @@
from __future__ import annotations
import asyncio
import datetime
import hashlib
import json
import time
from collections.abc import Callable, Iterable
from typing import TYPE_CHECKING, Any
import sqlalchemy
from sqlalchemy.dialects import postgresql, sqlite
from ..entity.persistence.cloud_directory import DirectoryProjectionInbox, DirectoryProjectionState
from ..entity.persistence.user import AccountSource, AccountStatus, User
from ..entity.persistence.workspace import (
MembershipRole,
MembershipStatus,
Workspace,
WorkspaceExecutionSource,
WorkspaceExecutionState,
WorkspaceExecutionStatus,
WorkspaceMembership,
WorkspaceSource,
WorkspaceStatus,
)
from .directory import (
DirectoryDelta,
DirectoryEvent,
DirectoryEventBatch,
DirectoryMember,
DirectoryProjectionProvider,
DirectoryProjectionUnavailableError,
DirectorySnapshot,
DirectoryWorkspace,
)
if TYPE_CHECKING:
from ..core.app import Application
_ROLE_MAP = {
'owner': MembershipRole.OWNER.value,
'admin': MembershipRole.ADMIN.value,
# Space deliberately exposes a smaller product role vocabulary. A regular
# SaaS member receives the Core developer role; operator/viewer can be
# introduced later without changing the signed directory contract.
'member': MembershipRole.DEVELOPER.value,
'developer': MembershipRole.DEVELOPER.value,
'operator': MembershipRole.OPERATOR.value,
'viewer': MembershipRole.VIEWER.value,
}
_ACCOUNT_STATUS_MAP = {
'active': AccountStatus.ACTIVE.value,
'blocked': AccountStatus.DISABLED.value,
'disabled': AccountStatus.DISABLED.value,
'deleted': AccountStatus.DELETED.value,
}
_MEMBERSHIP_STATUS_MAP = {
'active': MembershipStatus.ACTIVE.value,
'invited': MembershipStatus.DISABLED.value,
'disabled': MembershipStatus.DISABLED.value,
'removed': MembershipStatus.REMOVED.value,
}
_INCREMENTAL_PROJECTION_FINGERPRINT = hashlib.sha256(b'langbot-directory-incremental-v1').hexdigest()
class _DirectorySnapshotSuperseded(DirectoryProjectionUnavailableError):
"""A valid snapshot lost a race with a newer shared projection."""
class DirectoryProjectionService:
"""Project a verified SaaS directory into Core-owned tenant tables.
The closed adapter verifies transport signatures and returns immutable
models. Core owns database transactions, revision checks, execution fences,
and readiness. This keeps the ORM and PostgreSQL RLS boundary out of the
closed control-plane package.
"""
def __init__(
self,
ap: Application,
provider: DirectoryProjectionProvider,
instance_uuid: str,
*,
sync_interval_seconds: float = 5.0,
max_staleness_seconds: float = 60.0,
event_limit: int = 100,
monotonic_time: Callable[[], float] = time.monotonic,
) -> None:
if not isinstance(provider, DirectoryProjectionProvider):
raise TypeError('Cloud directory projection requires a DirectoryProjectionProvider')
if not instance_uuid.strip():
raise ValueError('Cloud directory projection requires an instance UUID')
if sync_interval_seconds <= 0:
raise ValueError('Directory sync interval must be positive')
if max_staleness_seconds <= sync_interval_seconds:
raise ValueError('Directory max staleness must exceed the sync interval')
if event_limit <= 0 or event_limit > 100:
raise ValueError('Directory event limit must be between 1 and 100')
self.ap = ap
self.provider = provider
self.instance_uuid = instance_uuid.strip()
self.sync_interval_seconds = float(sync_interval_seconds)
self.max_staleness_seconds = float(max_staleness_seconds)
self.event_limit = event_limit
self._monotonic_time = monotonic_time
self._last_success_monotonic: float | None = None
self._ready = False
# Every runtime replica must consume the event stream independently:
# entitlement snapshots live in the closed adapter's process memory.
# The database cursor remains the shared projection high-water mark,
# while this cursor tracks what this process has actually observed.
self._consumer_cursor: int | None = None
async def initialize(self) -> None:
"""Block Cloud startup until one full signed snapshot is committed."""
last_superseded: _DirectorySnapshotSuperseded | None = None
for _attempt in range(5):
snapshot = await self.provider.fetch_snapshot(self.instance_uuid)
try:
await self.apply_snapshot(snapshot)
except _DirectorySnapshotSuperseded as exc:
last_superseded = exc
continue
self._consumer_cursor = snapshot.cursor
return
raise DirectoryProjectionUnavailableError(
'Directory snapshot was repeatedly superseded by another runtime replica'
) from last_superseded
async def run(self) -> None:
"""Continuously refresh the directory and fail closed when it goes stale."""
delay = self.sync_interval_seconds
while True:
try:
await asyncio.sleep(delay)
await self.sync_once()
delay = self.sync_interval_seconds
except asyncio.CancelledError:
raise
except Exception:
self.ap.logger.exception('Cloud directory synchronization failed')
delay = min(max(delay * 2, self.sync_interval_seconds), self.max_staleness_seconds / 2)
async def sync_once(self) -> None:
cursor = self._consumer_cursor
if cursor is None:
await self.initialize()
return
batch = await self.provider.fetch_events(
self.instance_uuid,
cursor,
self.event_limit,
)
batch = DirectoryEventBatch.model_validate(batch.model_dump())
self._validate_batch(batch, expected_after_cursor=cursor)
if batch.events:
directory_revisions = self._directory_event_revisions(batch.events)
if directory_revisions:
requested_workspace_uuids = tuple(sorted(directory_revisions))
delta = await self.provider.fetch_workspaces(
self.instance_uuid,
requested_workspace_uuids,
)
await self.apply_delta(delta, batch)
else:
await self.apply_event_batch(batch)
self._consumer_cursor = batch.cursor
return
await self._touch_freshness(cursor)
def require_ready(self) -> None:
"""Fail synchronously at execution admission when projection is stale."""
last_success = self._last_success_monotonic
if not self._ready or last_success is None:
raise DirectoryProjectionUnavailableError('Cloud directory projection is not ready')
if self._monotonic_time() - last_success >= self.max_staleness_seconds:
raise DirectoryProjectionUnavailableError('Cloud directory projection is stale')
async def apply_snapshot(
self,
snapshot: DirectorySnapshot,
*,
events: Iterable[DirectoryEvent] = (),
) -> None:
"""Atomically apply one monotonic full snapshot and its event receipts."""
if not isinstance(snapshot, DirectorySnapshot):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid snapshot')
snapshot = DirectorySnapshot.model_validate(snapshot.model_dump())
if snapshot.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory snapshot targets another LangBot instance')
fingerprint = self._snapshot_fingerprint(snapshot)
now = self._utcnow()
lease_expires_at = now + datetime.timedelta(seconds=self.max_staleness_seconds)
event_list = tuple(DirectoryEvent.model_validate(event.model_dump()) for event in events)
directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None)
if not callable(directory_uow):
raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable')
async with directory_uow(self.instance_uuid) as uow:
session = uow.session
state_values = {
'instance_uuid': self.instance_uuid,
'cursor': snapshot.cursor,
'snapshot_coverage_cursor': snapshot.cursor,
'snapshot_fingerprint': fingerprint,
'last_applied_at': now,
'lease_expires_at': lease_expires_at,
}
dialect_name = self.ap.persistence_mgr.get_db_engine().dialect.name
if dialect_name == 'postgresql':
insert_state = postgresql.insert(DirectoryProjectionState)
elif dialect_name == 'sqlite':
insert_state = sqlite.insert(DirectoryProjectionState)
else: # pragma: no cover - Cloud supports PostgreSQL; tests use SQLite.
raise DirectoryProjectionUnavailableError('Directory projection database is unsupported')
await session.execute(
insert_state.values(**state_values).on_conflict_do_nothing(
index_elements=[DirectoryProjectionState.instance_uuid]
)
)
state = await session.scalar(
sqlalchemy.select(DirectoryProjectionState)
.where(DirectoryProjectionState.instance_uuid == self.instance_uuid)
.with_for_update()
)
if state is None: # pragma: no cover - insert/select are one transaction.
raise DirectoryProjectionUnavailableError('Directory projection state could not be locked')
if snapshot.cursor < state.cursor:
raise _DirectorySnapshotSuperseded('Directory snapshot cursor rolled back')
if snapshot.cursor == state.cursor and state.snapshot_fingerprint not in {
fingerprint,
_INCREMENTAL_PROJECTION_FINGERPRINT,
}:
raise DirectoryProjectionUnavailableError('Directory snapshot cursor has conflicting contents')
await self._record_events(session, event_list, now=now)
await self._apply_accounts(session, snapshot)
await self._apply_workspaces(session, snapshot)
await self._fence_absent_workspaces(session, snapshot)
state.cursor = snapshot.cursor
state.snapshot_coverage_cursor = snapshot.cursor
state.snapshot_fingerprint = fingerprint
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await self._mark_events_applied(session, event_list, now=now)
await session.flush()
self._record_success()
self._consumer_cursor = snapshot.cursor
async def apply_delta(self, delta: DirectoryDelta, batch: DirectoryEventBatch) -> None:
"""Apply only Workspaces named by directory events in one signed page."""
if not isinstance(delta, DirectoryDelta):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid delta')
if not isinstance(batch, DirectoryEventBatch):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid event batch')
delta = DirectoryDelta.model_validate(delta.model_dump())
batch = DirectoryEventBatch.model_validate(batch.model_dump())
self._validate_batch(batch, expected_after_cursor=batch.after_cursor)
if delta.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory delta targets another LangBot instance')
required_revisions = self._directory_event_revisions(batch.events)
requested = set(delta.requested_workspace_uuids)
if not required_revisions or requested != set(required_revisions):
raise DirectoryProjectionUnavailableError('Directory delta does not match its event batch')
returned = {workspace.uuid: workspace for workspace in delta.workspaces}
for workspace_uuid, workspace in returned.items():
if workspace.projection_revision < required_revisions[workspace_uuid]:
raise DirectoryProjectionUnavailableError(
'Directory Workspace delta is older than its signed event notification'
)
now = self._utcnow()
lease_expires_at = now + datetime.timedelta(seconds=self.max_staleness_seconds)
directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None)
if not callable(directory_uow):
raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable')
projection_caught_up = False
async with directory_uow(self.instance_uuid) as uow:
session = uow.session
state = await session.scalar(
sqlalchemy.select(DirectoryProjectionState)
.where(DirectoryProjectionState.instance_uuid == self.instance_uuid)
.with_for_update()
)
if state is None:
raise DirectoryProjectionUnavailableError('Directory projection state disappeared')
state_cursor = int(state.cursor)
if state_cursor < batch.after_cursor:
raise DirectoryProjectionUnavailableError('Directory projection state cursor rolled back')
# A different runtime replica may already have applied this page.
# In that case every receipt through the shared cursor must exist;
# this replica still fetched the delta and refreshed its own
# entitlement cache before advancing its process-local cursor.
await self._record_events(
session,
batch.events,
now=now,
allow_missing_through_cursor=int(state.snapshot_coverage_cursor),
reject_missing_through_cursor=state_cursor,
)
if state_cursor < batch.cursor:
projected_delta = DirectorySnapshot(
instance_uuid=self.instance_uuid,
cursor=batch.cursor,
generated_at=delta.generated_at,
workspaces=delta.workspaces,
)
await self._apply_accounts(session, projected_delta)
await self._apply_workspaces(session, projected_delta)
await self._fence_workspaces(
session,
{
workspace_uuid: required_revisions[workspace_uuid]
for workspace_uuid in requested - set(returned)
},
)
state.cursor = batch.cursor
# A per-Workspace delta cannot prove a full-directory
# fingerprint. Event receipts and entity revisions protect the
# incremental path; a later full snapshot replaces this marker.
state.snapshot_fingerprint = _INCREMENTAL_PROJECTION_FINGERPRINT
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await self._mark_events_applied(session, batch.events, now=now)
await session.flush()
projection_caught_up = batch.cursor == batch.high_water_cursor and int(state.cursor) == batch.cursor
if projection_caught_up:
self._record_success()
self._consumer_cursor = batch.cursor
async def apply_event_batch(self, batch: DirectoryEventBatch) -> None:
"""Advance non-directory events after the adapter refreshes local caches."""
if not isinstance(batch, DirectoryEventBatch):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid event batch')
batch = DirectoryEventBatch.model_validate(batch.model_dump())
self._validate_batch(batch, expected_after_cursor=batch.after_cursor)
if any(event.event_type == 'directory.changed' for event in batch.events):
raise DirectoryProjectionUnavailableError('Directory changes require an authoritative full snapshot')
now = self._utcnow()
lease_expires_at = now + datetime.timedelta(seconds=self.max_staleness_seconds)
directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None)
if not callable(directory_uow):
raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable')
projection_caught_up = False
async with directory_uow(self.instance_uuid) as uow:
state = await uow.session.scalar(
sqlalchemy.select(DirectoryProjectionState)
.where(DirectoryProjectionState.instance_uuid == self.instance_uuid)
.with_for_update()
)
if state is None:
raise DirectoryProjectionUnavailableError('Directory projection state disappeared')
if state.cursor < batch.after_cursor:
raise DirectoryProjectionUnavailableError('Directory projection state cursor rolled back')
await self._record_events(
uow.session,
batch.events,
now=now,
allow_missing_through_cursor=int(state.snapshot_coverage_cursor),
reject_missing_through_cursor=int(state.cursor),
)
state.cursor = max(int(state.cursor), batch.cursor)
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await self._mark_events_applied(uow.session, batch.events, now=now)
await uow.session.flush()
projection_caught_up = batch.cursor == batch.high_water_cursor and int(state.cursor) == batch.cursor
if projection_caught_up:
self._record_success()
self._consumer_cursor = batch.cursor
async def _touch_freshness(self, requested_cursor: int) -> None:
now = self._utcnow()
lease_expires_at = now + datetime.timedelta(seconds=self.max_staleness_seconds)
directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None)
if not callable(directory_uow):
raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable')
async with directory_uow(self.instance_uuid) as uow:
state = await uow.session.scalar(
sqlalchemy.select(DirectoryProjectionState)
.where(DirectoryProjectionState.instance_uuid == self.instance_uuid)
.with_for_update()
)
if state is None:
raise DirectoryProjectionUnavailableError('Directory projection state disappeared')
if state.cursor < requested_cursor:
raise DirectoryProjectionUnavailableError('Directory projection state cursor rolled back')
if state.cursor > requested_cursor:
raise DirectoryProjectionUnavailableError(
'This runtime replica has not consumed the shared directory high-water mark'
)
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await uow.session.flush()
self._record_success()
def _validate_batch(self, batch: DirectoryEventBatch, *, expected_after_cursor: int) -> None:
if batch.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory event batch targets another LangBot instance')
if batch.after_cursor != expected_after_cursor:
raise DirectoryProjectionUnavailableError('Directory event batch does not match the requested cursor')
supported_event_types = {'directory.changed', 'entitlement.changed'}
if any(event.event_type not in supported_event_types for event in batch.events):
raise DirectoryProjectionUnavailableError('Directory event batch contains an unsupported event type')
for event in batch.events:
if event.payload.get('workspace_uuid') != event.aggregate_uuid:
raise DirectoryProjectionUnavailableError('Directory event payload has a conflicting Workspace scope')
revision_key = 'directory_revision' if event.event_type == 'directory.changed' else 'entitlement_revision'
payload_revision = event.payload.get(revision_key)
if type(payload_revision) is not int or payload_revision != event.revision:
raise DirectoryProjectionUnavailableError('Directory event payload has a conflicting revision')
@staticmethod
def _directory_event_revisions(events: Iterable[DirectoryEvent]) -> dict[str, int]:
revisions: dict[str, int] = {}
for event in events:
if event.event_type == 'directory.changed':
revisions[event.aggregate_uuid] = max(revisions.get(event.aggregate_uuid, 0), event.revision)
return revisions
async def _record_events(
self,
session: Any,
events: tuple[DirectoryEvent, ...],
*,
now: datetime.datetime,
allow_missing_through_cursor: int = -1,
reject_missing_through_cursor: int | None = None,
) -> None:
for event in events:
fingerprint = self._fingerprint(event.model_dump(mode='json'))
existing = await session.scalar(
sqlalchemy.select(DirectoryProjectionInbox).where(
DirectoryProjectionInbox.instance_uuid == self.instance_uuid,
DirectoryProjectionInbox.event_uuid == event.uuid,
)
)
if existing is not None:
if existing.cursor != event.cursor or existing.fingerprint != fingerprint:
raise DirectoryProjectionUnavailableError('Directory event UUID has conflicting contents')
continue
if (
reject_missing_through_cursor is not None
and allow_missing_through_cursor < event.cursor <= reject_missing_through_cursor
):
raise DirectoryProjectionUnavailableError(
'Directory projection cursor advanced without a matching event receipt'
)
session.add(
DirectoryProjectionInbox(
instance_uuid=self.instance_uuid,
event_uuid=event.uuid,
cursor=event.cursor,
event_type=event.event_type,
revision=event.revision,
fingerprint=fingerprint,
received_at=now,
applied_at=None,
)
)
async def _mark_events_applied(
self,
session: Any,
events: Iterable[DirectoryEvent],
*,
now: datetime.datetime,
) -> None:
event_uuids = [event.uuid for event in events]
if not event_uuids:
return
inbox_rows = (
await session.scalars(
sqlalchemy.select(DirectoryProjectionInbox).where(
DirectoryProjectionInbox.instance_uuid == self.instance_uuid,
DirectoryProjectionInbox.event_uuid.in_(event_uuids),
)
)
).all()
if len(inbox_rows) != len(event_uuids):
raise DirectoryProjectionUnavailableError('Directory event receipt could not be persisted')
for row in inbox_rows:
row.applied_at = now
async def _apply_accounts(self, session: Any, snapshot: DirectorySnapshot) -> None:
selected: dict[str, DirectoryMember] = {}
emails: dict[str, str] = {}
for workspace in snapshot.workspaces:
for member in workspace.members:
email_owner = emails.setdefault(member.normalized_email, member.account_uuid)
if email_owner != member.account_uuid:
raise DirectoryProjectionUnavailableError(
'Directory snapshot maps one normalized email to multiple accounts'
)
previous = selected.get(member.account_uuid)
if previous is not None and self._account_projection(previous) != self._account_projection(member):
raise DirectoryProjectionUnavailableError('Directory snapshot has conflicting account projections')
if previous is None:
selected[member.account_uuid] = member
for account_uuid, member in selected.items():
account = await session.scalar(sqlalchemy.select(User).where(User.uuid == account_uuid))
email_account = await session.scalar(
sqlalchemy.select(User).where(User.normalized_email == member.normalized_email)
)
if email_account is not None and email_account.uuid != account_uuid:
raise DirectoryProjectionUnavailableError('Directory account email collides with another Core account')
if account is None:
account = User(
uuid=account_uuid,
user=member.display_name,
normalized_email=member.normalized_email,
password='',
status=_ACCOUNT_STATUS_MAP[member.account_status],
source=AccountSource.CLOUD_PROJECTION.value,
projection_revision=snapshot.cursor,
account_type='space',
space_account_uuid=account_uuid,
)
session.add(account)
continue
if account.source != AccountSource.CLOUD_PROJECTION.value:
raise DirectoryProjectionUnavailableError('Directory account UUID collides with a local Core account')
if account.projection_revision > snapshot.cursor:
raise DirectoryProjectionUnavailableError('Directory account revision rolled back')
projected_account = self._account_projection(member)
persisted_account = self._persisted_account_projection(account)
if account.projection_revision == snapshot.cursor and persisted_account != projected_account:
raise DirectoryProjectionUnavailableError('Directory account revision has conflicting contents')
if persisted_account == projected_account:
# A Workspace rename, role update, or another member's change
# must not revoke this Account's JWT. Account revisions advance
# only when the Account projection itself changes.
continue
account.user = member.display_name
account.normalized_email = member.normalized_email
account.status = _ACCOUNT_STATUS_MAP[member.account_status]
account.projection_revision = snapshot.cursor
account.account_type = 'space'
account.space_account_uuid = account_uuid
await session.flush()
async def _apply_workspaces(self, session: Any, snapshot: DirectorySnapshot) -> None:
for candidate in snapshot.workspaces:
workspace = await session.get(Workspace, candidate.uuid)
if workspace is None:
workspace = Workspace(
uuid=candidate.uuid,
instance_uuid=self.instance_uuid,
name=candidate.name,
slug=candidate.slug,
type=candidate.type,
status=candidate.status,
created_by_account_uuid=await self._projected_creator_uuid(session, candidate),
source=WorkspaceSource.CLOUD_PROJECTION.value,
projection_revision=candidate.projection_revision,
)
session.add(workspace)
await session.flush()
else:
self._validate_existing_workspace(workspace, candidate)
workspace.name = candidate.name
workspace.slug = candidate.slug
workspace.type = candidate.type
workspace.status = candidate.status
workspace.created_by_account_uuid = await self._projected_creator_uuid(session, candidate)
workspace.projection_revision = candidate.projection_revision
await self._apply_memberships(session, workspace, candidate)
await self._apply_execution_state(session, workspace, candidate)
async def _projected_creator_uuid(self, session: Any, candidate: DirectoryWorkspace) -> str | None:
creator = await session.scalar(sqlalchemy.select(User).where(User.uuid == candidate.created_by_account_uuid))
if creator is None:
if candidate.status == WorkspaceStatus.ACTIVE.value:
raise DirectoryProjectionUnavailableError('Active Directory Workspace creator is not projected')
return None
return candidate.created_by_account_uuid
def _validate_existing_workspace(self, workspace: Workspace, candidate: DirectoryWorkspace) -> None:
if workspace.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory Workspace belongs to another LangBot instance')
if workspace.source != WorkspaceSource.CLOUD_PROJECTION.value:
raise DirectoryProjectionUnavailableError('Directory Workspace UUID collides with a local Workspace')
if workspace.projection_revision > candidate.projection_revision:
raise DirectoryProjectionUnavailableError('Directory Workspace revision rolled back')
if workspace.projection_revision == candidate.projection_revision and self._workspace_projection(
workspace
) != self._candidate_workspace_projection(candidate):
raise DirectoryProjectionUnavailableError('Directory Workspace revision has conflicting contents')
async def _apply_memberships(
self,
session: Any,
workspace: Workspace,
candidate: DirectoryWorkspace,
) -> None:
existing = {
membership.account_uuid: membership
for membership in (
await session.scalars(
sqlalchemy.select(WorkspaceMembership).where(WorkspaceMembership.workspace_uuid == workspace.uuid)
)
).all()
}
included_accounts: set[str] = set()
for member in candidate.members:
included_accounts.add(member.account_uuid)
membership = existing.get(member.account_uuid)
joined_at = self._naive_utc(member.joined_at)
role = _ROLE_MAP[member.role]
status = _MEMBERSHIP_STATUS_MAP[member.membership_status]
if candidate.status != WorkspaceStatus.ACTIVE.value:
status = (
MembershipStatus.REMOVED.value
if candidate.status in {WorkspaceStatus.ARCHIVED.value, WorkspaceStatus.DELETED.value}
else MembershipStatus.DISABLED.value
)
if membership is None:
session.add(
WorkspaceMembership(
uuid=member.membership_uuid,
workspace_uuid=workspace.uuid,
account_uuid=member.account_uuid,
role=role,
status=status,
joined_at=joined_at,
projection_revision=member.projection_revision,
)
)
continue
if membership.uuid != member.membership_uuid:
raise DirectoryProjectionUnavailableError('Directory membership UUID changed for one account')
if membership.projection_revision > member.projection_revision:
raise DirectoryProjectionUnavailableError('Directory membership revision rolled back')
if membership.projection_revision == member.projection_revision and self._membership_projection(
membership
) != (role, status, self._datetime_fingerprint(joined_at)):
raise DirectoryProjectionUnavailableError('Directory membership revision has conflicting contents')
membership.role = role
membership.status = status
membership.joined_at = joined_at
membership.projection_revision = member.projection_revision
for account_uuid, membership in existing.items():
if account_uuid not in included_accounts:
membership.status = MembershipStatus.REMOVED.value
membership.projection_revision = max(
int(membership.projection_revision),
candidate.projection_revision,
)
await session.flush()
async def _apply_execution_state(
self,
session: Any,
workspace: Workspace,
candidate: DirectoryWorkspace,
) -> None:
active = candidate.status == WorkspaceStatus.ACTIVE.value
desired_state = WorkspaceExecutionStatus.ACTIVE.value if active else WorkspaceExecutionStatus.INACTIVE.value
execution = await session.get(WorkspaceExecutionState, workspace.uuid)
if execution is None:
session.add(
WorkspaceExecutionState(
workspace_uuid=workspace.uuid,
instance_uuid=self.instance_uuid,
active_generation=candidate.execution_generation,
state=desired_state,
write_fenced=not active,
source=WorkspaceExecutionSource.CLOUD.value,
desired_state_revision=candidate.projection_revision,
)
)
await session.flush()
return
if execution.instance_uuid != self.instance_uuid or execution.source != WorkspaceExecutionSource.CLOUD.value:
raise DirectoryProjectionUnavailableError('Directory execution state has an invalid owner')
if execution.active_generation > candidate.execution_generation:
raise DirectoryProjectionUnavailableError('Directory execution generation rolled back')
if execution.desired_state_revision > candidate.projection_revision:
raise DirectoryProjectionUnavailableError('Directory desired-state revision rolled back')
if execution.desired_state_revision == candidate.projection_revision and (
execution.active_generation != candidate.execution_generation
or execution.state != desired_state
or execution.write_fenced != (not active)
):
raise DirectoryProjectionUnavailableError(
'Directory execution state has conflicting contents at one revision'
)
execution.active_generation = candidate.execution_generation
execution.state = desired_state
execution.write_fenced = not active
execution.desired_state_revision = candidate.projection_revision
await session.flush()
async def _fence_absent_workspaces(self, session: Any, snapshot: DirectorySnapshot) -> None:
included = {workspace.uuid for workspace in snapshot.workspaces}
projected = (
await session.scalars(
sqlalchemy.select(Workspace).where(
Workspace.instance_uuid == self.instance_uuid,
Workspace.source == WorkspaceSource.CLOUD_PROJECTION.value,
)
)
).all()
for workspace in projected:
if workspace.uuid in included:
continue
workspace.status = WorkspaceStatus.ARCHIVED.value
await self._remove_workspace_memberships(session, workspace.uuid)
execution = await session.get(WorkspaceExecutionState, workspace.uuid)
if execution is not None:
execution.state = WorkspaceExecutionStatus.INACTIVE.value
execution.write_fenced = True
await session.flush()
async def _fence_workspaces(self, session: Any, workspace_revisions: dict[str, int]) -> None:
"""Fence requested Workspaces omitted from an authoritative delta."""
if not workspace_revisions:
return
projected = (
await session.scalars(
sqlalchemy.select(Workspace).where(
Workspace.instance_uuid == self.instance_uuid,
Workspace.source == WorkspaceSource.CLOUD_PROJECTION.value,
Workspace.uuid.in_(workspace_revisions),
)
)
).all()
for workspace in projected:
tombstone_revision = workspace_revisions[workspace.uuid]
if int(workspace.projection_revision) > tombstone_revision:
raise DirectoryProjectionUnavailableError('Directory Workspace tombstone revision rolled back')
memberships = (
await session.scalars(
sqlalchemy.select(WorkspaceMembership).where(WorkspaceMembership.workspace_uuid == workspace.uuid)
)
).all()
if any(int(membership.projection_revision) > tombstone_revision for membership in memberships):
raise DirectoryProjectionUnavailableError('Directory membership tombstone revision rolled back')
execution = await session.get(WorkspaceExecutionState, workspace.uuid)
if execution is not None and int(execution.desired_state_revision) > tombstone_revision:
raise DirectoryProjectionUnavailableError('Directory execution tombstone revision rolled back')
workspace.status = WorkspaceStatus.ARCHIVED.value
workspace.projection_revision = max(int(workspace.projection_revision), tombstone_revision)
await self._remove_workspace_memberships(
session,
workspace.uuid,
projection_revision=tombstone_revision,
memberships=memberships,
)
if execution is not None:
execution.state = WorkspaceExecutionStatus.INACTIVE.value
execution.write_fenced = True
execution.desired_state_revision = max(
int(execution.desired_state_revision),
tombstone_revision,
)
await session.flush()
async def _remove_workspace_memberships(
self,
session: Any,
workspace_uuid: str,
*,
projection_revision: int | None = None,
memberships: Iterable[WorkspaceMembership] | None = None,
) -> None:
if memberships is None:
memberships = (
await session.scalars(
sqlalchemy.select(WorkspaceMembership).where(WorkspaceMembership.workspace_uuid == workspace_uuid)
)
).all()
for membership in memberships:
membership.status = MembershipStatus.REMOVED.value
if projection_revision is not None:
membership.projection_revision = max(
int(membership.projection_revision),
projection_revision,
)
def _record_success(self) -> None:
self._last_success_monotonic = self._monotonic_time()
self._ready = True
@classmethod
def _snapshot_fingerprint(cls, snapshot: DirectorySnapshot) -> str:
workspaces = []
for workspace in sorted(snapshot.workspaces, key=lambda item: item.uuid):
data = workspace.model_dump(mode='json')
data['members'] = sorted(data['members'], key=lambda item: item['membership_uuid'])
workspaces.append(data)
return cls._fingerprint(
{
'instance_uuid': snapshot.instance_uuid,
'workspaces': workspaces,
}
)
@staticmethod
def _fingerprint(value: Any) -> str:
encoded = json.dumps(value, sort_keys=True, separators=(',', ':'), ensure_ascii=True).encode()
return hashlib.sha256(encoded).hexdigest()
@staticmethod
def _account_projection(member: DirectoryMember) -> tuple[str, str, str]:
return member.normalized_email, member.display_name, _ACCOUNT_STATUS_MAP[member.account_status]
@staticmethod
def _persisted_account_projection(account: User) -> tuple[str, str, str]:
return account.normalized_email, account.user, account.status
@staticmethod
def _workspace_projection(workspace: Workspace) -> tuple[Any, ...]:
return (
workspace.name,
workspace.slug,
workspace.type,
workspace.status,
workspace.created_by_account_uuid,
)
@staticmethod
def _candidate_workspace_projection(candidate: DirectoryWorkspace) -> tuple[Any, ...]:
return (
candidate.name,
candidate.slug,
candidate.type,
candidate.status,
candidate.created_by_account_uuid,
)
@classmethod
def _membership_projection(cls, membership: WorkspaceMembership) -> tuple[Any, ...]:
return (
membership.role,
membership.status,
cls._datetime_fingerprint(membership.joined_at),
)
@staticmethod
def _datetime_fingerprint(value: datetime.datetime | None) -> str | None:
if value is None:
return None
return DirectoryProjectionService._naive_utc(value).isoformat(timespec='microseconds')
@staticmethod
def _naive_utc(value: datetime.datetime | None) -> datetime.datetime | None:
if value is None:
return None
if value.tzinfo is None:
return value
return value.astimezone(datetime.UTC).replace(tzinfo=None)
@staticmethod
def _utcnow() -> datetime.datetime:
return datetime.datetime.now(datetime.UTC)
+22
View File
@@ -48,6 +48,7 @@ from ..skill import manager as skill_mgr
from ..workspace import service as workspace_service_module
from ..workspace import collaboration as workspace_collaboration_module
from ..cloud import bootstrap as cloud_bootstrap_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
@@ -132,8 +133,12 @@ class Application:
deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None
manifest_refresh_service: cloud_bootstrap_module.CloudManifestRefreshService | None = None
entitlement_resolver: cloud_entitlements_module.EntitlementResolver | None = None
directory_projection_service: cloud_directory_projection_module.DirectoryProjectionService | None = None
vector_db_mgr: vectordb_mgr.VectorDBManager = None
http_ctrl: http_controller.HTTPController = None
@@ -189,6 +194,19 @@ class Application:
async def run(self):
try:
if self.directory_projection_service is not None:
self.task_mgr.create_task(
self.directory_projection_service.run(),
name='cloud-directory-projection',
scopes=[core_entities.LifecycleControlScope.APPLICATION],
)
if self.manifest_refresh_service is not None:
self.task_mgr.create_task(
self.manifest_refresh_service.run(),
name='cloud-manifest-refresh',
scopes=[core_entities.LifecycleControlScope.APPLICATION],
)
await self.plugin_connector.initialize_plugins()
# 后续可能会允许动态重启其他任务
@@ -374,6 +392,10 @@ class Application:
if self.plugin_connector is not None:
with contextlib.suppress(Exception):
await self.plugin_connector.aclose()
manifest_provider = getattr(self.deployment, 'manifest_provider', None)
if manifest_provider is not None:
with contextlib.suppress(Exception):
await manifest_provider.aclose()
if self.task_mgr is not None:
tasks = [wrapper.task for wrapper in self.task_mgr.tasks if not wrapper.task.done()]
+19
View File
@@ -40,6 +40,7 @@ from ...survey import manager as survey_module
from ...workspace import service as workspace_service_module
from ...workspace import collaboration as workspace_collaboration_module
from ...cloud import bootstrap as cloud_bootstrap
from ...cloud.directory_projection import DirectoryProjectionService
from ...cloud.entitlements import EntitlementResolver
from ...api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from ...api.http.authz import WorkspaceRequiredError
@@ -63,6 +64,15 @@ class BuildAppStage(stage.BootingStage):
constants.instance_id,
deployment,
)
ap.manifest_refresh_service = (
cloud_bootstrap.CloudManifestRefreshService(
ap.deployment_admission,
deployment.manifest_provider,
ap.logger,
)
if deployment.multi_workspace_enabled
else None
)
ap.entitlement_resolver = (
EntitlementResolver(
constants.instance_id,
@@ -137,6 +147,15 @@ class BuildAppStage(stage.BootingStage):
ap.persistence_mgr = persistence_mgr_inst
await persistence_mgr_inst.initialize()
if deployment.multi_workspace_enabled:
directory_projection_service = DirectoryProjectionService(
ap,
deployment.directory_provider,
constants.instance_id,
)
await directory_projection_service.initialize()
ap.directory_projection_service = directory_projection_service
workspace_policy = deployment.workspace_policy
workspace_service_inst = workspace_service_module.WorkspaceService(
ap,
@@ -0,0 +1,69 @@
from __future__ import annotations
import sqlalchemy
from .base import Base
class DirectoryProjectionState(Base):
"""Durable cursor and lease for one verified Cloud directory."""
__tablename__ = 'directory_projection_states'
instance_uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
cursor = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False, server_default='0')
snapshot_coverage_cursor = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False, server_default='0')
snapshot_fingerprint = sqlalchemy.Column(sqlalchemy.Text, nullable=False)
last_applied_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), nullable=False)
lease_expires_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), nullable=True)
__table_args__ = (
sqlalchemy.CheckConstraint('cursor >= 0', name='ck_directory_projection_state_cursor'),
sqlalchemy.CheckConstraint(
'snapshot_coverage_cursor >= 0 AND snapshot_coverage_cursor <= cursor',
name='ck_directory_projection_state_snapshot_coverage',
),
sqlalchemy.CheckConstraint(
'length(snapshot_fingerprint) = 64',
name='ck_directory_projection_state_fingerprint',
),
)
class DirectoryProjectionInbox(Base):
"""Idempotency ledger for signed control-plane directory events."""
__tablename__ = 'directory_projection_inbox'
instance_uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
event_uuid = sqlalchemy.Column(sqlalchemy.String(36), primary_key=True)
cursor = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False)
event_type = sqlalchemy.Column(sqlalchemy.String(128), nullable=False)
revision = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False)
fingerprint = sqlalchemy.Column(sqlalchemy.Text, nullable=False)
received_at = sqlalchemy.Column(
sqlalchemy.DateTime(timezone=True),
nullable=False,
server_default=sqlalchemy.func.now(),
)
applied_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), nullable=True)
__table_args__ = (
sqlalchemy.UniqueConstraint(
'instance_uuid',
'cursor',
name='uq_directory_projection_inbox_cursor',
),
sqlalchemy.Index(
'ix_directory_projection_inbox_pending',
'instance_uuid',
'applied_at',
'cursor',
),
sqlalchemy.CheckConstraint('cursor > 0', name='ck_directory_projection_inbox_cursor'),
sqlalchemy.CheckConstraint('revision > 0', name='ck_directory_projection_inbox_revision'),
sqlalchemy.CheckConstraint(
'length(fingerprint) = 64',
name='ck_directory_projection_inbox_fingerprint',
),
)
@@ -0,0 +1,268 @@
"""add the Cloud directory projection persistence boundary
Revision ID: 0014_cloud_directory
Revises: 0013_tenant_pgvector
Create Date: 2026-07-24
The open Core projector receives already-verified control-plane data and is the
only runtime path allowed to mutate projected Workspace directory rows. Its
transaction-local instance setting is intentionally distinct from both normal
Workspace scope and the read-only instance discovery scope.
"""
from __future__ import annotations
import sqlalchemy as sa
from alembic import op
revision = '0014_cloud_directory'
down_revision = '0013_tenant_pgvector'
branch_labels = None
depends_on = None
_STATE_TABLE = 'directory_projection_states'
_INBOX_TABLE = 'directory_projection_inbox'
_DIRECTORY_POLICY_NAME = 'langbot_directory_projection'
_TENANT_POLICY_NAME = 'langbot_workspace_isolation'
_LOCAL_WRITE_POLICY_NAME = 'langbot_workspace_local_directory_write'
_DIRECTORY_SETTING = 'langbot.directory_instance_uuid'
_TENANT_SETTING = 'langbot.workspace_uuid'
_PROJECTED_TENANT_TABLES = (
'workspaces',
'workspace_memberships',
'workspace_execution_states',
)
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 _create_tables(conn: sa.Connection) -> None:
existing_tables = set(sa.inspect(conn).get_table_names())
if _STATE_TABLE not in existing_tables:
op.create_table(
_STATE_TABLE,
sa.Column('instance_uuid', sa.String(255), nullable=False),
sa.Column('cursor', sa.BigInteger(), server_default='0', nullable=False),
sa.Column('snapshot_coverage_cursor', sa.BigInteger(), server_default='0', nullable=False),
sa.Column('snapshot_fingerprint', sa.Text(), nullable=False),
sa.Column('last_applied_at', sa.DateTime(timezone=True), nullable=False),
sa.Column('lease_expires_at', sa.DateTime(timezone=True), nullable=True),
sa.CheckConstraint(
'cursor >= 0',
name='ck_directory_projection_state_cursor',
),
sa.CheckConstraint(
'snapshot_coverage_cursor >= 0 AND snapshot_coverage_cursor <= cursor',
name='ck_directory_projection_state_snapshot_coverage',
),
sa.CheckConstraint(
'length(snapshot_fingerprint) = 64',
name='ck_directory_projection_state_fingerprint',
),
sa.PrimaryKeyConstraint('instance_uuid'),
)
if _INBOX_TABLE not in existing_tables:
op.create_table(
_INBOX_TABLE,
sa.Column('instance_uuid', sa.String(255), nullable=False),
sa.Column('event_uuid', sa.String(36), nullable=False),
sa.Column('cursor', sa.BigInteger(), nullable=False),
sa.Column('event_type', sa.String(128), nullable=False),
sa.Column('revision', sa.BigInteger(), nullable=False),
sa.Column('fingerprint', sa.Text(), nullable=False),
sa.Column(
'received_at',
sa.DateTime(timezone=True),
server_default=sa.func.now(),
nullable=False,
),
sa.Column('applied_at', sa.DateTime(timezone=True), nullable=True),
sa.CheckConstraint(
'cursor > 0',
name='ck_directory_projection_inbox_cursor',
),
sa.CheckConstraint(
'revision > 0',
name='ck_directory_projection_inbox_revision',
),
sa.CheckConstraint(
'length(fingerprint) = 64',
name='ck_directory_projection_inbox_fingerprint',
),
sa.PrimaryKeyConstraint('instance_uuid', 'event_uuid'),
sa.UniqueConstraint(
'instance_uuid',
'cursor',
name='uq_directory_projection_inbox_cursor',
),
)
op.create_index(
'ix_directory_projection_inbox_pending',
_INBOX_TABLE,
['instance_uuid', 'applied_at', 'cursor'],
unique=False,
)
def _drop_policy(conn: sa.Connection, table_name: str, policy_name: str) -> None:
table = _quote(conn, table_name)
policy = _quote(conn, policy_name)
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
def _create_policy(
conn: sa.Connection,
table_name: str,
policy_name: str,
expression: str,
*,
command: str = 'ALL',
) -> None:
table = _quote(conn, table_name)
policy = _quote(conn, policy_name)
_drop_policy(conn, table_name, policy_name)
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
if command == 'SELECT':
sql = f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR SELECT TO PUBLIC USING ({expression})'
elif command == 'ALL':
sql = (
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
f'USING ({expression}) WITH CHECK ({expression})'
)
else: # pragma: no cover - migration-local invariant.
raise AssertionError(f'Unsupported RLS policy command: {command}')
op.execute(sa.text(sql))
def _install_postgres_policies(conn: sa.Connection) -> None:
existing_tables = set(sa.inspect(conn).get_table_names())
required_tables = set(_PROJECTED_TENANT_TABLES) | {_STATE_TABLE, _INBOX_TABLE}
missing_tables = required_tables - existing_tables
if missing_tables:
raise RuntimeError(
f'Cannot enable Cloud directory projection RLS before all required tables exist: {sorted(missing_tables)!r}'
)
directory_setting = _setting(_DIRECTORY_SETTING)
tenant_setting = _setting(_TENANT_SETTING)
directory_expressions = {
'workspaces': (f"instance_uuid::text = {directory_setting} AND source = 'cloud_projection'"),
'workspace_memberships': (
'EXISTS ('
'SELECT 1 FROM workspaces AS directory_workspace '
'WHERE directory_workspace.uuid = workspace_memberships.workspace_uuid '
f'AND directory_workspace.instance_uuid::text = {directory_setting} '
"AND directory_workspace.source = 'cloud_projection'"
')'
),
'workspace_execution_states': (
f"instance_uuid::text = {directory_setting} AND source = 'cloud' AND EXISTS ("
'SELECT 1 FROM workspaces AS directory_workspace '
'WHERE directory_workspace.uuid = workspace_execution_states.workspace_uuid '
f'AND directory_workspace.instance_uuid::text = {directory_setting} '
"AND directory_workspace.source = 'cloud_projection'"
')'
),
_STATE_TABLE: f'instance_uuid::text = {directory_setting}',
_INBOX_TABLE: f'instance_uuid::text = {directory_setting}',
}
tenant_expressions = {
'workspaces': f'uuid::text = {tenant_setting}',
'workspace_memberships': f'workspace_uuid::text = {tenant_setting}',
'workspace_execution_states': f'workspace_uuid::text = {tenant_setting}',
}
local_write_expressions = {
'workspaces': f"uuid::text = {tenant_setting} AND source = 'local'",
'workspace_memberships': (
f'workspace_uuid::text = {tenant_setting} AND EXISTS ('
'SELECT 1 FROM workspaces AS local_workspace '
'WHERE local_workspace.uuid = workspace_memberships.workspace_uuid '
"AND local_workspace.source = 'local'"
')'
),
'workspace_execution_states': (
f'workspace_uuid::text = {tenant_setting} AND EXISTS ('
'SELECT 1 FROM workspaces AS local_workspace '
'WHERE local_workspace.uuid = workspace_execution_states.workspace_uuid '
"AND local_workspace.source = 'local'"
')'
),
}
for table_name in _PROJECTED_TENANT_TABLES:
_create_policy(
conn,
table_name,
_TENANT_POLICY_NAME,
tenant_expressions[table_name],
command='SELECT',
)
_create_policy(
conn,
table_name,
_LOCAL_WRITE_POLICY_NAME,
local_write_expressions[table_name],
)
_create_policy(
conn,
table_name,
_DIRECTORY_POLICY_NAME,
directory_expressions[table_name],
)
for table_name in (_STATE_TABLE, _INBOX_TABLE):
_create_policy(
conn,
table_name,
_DIRECTORY_POLICY_NAME,
directory_expressions[table_name],
)
def upgrade() -> None:
conn = op.get_bind()
_create_tables(conn)
if conn.dialect.name == 'postgresql':
_install_postgres_policies(conn)
def downgrade() -> None:
conn = op.get_bind()
existing_tables = set(sa.inspect(conn).get_table_names())
if conn.dialect.name == 'postgresql':
tenant_setting = _setting(_TENANT_SETTING)
tenant_columns = {
'workspaces': 'uuid',
'workspace_memberships': 'workspace_uuid',
'workspace_execution_states': 'workspace_uuid',
}
for table_name in _PROJECTED_TENANT_TABLES:
if table_name not in existing_tables:
continue
_drop_policy(conn, table_name, _DIRECTORY_POLICY_NAME)
_drop_policy(conn, table_name, _LOCAL_WRITE_POLICY_NAME)
_create_policy(
conn,
table_name,
_TENANT_POLICY_NAME,
f'{tenant_columns[table_name]}::text = {tenant_setting}',
)
for table_name in (_STATE_TABLE, _INBOX_TABLE):
if table_name not in existing_tables:
continue
_drop_policy(conn, table_name, _DIRECTORY_POLICY_NAME)
table = _quote(conn, table_name)
op.execute(sa.text(f'ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY'))
op.execute(sa.text(f'ALTER TABLE {table} DISABLE ROW LEVEL SECURITY'))
if _INBOX_TABLE in existing_tables:
op.drop_table(_INBOX_TABLE)
if _STATE_TABLE in existing_tables:
op.drop_table(_STATE_TABLE)
+90 -4
View File
@@ -27,10 +27,15 @@ from .tenant_uow import (
ACCOUNT_DISCOVERY_POLICY_NAME,
INSTANCE_DISCOVERY_POLICY_NAME,
INVITATION_DISCOVERY_POLICY_NAME,
LOCAL_DIRECTORY_WRITE_POLICY_NAME,
TENANT_POLICY_NAME,
TENANT_SETTING,
TENANT_TABLE_COLUMNS,
CrossScopeTransactionError,
DIRECTORY_INSTANCE_SETTING,
DIRECTORY_PROJECTED_TENANT_TABLES,
DIRECTORY_PROJECTION_POLICY_NAME,
DIRECTORY_PROJECTION_TABLE_COLUMNS,
PersistenceScope,
PersistenceScopeBoundary,
PersistenceScopeKind,
@@ -74,6 +79,8 @@ _ALEMBIC_TENANT_TABLES = {
'monitoring_embedding_calls',
'monitoring_feedback',
'langbot_vectors',
'directory_projection_states',
'directory_projection_inbox',
}
_PRE_WORKSPACE_ALEMBIC_REVISIONS = {
@@ -1345,6 +1352,7 @@ class PersistenceManager:
engine = self.get_db_engine()
if engine.dialect.name != 'postgresql':
raise RuntimeError('PostgreSQL tenant schema validation requires PostgreSQL')
rls_table_names = tuple(sorted(set(TENANT_TABLE_COLUMNS) | set(DIRECTORY_PROJECTION_TABLE_COLUMNS)))
table_query = sqlalchemy.text(
"""
@@ -1408,7 +1416,7 @@ class PersistenceManager:
await conn.execute(
table_query,
{
'table_names': tuple(TENANT_TABLE_COLUMNS),
'table_names': rls_table_names,
},
)
)
@@ -1419,7 +1427,7 @@ class PersistenceManager:
(
await conn.execute(
policy_query,
{'table_names': tuple(TENANT_TABLE_COLUMNS)},
{'table_names': rls_table_names},
)
)
.mappings()
@@ -1427,7 +1435,7 @@ class PersistenceManager:
)
by_table = {row['table_name']: row for row in rows}
missing_tables = set(TENANT_TABLE_COLUMNS) - set(by_table)
missing_tables = set(rls_table_names) - set(by_table)
if missing_tables:
raise RuntimeError(f'PostgreSQL tenant tables are missing: {sorted(missing_tables)!r}')
@@ -1472,7 +1480,7 @@ class PersistenceManager:
@staticmethod
def _expected_postgres_tenant_policies() -> dict[str, dict[str, dict[str, str | None]]]:
"""Return the exact PostgreSQL 16 policy expressions emitted by 0011."""
"""Return the exact PostgreSQL 16 policy expressions emitted by 0011/0014."""
def setting(name: str) -> str:
return f"NULLIF(current_setting('{name}'::text, true), ''::text)"
@@ -1488,6 +1496,39 @@ class PersistenceManager:
}
}
local_workspace_expression = (
f"(((uuid)::text = {setting(TENANT_SETTING)}) AND ((source)::text = 'local'::text))"
)
local_membership_expression = (
f'(((workspace_uuid)::text = {setting(TENANT_SETTING)}) AND (EXISTS ( SELECT 1\n'
' FROM workspaces local_workspace\n'
' WHERE (((local_workspace.uuid)::text = (workspace_memberships.workspace_uuid)::text) '
"AND ((local_workspace.source)::text = 'local'::text)))))"
)
local_execution_expression = (
f'(((workspace_uuid)::text = {setting(TENANT_SETTING)}) AND (EXISTS ( SELECT 1\n'
' FROM workspaces local_workspace\n'
' WHERE (((local_workspace.uuid)::text = (workspace_execution_states.workspace_uuid)::text) '
"AND ((local_workspace.source)::text = 'local'::text)))))"
)
for table_name, local_write_expression in {
'workspaces': local_workspace_expression,
'workspace_memberships': local_membership_expression,
'workspace_execution_states': local_execution_expression,
}.items():
tenant_column = TENANT_TABLE_COLUMNS[table_name]
tenant_expression = f'(({tenant_column})::text = {setting(TENANT_SETTING)})'
policies[table_name][TENANT_POLICY_NAME] = {
'command': 'r',
'using_expression': tenant_expression,
'check_expression': None,
}
policies[table_name][LOCAL_DIRECTORY_WRITE_POLICY_NAME] = {
'command': '*',
'using_expression': local_write_expression,
'check_expression': local_write_expression,
}
policies['workspace_memberships'][ACCOUNT_DISCOVERY_POLICY_NAME] = {
'command': 'r',
'using_expression': (
@@ -1517,6 +1558,48 @@ class PersistenceManager:
),
'check_expression': None,
}
directory_setting = setting(DIRECTORY_INSTANCE_SETTING)
workspace_expression = (
f"(((instance_uuid)::text = {directory_setting}) AND ((source)::text = 'cloud_projection'::text))"
)
membership_expression = (
'(EXISTS ( SELECT 1\n'
' FROM workspaces directory_workspace\n'
' WHERE (((directory_workspace.uuid)::text = (workspace_memberships.workspace_uuid)::text) '
f'AND ((directory_workspace.instance_uuid)::text = {directory_setting}) '
"AND ((directory_workspace.source)::text = 'cloud_projection'::text))))"
)
execution_expression = (
f'(((instance_uuid)::text = {directory_setting}) '
"AND ((source)::text = 'cloud'::text) "
'AND (EXISTS ( SELECT 1\n'
' FROM workspaces directory_workspace\n'
' WHERE (((directory_workspace.uuid)::text = (workspace_execution_states.workspace_uuid)::text) '
f'AND ((directory_workspace.instance_uuid)::text = {directory_setting}) '
"AND ((directory_workspace.source)::text = 'cloud_projection'::text)))))"
)
directory_tenant_expressions = {
'workspaces': workspace_expression,
'workspace_memberships': membership_expression,
'workspace_execution_states': execution_expression,
}
for table_name in DIRECTORY_PROJECTED_TENANT_TABLES:
expression = directory_tenant_expressions[table_name]
policies[table_name][DIRECTORY_PROJECTION_POLICY_NAME] = {
'command': '*',
'using_expression': expression,
'check_expression': expression,
}
for table_name, instance_column in DIRECTORY_PROJECTION_TABLE_COLUMNS.items():
expression = f'(({instance_column})::text = {directory_setting})'
policies[table_name] = {
DIRECTORY_PROJECTION_POLICY_NAME: {
'command': '*',
'using_expression': expression,
'check_expression': expression,
}
}
return policies
async def write_space_model_providers(self):
@@ -1729,6 +1812,9 @@ class PersistenceManager:
def instance_discovery_uow(self, instance_uuid: str) -> TenantUnitOfWork:
return self._scoped_uow(PersistenceScope.instance(instance_uuid))
def directory_projection_uow(self, instance_uuid: str) -> TenantUnitOfWork:
return self._scoped_uow(PersistenceScope.directory(instance_uuid))
def identity_discovery_uow(self, identity_digest: str) -> TenantUnitOfWork:
return self._scoped_uow(PersistenceScope.identity(identity_digest))
+38 -1
View File
@@ -14,7 +14,9 @@ import sqlalchemy
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
import sqlalchemy.orm as sqlalchemy_orm
from pgvector.sqlalchemy import Vector
from sqlalchemy.dialects.postgresql.dml import OnConflictDoNothing as PostgreSQLOnConflictDoNothing
from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLOnConflictDoUpdate
from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing
from sqlalchemy.dialects.sqlite.dml import OnConflictDoUpdate as SQLiteOnConflictDoUpdate
@@ -24,12 +26,15 @@ API_KEY_HASH_SETTING = 'langbot.api_key_hash'
INVITATION_HASH_SETTING = 'langbot.invitation_hash'
INSTANCE_SETTING = 'langbot.instance_uuid'
IDENTITY_DIGEST_SETTING = 'langbot.identity_digest'
DIRECTORY_INSTANCE_SETTING = 'langbot.directory_instance_uuid'
TENANT_POLICY_NAME = 'langbot_workspace_isolation'
LOCAL_DIRECTORY_WRITE_POLICY_NAME = 'langbot_workspace_local_directory_write'
ACCOUNT_DISCOVERY_POLICY_NAME = 'langbot_account_discovery'
API_KEY_DISCOVERY_POLICY_NAME = 'langbot_api_key_discovery'
INVITATION_DISCOVERY_POLICY_NAME = 'langbot_invitation_discovery'
INSTANCE_DISCOVERY_POLICY_NAME = 'langbot_instance_discovery'
DIRECTORY_PROJECTION_POLICY_NAME = 'langbot_directory_projection'
# Keep this contract explicit. A new tenant-owned table must be added to both
# this runtime list and the corresponding Alembic migration before release.
@@ -67,6 +72,19 @@ TENANT_TABLE_COLUMNS: dict[str, str] = {
'langbot_vectors': 'workspace_uuid',
}
DIRECTORY_PROJECTION_TABLE_COLUMNS: dict[str, str] = {
'directory_projection_states': 'instance_uuid',
'directory_projection_inbox': 'instance_uuid',
}
DIRECTORY_PROJECTED_TENANT_TABLES = frozenset(
{
'workspaces',
'workspace_memberships',
'workspace_execution_states',
}
)
class PersistenceScopeKind(enum.StrEnum):
WORKSPACE = 'workspace'
@@ -75,6 +93,7 @@ class PersistenceScopeKind(enum.StrEnum):
INVITATION_DISCOVERY = 'invitation_discovery'
INSTANCE_DISCOVERY = 'instance_discovery'
IDENTITY_DISCOVERY = 'identity_discovery'
DIRECTORY_PROJECTION = 'directory_projection'
@dataclasses.dataclass(frozen=True, slots=True)
@@ -108,6 +127,14 @@ class PersistenceScope:
def identity(cls, identity_digest: str) -> PersistenceScope:
return cls._one(PersistenceScopeKind.IDENTITY_DISCOVERY, IDENTITY_DIGEST_SETTING, identity_digest)
@classmethod
def directory(cls, instance_uuid: str) -> PersistenceScope:
return cls._one(
PersistenceScopeKind.DIRECTORY_PROJECTION,
DIRECTORY_INSTANCE_SETTING,
instance_uuid,
)
@classmethod
def _one(cls, kind: PersistenceScopeKind, setting: str, value: str) -> PersistenceScope:
return cls(kind, ((setting, cls._normalize(value, kind.value)),))
@@ -188,7 +215,9 @@ _ALLOWED_SCOPED_STATEMENT_TYPES = (
sqlalchemy.sql.selectable.SelectBase,
)
_ALLOWED_SCOPED_POST_VALUES_TYPES = (
PostgreSQLOnConflictDoNothing,
PostgreSQLOnConflictDoUpdate,
SQLiteOnConflictDoNothing,
SQLiteOnConflictDoUpdate,
)
@@ -853,6 +882,7 @@ class TenantScopedAsyncSession(sqlalchemy_asyncio.AsyncSession):
INVITATION_HASH_SETTING,
INSTANCE_SETTING,
IDENTITY_DIGEST_SETTING,
DIRECTORY_INSTANCE_SETTING,
}:
self._reject_transaction_escape('unknown tenant scope configuration')
self._require_owner_task()
@@ -1230,7 +1260,14 @@ class TenantUnitOfWork:
# It must still never switch directly to another Workspace.
workspace_to_discovery = (
active_scope.scope.kind == PersistenceScopeKind.WORKSPACE
and self.scope.kind != PersistenceScopeKind.WORKSPACE
and self.scope.kind
in {
PersistenceScopeKind.ACCOUNT_DISCOVERY,
PersistenceScopeKind.API_KEY_DISCOVERY,
PersistenceScopeKind.INVITATION_DISCOVERY,
PersistenceScopeKind.INSTANCE_DISCOVERY,
PersistenceScopeKind.IDENTITY_DISCOVERY,
}
)
if not workspace_to_discovery:
raise CrossScopeTransactionError(
+287 -177
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
import contextlib
import contextvars
import hashlib
import io
@@ -73,6 +74,10 @@ _GITHUB_ASSET_HOSTS = frozenset(
}
)
_HTTP_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308})
_CONNECT_TIMEOUT_SEC = 30.0
_HEARTBEAT_INTERVAL_SEC = 20.0
_HEARTBEAT_FAILURE_THRESHOLD = 3
_RECONNECT_MAX_DELAY_SEC = 60.0
class PluginRuntimeNotConnectedError(RuntimeError):
@@ -139,6 +144,24 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
self._installation_failures: dict[str, dict[str, str]] = {}
self._state_lock = asyncio.Lock()
self._control_token = str(os.environ.get(PLUGIN_RUNTIME_CONTROL_TOKEN_ENV) or '').strip()
self._transport_task: asyncio.Task | None = None
self._reconnect_task: asyncio.Task | None = None
self._generation = 0
self._connected = asyncio.Event()
def _runtime_handler(self) -> handler.RuntimeConnectionHandler:
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is None:
raise PluginRuntimeNotConnectedError('Plugin runtime is not connected')
return runtime_handler
def _runtime_available(self) -> bool:
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is None:
return False
# Unit-level and explicitly injected handlers do not own a transport.
# A managed transport must also have completed its handshake.
return self._transport_task is None or self._connected.is_set()
def _load_worker_policy(self) -> PluginWorkerPolicy:
"""Validate the instance policy without consulting plugin manifests."""
@@ -351,8 +374,9 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
desired_states: list[PluginInstallationDesiredState] = []
for setting in settings:
binding = self._binding_from_setting(execution_context, setting)
if hasattr(self, 'handler'):
self.handler.register_installation_binding(
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is not None:
runtime_handler.register_installation_binding(
binding,
plugin_author=setting.plugin_author,
plugin_name=setting.plugin_name,
@@ -371,7 +395,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
*,
artifact_package: bytes | None = None,
) -> dict[str, Any]:
result = await self.handler.apply_plugin_installation(
runtime_handler = self._runtime_handler()
result = await runtime_handler.apply_plugin_installation(
desired.binding,
artifact_package=artifact_package,
enabled=desired.enabled,
@@ -396,7 +421,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
f'Durable plugin artifact {desired.binding.artifact_digest} is missing for '
f'installation {desired.binding.installation_uuid}'
)
repaired = await self.handler.apply_plugin_installation(
repaired = await runtime_handler.apply_plugin_installation(
desired.binding,
artifact_package=persisted_package,
enabled=desired.enabled,
@@ -559,6 +584,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
async def _prepare_connected_runtime(self) -> None:
"""Handshake follow-up: pin OSS compatibility, then replay authority."""
runtime_handler = self._runtime_handler()
workspace_service = getattr(self.ap, 'workspace_service', None)
if workspace_service is None:
raise RuntimeError('Plugin Runtime requires the Workspace projection service')
@@ -580,15 +606,15 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
# One fully-bound action releases the SDK's deliberately retained
# pre-v4 data/plugins/debug compatibility path. Shared mode never
# creates this bridge.
with self.handler.installation_scope(bridge):
await self.handler.list_plugins()
with runtime_handler.installation_scope(bridge):
await runtime_handler.list_plugins()
desired_states = await self._load_workspace_desired_states(execution_context)
self._workspace_installations[execution_context.workspace_uuid] = {
state.binding.installation_uuid for state in desired_states
}
self._known_desired_states.update({state.binding.installation_uuid: state for state in desired_states})
result = await self.handler.reconcile_plugin_installations(tuple(self._known_desired_states.values()))
result = await runtime_handler.reconcile_plugin_installations(tuple(self._known_desired_states.values()))
await self._repair_reconcile_missing_artifacts(self._known_desired_states, result)
self._record_reconcile_failures(self._known_desired_states, result)
@@ -604,8 +630,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
Runtime reconcile.
"""
if not hasattr(self, 'handler'):
raise PluginRuntimeNotConnectedError('Plugin runtime is not connected')
runtime_handler = self._runtime_handler()
async with self._state_lock:
all_states: dict[str, PluginInstallationDesiredState] = {}
workspace_installations: dict[str, set[str]] = {}
@@ -619,12 +644,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if state.binding.installation_uuid in all_states:
raise ValueError('Duplicate plugin installation UUID across projected Workspaces')
all_states[state.binding.installation_uuid] = state
result = await self.handler.reconcile_plugin_installations(tuple(all_states.values()))
result = await runtime_handler.reconcile_plugin_installations(tuple(all_states.values()))
await self._repair_reconcile_missing_artifacts(all_states, result)
self._record_reconcile_failures(all_states, result)
for installation_uuid, previous in tuple(self._known_desired_states.items()):
if installation_uuid not in all_states:
self.handler.unregister_installation_binding(previous.binding)
runtime_handler.unregister_installation_binding(previous.binding)
self._known_desired_states = all_states
self._workspace_installations = workspace_installations
return result
@@ -652,6 +677,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
async def _synchronize_workspace(self, execution_context: ExecutionContext) -> None:
if not self.is_enable_plugin or not hasattr(self, 'handler'):
return
runtime_handler = self._runtime_handler()
desired_states = await self._load_workspace_desired_states(execution_context)
desired_by_uuid = {state.binding.installation_uuid: state for state in desired_states}
async with self._state_lock:
@@ -659,8 +685,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
for installation_uuid in previous_ids - set(desired_by_uuid):
previous = self._known_desired_states.get(installation_uuid)
if previous is not None:
await self.handler.remove_plugin_installation(previous.binding)
self.handler.unregister_installation_binding(previous.binding)
await runtime_handler.remove_plugin_installation(previous.binding)
runtime_handler.unregister_installation_binding(previous.binding)
self._known_desired_states.pop(installation_uuid, None)
self._installation_failures.pop(installation_uuid, None)
@@ -755,48 +781,152 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
await notify_disconnect()
return False
self.handler = handler.RuntimeConnectionHandler(
connection,
disconnect_callback,
self.ap,
)
self.handler_task = asyncio.create_task(self.handler.run())
_ = await self.handler.ping()
# Push the configured marketplace (Space) URL to the runtime so it
# downloads plugins from the same Space LangBot is bound to, rather
# than relying on the runtime's own env/default.
space_url = self.ap.instance_config.data.get('space', {}).get('url', '').rstrip('/')
try:
if self.runtime_identity is None or self.worker_policy is None: # pragma: no cover
raise RuntimeError('Plugin Runtime identity or worker policy was not loaded')
await self.handler.set_runtime_config(
runtime_identity=self.runtime_identity,
worker_policy=self.worker_policy,
runtime_profile=self.runtime_profile,
cloud_service_url=space_url or None,
runtime_handler = handler.RuntimeConnectionHandler(
connection,
disconnect_callback,
self.ap,
)
if space_url:
self.ap.logger.info(f'Pushed marketplace URL to plugin runtime: {space_url}')
except Exception as e:
self.ap.logger.warning(f'Failed to bind plugin runtime config: {e}')
raise
await self._prepare_connected_runtime()
self.ap.logger.info('Connected to instance-scoped plugin runtime.')
await self.handler_task
self.handler = runtime_handler
self.handler_task = asyncio.create_task(runtime_handler.run())
try:
await runtime_handler.ping()
if self.runtime_identity is None or self.worker_policy is None: # pragma: no cover
raise RuntimeError('Plugin Runtime identity or worker policy was not loaded')
space_url = self.ap.instance_config.data.get('space', {}).get('url', '').rstrip('/')
await runtime_handler.set_runtime_config(
runtime_identity=self.runtime_identity,
worker_policy=self.worker_policy,
runtime_profile=self.runtime_profile,
cloud_service_url=space_url or None,
)
if space_url:
self.ap.logger.info(f'Pushed marketplace URL to plugin runtime: {space_url}')
await self._prepare_connected_runtime()
if generation == self._generation and not self._closing:
connection_ready = True
self._connected.set()
self.ap.logger.info('Connected to instance-scoped plugin runtime.')
await self.handler_task
except asyncio.CancelledError:
raise
except Exception as exc:
if not self._connected.is_set():
connect_errors.append(exc)
self._connected.set()
finally:
if generation == self._generation and not self._closing:
self._connected.clear()
if getattr(self, 'handler', None) is runtime_handler:
del self.handler
await notify_disconnect()
task_coro: typing.Coroutine
task_coro: typing.Coroutine[Any, Any, Any]
if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime():
self.ap.logger.info('use websocket to connect to plugin runtime')
control_headers = self._control_headers(allow_generate=False)
ws_url = self.ap.instance_config.data.get('plugin', {}).get(
'runtime_ws_url',
'ws://langbot_plugin_runtime:5400/control/ws',
)
if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime(): # use websocket
self.ap.logger.info('use websocket to connect to plugin runtime')
control_headers = self._control_headers(allow_generate=False)
ws_url = self.ap.instance_config.data.get('plugin', {}).get(
'runtime_ws_url', 'ws://langbot_plugin_runtime:5400/control/ws'
async def connection_failed(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception | None = None,
) -> None:
del ctrl
connect_errors.append(exc or RuntimeError('WebSocket connection failed'))
self._connected.set()
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=connection_failed,
additional_headers=control_headers,
)
task_coro = self.ctrl.run(new_connection_callback)
elif platform.get_platform() == 'win32':
# Windows cannot use the stdio subprocess transport, so launch
# a managed runtime and authenticate its WebSocket controller.
self.ap.logger.info('(windows) use cmd to launch plugin runtime and communicate via ws')
control_headers = self._control_headers(allow_generate=True)
await self._start_runtime_subprocess(
'-m',
'langbot_plugin.cli.__init__',
'rt',
env_overrides={PLUGIN_RUNTIME_CONTROL_TOKEN_ENV: self._control_token},
)
ws_url = 'ws://localhost:5400/control/ws'
async def connection_failed(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception | None = None,
) -> None:
del ctrl
connect_errors.append(exc or RuntimeError('WebSocket connection failed'))
self._connected.set()
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=connection_failed,
additional_headers=control_headers,
)
task_coro = self.ctrl.run(new_connection_callback)
else:
self.ap.logger.info('use stdio to connect to plugin runtime')
self.ctrl = stdio_client_controller.StdioClientController(
command=sys.executable,
args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'],
env=os.environ.copy(),
capture_stderr=False,
)
task_coro = self.ctrl.run(new_connection_callback)
self._transport_task = asyncio.create_task(task_coro)
try:
await asyncio.wait_for(self._connected.wait(), timeout=_CONNECT_TIMEOUT_SEC)
except asyncio.TimeoutError as exc:
await self._stop_transport()
raise PluginRuntimeNotConnectedError('Plugin runtime did not become ready within 30 seconds') from exc
if connect_errors:
await self._stop_transport()
raise PluginRuntimeNotConnectedError(f'Plugin runtime connection failed: {connect_errors[-1]}')
if self.heartbeat_task is None or self.heartbeat_task.done():
self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
def schedule_reconnect(self) -> None:
if self._closing or not self.is_enable_plugin:
return
if self._reconnect_task is not None and not self._reconnect_task.done():
return
self._reconnect_task = asyncio.create_task(self._reconnect_loop())
async def _reconnect_loop(self) -> None:
delay = 1.0
try:
while not self._closing:
try:
await self.initialize()
return
except Exception as exc:
self.ap.logger.warning(f'Plugin runtime reconnection failed: {exc}; retrying in {delay:.0f}s')
await asyncio.sleep(delay)
delay = min(delay * 2, _RECONNECT_MAX_DELAY_SEC)
finally:
self._reconnect_task = None
async def _stop_transport(self) -> None:
self._connected.clear()
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is not None:
with contextlib.suppress(Exception):
await runtime_handler.close()
if getattr(self, 'handler', None) is runtime_handler:
del self.handler
tasks = [
task
for task in (
getattr(self, 'handler_task', None),
self._transport_task,
)
if task is not None and task is not asyncio.current_task()
]
@@ -812,73 +942,20 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
with contextlib.suppress(Exception):
await close_ctrl()
async def make_connection_failed_callback(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception = None,
) -> None:
if exc is not None:
self.ap.logger.error(f'Failed to connect to plugin runtime({ws_url}): {exc}')
else:
self.ap.logger.error(f'Failed to connect to plugin runtime({ws_url}), trying to reconnect...')
await self.runtime_disconnect_callback(self)
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=make_connection_failed_callback,
additional_headers=control_headers,
)
task = self.ctrl.run(new_connection_callback)
elif platform.get_platform() == 'win32':
# Due to Windows's lack of supports for both stdio and subprocess:
# See also: https://docs.python.org/zh-cn/3.13/library/asyncio-platforms.html
# We have to launch runtime via cmd but communicate via ws.
self.ap.logger.info('(windows) use cmd to launch plugin runtime and communicate via ws')
control_headers = self._control_headers(allow_generate=True)
await self._start_runtime_subprocess(
'-m',
'langbot_plugin.cli.__init__',
'rt',
env_overrides={PLUGIN_RUNTIME_CONTROL_TOKEN_ENV: self._control_token},
)
ws_url = 'ws://localhost:5400/control/ws'
async def make_connection_failed_callback(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception = None,
) -> None:
if exc is not None:
self.ap.logger.error(f'(windows) Failed to connect to plugin runtime({ws_url}): {exc}')
else:
self.ap.logger.error(
f'(windows) Failed to connect to plugin runtime({ws_url}), trying to reconnect...'
)
await self.runtime_disconnect_callback(self)
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=make_connection_failed_callback,
additional_headers=control_headers,
)
task = self.ctrl.run(new_connection_callback)
else: # stdio
self.ap.logger.info('use stdio to connect to plugin runtime')
# cmd: lbp rt -s
python_path = sys.executable
env = os.environ.copy()
self.ctrl = stdio_client_controller.StdioClientController(
command=python_path,
args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'],
env=env,
)
task = self.ctrl.run(new_connection_callback)
if self.heartbeat_task is None:
self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
asyncio.create_task(task)
async def aclose(self) -> None:
self._closing = True
self._generation += 1
reconnect_task = self._reconnect_task
self._reconnect_task = None
if reconnect_task is not None and reconnect_task is not asyncio.current_task():
reconnect_task.cancel()
await asyncio.gather(reconnect_task, return_exceptions=True)
if self.heartbeat_task is not None:
self.heartbeat_task.cancel()
await asyncio.gather(self.heartbeat_task, return_exceptions=True)
self.heartbeat_task = None
await self._stop_transport()
await self._close_managed_subprocess()
async def initialize_plugins(self):
pass
@@ -1001,12 +1078,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
component_kind: typing.Literal['tool', 'command'],
include_plugins: list[str] | None,
) -> InstallationBinding:
runtime_handler = self._runtime_handler()
for binding in await self._operation_bindings(include_plugins=include_plugins):
with self.handler.installation_scope(binding):
with runtime_handler.installation_scope(binding):
components = (
await self.handler.list_tools(include_plugins=include_plugins)
await runtime_handler.list_tools(include_plugins=include_plugins)
if component_kind == 'tool'
else await self.handler.list_commands(include_plugins=include_plugins)
else await runtime_handler.list_commands(include_plugins=include_plugins)
)
for component in components:
manifest = ComponentManifest.model_validate(component)
@@ -1467,6 +1545,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
install_info: dict[str, Any],
task_context: taskmgr.TaskContext | None = None,
) -> None:
runtime_handler = self._runtime_handler()
execution_context = await self._current_execution_context()
plugin_author = str(install_info.get('plugin_author') or '')
plugin_name = str(install_info.get('plugin_name') or '')
@@ -1521,7 +1600,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
except Exception:
await self._delete_artifact_if_unreferenced(execution_context, artifact_digest)
raise
self.handler.register_installation_binding(
runtime_handler.register_installation_binding(
binding,
plugin_author=plugin_author,
plugin_name=plugin_name,
@@ -1538,8 +1617,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if previous_digest is not None and not previous_was_durable and self.runtime_profile == 'oss_dev':
bridge = self._legacy_oss_bridge_binding(execution_context)
try:
with self.handler.installation_scope(bridge):
async for _ in self.handler.delete_plugin(plugin_author, plugin_name):
with runtime_handler.installation_scope(bridge):
async for _ in runtime_handler.delete_plugin(plugin_author, plugin_name):
pass
except Exception as exc:
self.ap.logger.debug(f'Legacy OSS plugin cleanup skipped: {exc}')
@@ -1570,6 +1649,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
delete_data: bool = False,
task_context: taskmgr.TaskContext | None = None,
) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
execution_context, setting = await self._setting_for_plugin(plugin_author, plugin_name)
binding = self._binding_from_setting(execution_context, setting)
is_legacy_oss = self.runtime_profile == 'oss_dev' and (
@@ -1578,11 +1658,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
)
if is_legacy_oss:
bridge = self._legacy_oss_bridge_binding(execution_context)
with self.handler.installation_scope(bridge):
async for _ in self.handler.delete_plugin(plugin_author, plugin_name):
with runtime_handler.installation_scope(bridge):
async for _ in runtime_handler.delete_plugin(plugin_author, plugin_name):
pass
await self.handler.remove_plugin_installation(binding)
self.handler.unregister_installation_binding(binding)
await runtime_handler.remove_plugin_installation(binding)
runtime_handler.unregister_installation_binding(binding)
async def delete(execute):
await execute(
@@ -1629,11 +1709,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
plugins: list[dict[str, Any]] = []
seen_plugin_ids: set[str] = set()
for binding in await self._operation_bindings():
with self.handler.installation_scope(binding):
scoped_plugins = await self.handler.list_plugins()
with runtime_handler.installation_scope(binding):
scoped_plugins = await runtime_handler.list_plugins()
for plugin in scoped_plugins:
metadata = plugin.get('manifest', {}).get('manifest', {}).get('metadata', {})
plugin_id = f'{metadata.get("author", "")}/{metadata.get("name", "")}'
@@ -1685,11 +1766,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
return plugins
async def get_plugin_info(self, author: str, plugin_name: str) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_info(author, plugin_name)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_info(author, plugin_name)
async def set_plugin_config(self, plugin_author: str, plugin_name: str, config: dict[str, Any]) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
execution_context, setting = await self._setting_for_plugin(plugin_author, plugin_name)
next_revision = setting.runtime_revision + 1
statement = (
@@ -1716,7 +1799,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
runtime_revision=next_revision,
artifact_digest=setting.artifact_digest,
)
self.handler.register_installation_binding(
runtime_handler.register_installation_binding(
binding,
plugin_author=plugin_author,
plugin_name=plugin_name,
@@ -1731,22 +1814,24 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
)
if is_legacy_oss:
bridge = self._legacy_oss_bridge_binding(execution_context)
with self.handler.installation_scope(bridge):
await self.handler.set_plugin_config(plugin_author, plugin_name, config)
with runtime_handler.installation_scope(bridge):
await runtime_handler.set_plugin_config(plugin_author, plugin_name, config)
else:
await self._apply_desired_state(desired)
self._known_desired_states[binding.installation_uuid] = desired
return {}
async def get_plugin_icon(self, plugin_author: str, plugin_name: str) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_icon(plugin_author, plugin_name)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_icon(plugin_author, plugin_name)
async def get_plugin_readme(self, plugin_author: str, plugin_name: str, language: str = 'en') -> str:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_readme(plugin_author, plugin_name, language)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_readme(plugin_author, plugin_name, language)
async def get_plugin_logs(
self,
@@ -1755,14 +1840,16 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
limit: int = 200,
level: str | None = None,
) -> list[dict[str, Any]]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_logs(plugin_author, plugin_name, limit, level)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_logs(plugin_author, plugin_name, limit, level)
async def get_plugin_assets(self, plugin_author: str, plugin_name: str, filepath: str) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_assets(plugin_author, plugin_name, filepath)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_assets(plugin_author, plugin_name, filepath)
async def handle_page_api(
self,
@@ -1773,9 +1860,10 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
method: str,
body: Any = None,
) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body)
with runtime_handler.installation_scope(binding):
return await runtime_handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body)
async def get_debug_info(self) -> dict[str, Any]:
"""Get debug information including debug key and WS URL"""
@@ -1800,11 +1888,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
event_ctx._response_sources = []
return event_ctx
runtime_handler = self._runtime_handler()
emitted_plugins: list[Any] = []
response_sources: list[dict[str, Any]] = []
for binding in await self._operation_bindings(include_plugins=bound_plugins):
with self.handler.installation_scope(binding):
result = await self.handler.emit_event(
with runtime_handler.installation_scope(binding):
result = await runtime_handler.emit_event(
event_ctx.model_dump(serialize_as_any=False),
include_plugins=bound_plugins,
)
@@ -1821,18 +1910,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return
try:
runtime_handler = self._runtime_handler()
plugin_ref = diagnostic.get('plugin') if isinstance(diagnostic, dict) else None
if isinstance(plugin_ref, dict):
author = plugin_ref.get('author') or plugin_ref.get('plugin_author')
name = plugin_ref.get('name') or plugin_ref.get('plugin_name')
if author and name:
binding = await self._target_binding(str(author), str(name))
with self.handler.installation_scope(binding):
await self.handler.notify_plugin_diagnostic(diagnostic)
with runtime_handler.installation_scope(binding):
await runtime_handler.notify_plugin_diagnostic(diagnostic)
return
for binding in await self._operation_bindings():
with self.handler.installation_scope(binding):
await self.handler.notify_plugin_diagnostic(diagnostic)
with runtime_handler.installation_scope(binding):
await runtime_handler.notify_plugin_diagnostic(diagnostic)
except Exception as e:
self.ap.logger.debug(f'Plugin diagnostic forwarding skipped: {e}')
@@ -1840,11 +1930,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
tools: list[ComponentManifest] = []
seen: set[tuple[str, str]] = set()
for binding in await self._operation_bindings(include_plugins=bound_plugins):
with self.handler.installation_scope(binding):
scoped = await self.handler.list_tools(include_plugins=bound_plugins)
with runtime_handler.installation_scope(binding):
scoped = await runtime_handler.list_tools(include_plugins=bound_plugins)
for raw_tool in scoped:
tool = ComponentManifest.model_validate(raw_tool)
key = (str(tool.owner), tool.metadata.name)
@@ -1878,8 +1969,9 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
component_kind='tool',
include_plugins=bound_plugins,
)
with self.handler.installation_scope(binding):
return await self.handler.call_tool(
runtime_handler = self._runtime_handler()
with runtime_handler.installation_scope(binding):
return await runtime_handler.call_tool(
tool_name,
parameters,
session.model_dump(serialize_as_any=True),
@@ -1892,11 +1984,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
commands: list[ComponentManifest] = []
seen: set[tuple[str, str]] = set()
for binding in await self._operation_bindings(include_plugins=bound_plugins):
with self.handler.installation_scope(binding):
scoped = await self.handler.list_commands(include_plugins=bound_plugins)
with runtime_handler.installation_scope(binding):
scoped = await runtime_handler.list_commands(include_plugins=bound_plugins)
for raw_command in scoped:
command = ComponentManifest.model_validate(raw_command)
key = (str(command.owner), command.metadata.name)
@@ -1925,8 +2018,9 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
component_kind='command',
include_plugins=bound_plugins,
)
with self.handler.installation_scope(binding):
gen = self.handler.execute_command(
runtime_handler = self._runtime_handler()
with runtime_handler.installation_scope(binding):
gen = runtime_handler.execute_command(
command_ctx.model_dump(serialize_as_any=True),
include_plugins=bound_plugins,
)
@@ -1944,9 +2038,15 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return {'results': []}
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.retrieve_knowledge(plugin_author, plugin_name, retriever_name, retrieval_context)
with runtime_handler.installation_scope(binding):
return await runtime_handler.retrieve_knowledge(
plugin_author,
plugin_name,
retriever_name,
retrieval_context,
)
def dispose(self):
"""Best-effort synchronous compatibility wrapper; prefer ``aclose``."""
@@ -1997,41 +2097,47 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
context_data: IngestionContext data.
"""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.rag_ingest_document(plugin_author, plugin_name, context_data)
with runtime_handler.installation_scope(binding):
return await runtime_handler.rag_ingest_document(plugin_author, plugin_name, context_data)
async def call_rag_delete_document(self, plugin_id: str, document_id: str, kb_id: str) -> bool:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.rag_delete_document(plugin_author, plugin_name, document_id, kb_id)
with runtime_handler.installation_scope(binding):
return await runtime_handler.rag_delete_document(plugin_author, plugin_name, document_id, kb_id)
async def get_rag_creation_schema(self, plugin_id: str) -> dict[str, Any]:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_rag_creation_schema(plugin_author, plugin_name)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_rag_creation_schema(plugin_author, plugin_name)
async def get_rag_retrieval_schema(self, plugin_id: str) -> dict[str, Any]:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_rag_retrieval_schema(plugin_author, plugin_name)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_rag_retrieval_schema(plugin_author, plugin_name)
async def rag_on_kb_create(self, plugin_id: str, kb_id: str, config: dict[str, Any]) -> dict[str, Any]:
"""Notify plugin about KB creation."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.rag_on_kb_create(plugin_author, plugin_name, kb_id, config)
with runtime_handler.installation_scope(binding):
return await runtime_handler.rag_on_kb_create(plugin_author, plugin_name, kb_id, config)
async def rag_on_kb_delete(self, plugin_id: str, kb_id: str) -> dict[str, Any]:
"""Notify plugin about KB deletion."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.rag_on_kb_delete(plugin_author, plugin_name, kb_id)
with runtime_handler.installation_scope(binding):
return await runtime_handler.rag_on_kb_delete(plugin_author, plugin_name, kb_id)
async def call_rag_retrieve(self, plugin_id: str, retrieval_context: dict[str, Any]) -> dict[str, Any]:
"""Call plugin to retrieve knowledge.
@@ -2041,9 +2147,10 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
retrieval_context: RetrievalContext data.
"""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context)
with runtime_handler.installation_scope(binding):
return await runtime_handler.retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context)
async def list_knowledge_engines(self) -> list[dict[str, Any]]:
"""List all available Knowledge Engines from plugins.
@@ -2053,11 +2160,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
engines: list[dict[str, Any]] = []
seen: set[tuple[str, str]] = set()
for binding in await self._operation_bindings():
with self.handler.installation_scope(binding):
scoped = await self.handler.list_knowledge_engines()
with runtime_handler.installation_scope(binding):
scoped = await runtime_handler.list_knowledge_engines()
for engine in scoped:
key = (str(engine.get('plugin_id', '')), str(engine.get('name', '')))
if key not in seen:
@@ -2069,11 +2177,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
"""List all available parsers from plugins."""
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
parsers: list[dict[str, Any]] = []
seen: set[tuple[str, str]] = set()
for binding in await self._operation_bindings():
with self.handler.installation_scope(binding):
scoped = await self.handler.list_parsers()
with runtime_handler.installation_scope(binding):
scoped = await runtime_handler.list_parsers()
for parser in scoped:
key = (str(parser.get('plugin_id', '')), str(parser.get('name', '')))
if key not in seen:
@@ -2084,6 +2193,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
async def call_parser(self, plugin_id: str, context_data: dict[str, Any], file_bytes: bytes) -> dict[str, Any]:
"""Call plugin to parse a document."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.parse_document(plugin_author, plugin_name, context_data, file_bytes)
with runtime_handler.installation_scope(binding):
return await runtime_handler.parse_document(plugin_author, plugin_name, context_data, file_bytes)
@@ -224,12 +224,13 @@ class LocalAgentRunner(runner.RequestRunner):
) -> list[modelmgr_requester.RuntimeLLMModel]:
"""Build ordered list of models to try: primary model + fallback models."""
candidates = []
execution_context = get_query_execution_context(query)
# Primary model
if query.use_llm_model_uuid:
try:
primary = await self.ap.model_mgr.get_model_by_uuid(
get_query_execution_context(query),
execution_context,
query.use_llm_model_uuid,
)
candidates.append(primary)
@@ -241,7 +242,7 @@ class LocalAgentRunner(runner.RequestRunner):
for fb_uuid in fallback_uuids:
try:
fb_model = await self.ap.model_mgr.get_model_by_uuid(
get_query_execution_context(query),
execution_context,
fb_uuid,
)
candidates.append(fb_model)
+49 -4
View File
@@ -33,7 +33,13 @@ from ....workspace.errors import WorkspaceError, WorkspaceInvariantError
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from ....entity.persistence import mcp as persistence_mcp
from .mcp_stdio import BoxStdioSessionRuntime, MCPServerBoxConfig, MCPSessionErrorPhase, _ColdStartRetry # noqa: F401
from .mcp_stdio import (
BoxStdioSessionRuntime,
MCPServerBoxConfig as MCPServerBoxConfig, # noqa: F401 - public re-export
MCPSessionErrorPhase,
_ColdStartRetry,
_get_default_memory_mb,
)
from .mcp_policy import require_stdio_mcp_enabled, stdio_mcp_enabled
# Synthesized LLM tools for MCP resources (not from server tools/list).
@@ -321,6 +327,26 @@ class RuntimeMCPSession:
self._box_stdio_runtime = BoxStdioSessionRuntime(self)
self.box_config = self._box_stdio_runtime.config
def _parse_tool_call_timeout(self, value: typing.Any) -> float:
"""Return a safe tool-call timeout; zero explicitly disables it."""
try:
timeout = -1 if isinstance(value, bool) else float(value)
if timeout > 0:
# Validate the exact conversion used for each call here, so a
# finite-but-enormous manual config cannot fail at invocation.
timedelta(seconds=timeout)
except (TypeError, ValueError, OverflowError):
timeout = -1
if not math.isfinite(timeout) or timeout < 0:
self.ap.logger.warning(
f'Invalid MCP tool call timeout {value!r} for {self.server_name}; '
f'using {MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS:g} seconds'
)
return MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS
return timeout
async def _assert_execution_active(self) -> None:
"""Fail closed when this long-lived session belongs to a stale placement."""
@@ -723,7 +749,11 @@ class RuntimeMCPSession:
self.status = MCPSessionStatus.CONNECTING
self.error_message = None
self.error_phase = None
await asyncio.sleep(1)
try:
await self._sleep_with_execution_fence(1)
except WorkspaceError as fence_error:
self._stop_for_stale_execution(fence_error)
return
continue
# Explicitly disabled Box is a deliberate refusal, not a
# transient failure. Surface it immediately without log
@@ -1082,7 +1112,12 @@ class RuntimeMCPSession:
try:
await self._assert_execution_active()
result = await self.session.call_tool(tool_name, arguments)
read_timeout = timedelta(seconds=self.tool_call_timeout_sec) if self.tool_call_timeout_sec > 0 else None
result = await self.session.call_tool(
tool_name,
arguments,
read_timeout_seconds=read_timeout,
)
await self._assert_execution_active()
except Exception as e:
if self._is_tool_call_timeout(e):
@@ -2144,7 +2179,15 @@ class MCPLoader(loader.ToolLoader):
async def shutdown(self):
"""关闭所有工具"""
self.ap.logger.info('Shutting down all MCP sessions...')
for key, session in list(self.sessions.items()):
hosted_tasks = [task for task in self._hosted_mcp_tasks if not task.done()]
for task in hosted_tasks:
task.cancel()
if hosted_tasks:
await asyncio.gather(*hosted_tasks, return_exceptions=True)
self._hosted_mcp_tasks.clear()
async def shutdown_session(session: RuntimeMCPSession) -> None:
try:
await session.shutdown()
self.ap.logger.debug(f'Shutdown MCP session: {session.server_name}')
@@ -2152,5 +2195,7 @@ class MCPLoader(loader.ToolLoader):
self.ap.logger.error(
f'Error shutting down MCP session {session.server_name}: {e}\n{traceback.format_exc()}'
)
await asyncio.gather(*(shutdown_session(session) for session in list(self.sessions.values())))
self.sessions.clear()
self.ap.logger.info('All MCP sessions shutdown complete')
+16 -7
View File
@@ -25,7 +25,7 @@ from ..entity.persistence.workspace import (
WorkspaceStatus,
)
from .entities import WorkspaceExecutionBinding
from .errors import WorkspaceNotFoundError
from .errors import WorkspaceExecutionUnavailableError, WorkspaceInvariantError, WorkspaceNotFoundError
from .policy import CloudWorkspacePolicy, SingleWorkspacePolicy
from .service import WorkspaceService
@@ -231,13 +231,22 @@ class WorkspaceCollaborationService:
accesses: list[ResolvedWorkspaceAccess] = []
for workspace_uuid in workspace_uuids:
async with tenant_uow(workspace_uuid) as workspace_uow:
accesses.append(
await self.resolve_account_workspace(
account_uuid,
workspace_uuid,
session=workspace_uow.session,
try:
accesses.append(
await self.resolve_account_workspace(
account_uuid,
workspace_uuid,
session=workspace_uow.session,
)
)
except (
WorkspaceExecutionUnavailableError,
WorkspaceInvariantError,
WorkspaceNotFoundError,
) as exc:
self.ap.logger.warning(
f'Skipping inactive Workspace discovery projection {workspace_uuid!r}: {exc}'
)
)
accesses.sort(key=lambda access: (access.workspace.created_at, access.workspace.uuid))
return accesses
+14
View File
@@ -177,6 +177,7 @@ class WorkspaceService:
"""
self._require_deployment_admission()
self._require_directory_projection()
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if session is None and workspace_uuid is not None and callable(tenant_uow):
@@ -247,6 +248,7 @@ class WorkspaceService:
binding = await self._run(operation, session=session)
# The database lookup can cross the Manifest expiry boundary. Never
# return a binding that is already invalid at a side-effect boundary.
self._require_directory_projection()
self._require_deployment_admission()
return binding
@@ -500,3 +502,15 @@ class WorkspaceService:
guard = getattr(self.ap, 'deployment_admission', None)
if guard is not None:
guard.require_active()
def _require_directory_projection(self) -> None:
deployment = getattr(self.ap, 'deployment', None)
if deployment is None or not getattr(deployment, 'multi_workspace_enabled', False):
return
projection = getattr(self.ap, 'directory_projection_service', None)
if projection is None:
raise WorkspaceExecutionUnavailableError('Cloud directory projection is unavailable')
try:
projection.require_ready()
except RuntimeError as exc:
raise WorkspaceExecutionUnavailableError(str(exc)) from exc