mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
fix(cloud): show owner model balance and enforce single owner (#2384)
* fix(cloud): show owner model balance and enforce single owner * fix(migrations): create owner index idempotently --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
@@ -19,7 +19,6 @@ class Permission(enum.StrEnum):
|
||||
WORKSPACE_VIEW = 'workspace.view'
|
||||
WORKSPACE_UPDATE = 'workspace.update'
|
||||
WORKSPACE_DELETE = 'workspace.delete'
|
||||
OWNER_TRANSFER = 'owner.transfer'
|
||||
MEMBER_VIEW = 'member.view'
|
||||
MEMBER_INVITE = 'member.invite'
|
||||
MEMBER_UPDATE_ROLE = 'member.update_role'
|
||||
@@ -49,7 +48,6 @@ _ROLE_PERMISSIONS: typing.Final = types.MappingProxyType(
|
||||
if permission
|
||||
not in {
|
||||
Permission.WORKSPACE_DELETE,
|
||||
Permission.OWNER_TRANSFER,
|
||||
Permission.BILLING_LINK_MANAGE,
|
||||
}
|
||||
),
|
||||
|
||||
@@ -62,7 +62,6 @@ class AuthType(enum.Enum):
|
||||
|
||||
_SUPPORT_ADMIN_DENIED_PERMISSIONS = frozenset(
|
||||
{
|
||||
Permission.OWNER_TRANSFER.value,
|
||||
Permission.MEMBER_VIEW.value,
|
||||
Permission.MEMBER_INVITE.value,
|
||||
Permission.MEMBER_UPDATE_ROLE.value,
|
||||
|
||||
@@ -291,11 +291,19 @@ class UserRouterGroup(group.RouterGroup):
|
||||
# Workspace owner is already bound even when this Core has no local OAuth
|
||||
# token row (model billing uses the owner's control-plane API key).
|
||||
owner_space_bound = cloud_mode or owner_has_local_space_credentials
|
||||
credits = (
|
||||
await self.ap.space_service.get_credits(owner.user)
|
||||
if owner is not None and owner.space_account_uuid
|
||||
else None
|
||||
)
|
||||
if cloud_mode:
|
||||
catalog_service = getattr(self.ap, 'cloud_model_catalog_service', None)
|
||||
credits = (
|
||||
catalog_service.get_workspace_credits(access.workspace.uuid)
|
||||
if catalog_service is not None
|
||||
else None
|
||||
)
|
||||
else:
|
||||
credits = (
|
||||
await self.ap.space_service.get_credits(owner.user)
|
||||
if owner is not None and owner.space_account_uuid
|
||||
else None
|
||||
)
|
||||
return self.success(
|
||||
data={
|
||||
'credits': credits,
|
||||
|
||||
@@ -53,6 +53,7 @@ class CloudWorkspaceModelBilling(BaseModel):
|
||||
workspace_uuid: str = Field(min_length=36, max_length=36)
|
||||
owner_account_uuid: str | None = Field(default=None, min_length=36, max_length=36)
|
||||
api_key: SecretStr | None = None
|
||||
credits: int | None = None
|
||||
|
||||
@field_validator('workspace_uuid')
|
||||
@classmethod
|
||||
@@ -149,6 +150,11 @@ class CloudModelCatalogSyncService:
|
||||
# convergence marker so a failed runtime reload is retried even when the
|
||||
# following database reconciliation is a no-op.
|
||||
self._runtime_reload_pending = False
|
||||
self._workspace_credits: dict[str, int | None] = {}
|
||||
|
||||
def get_workspace_credits(self, workspace_uuid: str) -> int | None:
|
||||
"""Return the latest signed owner-credit projection for a Workspace."""
|
||||
return self._workspace_credits.get(str(uuid.UUID(workspace_uuid)))
|
||||
|
||||
async def initialize(self) -> None:
|
||||
await self.sync_once(reload_runtime=False)
|
||||
@@ -198,6 +204,7 @@ class CloudModelCatalogSyncService:
|
||||
self._runtime_reload_pending = True
|
||||
for key in ('created', 'updated', 'deleted'):
|
||||
summary[key] += counts[key]
|
||||
self._workspace_credits[binding.workspace_uuid] = billing_by_workspace[binding.workspace_uuid].credits
|
||||
except Exception as exc:
|
||||
sync_error = exc
|
||||
finally:
|
||||
|
||||
@@ -163,6 +163,13 @@ class WorkspaceMembership(Base):
|
||||
__table_args__ = (
|
||||
sqlalchemy.UniqueConstraint('workspace_uuid', 'account_uuid', name='uq_workspace_membership_account'),
|
||||
sqlalchemy.Index('ix_workspace_memberships_account_status', 'account_uuid', 'status'),
|
||||
sqlalchemy.Index(
|
||||
'uq_workspace_memberships_one_active_owner',
|
||||
'workspace_uuid',
|
||||
unique=True,
|
||||
sqlite_where=sqlalchemy.text("role = 'owner' AND status = 'active'"),
|
||||
postgresql_where=sqlalchemy.text("role = 'owner' AND status = 'active'"),
|
||||
),
|
||||
sqlalchemy.CheckConstraint(
|
||||
"role IN ('owner', 'admin', 'developer', 'operator', 'viewer')",
|
||||
name='ck_workspace_memberships_role',
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
"""enforce one active owner per Workspace
|
||||
|
||||
Revision ID: 0019_single_workspace_owner
|
||||
Revises: 0018_merge_launch_replay
|
||||
Create Date: 2026-08-02
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = '0019_single_workspace_owner'
|
||||
down_revision = '0018_merge_launch_replay'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_INDEX_NAME = 'uq_workspace_memberships_one_active_owner'
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
if 'workspace_memberships' not in inspector.get_table_names():
|
||||
return
|
||||
|
||||
# Ownership transfer used to promote a second member without demoting the
|
||||
# original owner. Preserve the Workspace creator where possible and demote
|
||||
# every historical extra owner before installing the database invariant.
|
||||
op.execute(
|
||||
sa.text(
|
||||
"""
|
||||
WITH ranked_owners AS (
|
||||
SELECT membership.uuid,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY membership.workspace_uuid
|
||||
ORDER BY
|
||||
CASE
|
||||
WHEN membership.account_uuid = workspace.created_by_account_uuid THEN 0
|
||||
ELSE 1
|
||||
END,
|
||||
COALESCE(membership.joined_at, membership.created_at),
|
||||
membership.uuid
|
||||
) AS owner_rank
|
||||
FROM workspace_memberships AS membership
|
||||
JOIN workspaces AS workspace
|
||||
ON workspace.uuid = membership.workspace_uuid
|
||||
WHERE membership.role = 'owner'
|
||||
AND membership.status = 'active'
|
||||
)
|
||||
UPDATE workspace_memberships
|
||||
SET role = 'admin'
|
||||
WHERE uuid IN (
|
||||
SELECT uuid
|
||||
FROM ranked_owners
|
||||
WHERE owner_rank > 1
|
||||
)
|
||||
"""
|
||||
)
|
||||
)
|
||||
# Fresh installations may already have this index because SQLAlchemy
|
||||
# metadata is created before Alembic advances the revision marker.
|
||||
op.execute(
|
||||
sa.text(
|
||||
'CREATE UNIQUE INDEX IF NOT EXISTS '
|
||||
'uq_workspace_memberships_one_active_owner '
|
||||
'ON workspace_memberships (workspace_uuid) '
|
||||
"WHERE role = 'owner' AND status = 'active'"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
if 'workspace_memberships' not in inspector.get_table_names():
|
||||
return
|
||||
index_names = {index['name'] for index in inspector.get_indexes('workspace_memberships')}
|
||||
if _INDEX_NAME in index_names:
|
||||
op.drop_index(_INDEX_NAME, table_name='workspace_memberships')
|
||||
@@ -606,6 +606,8 @@ class WorkspaceCollaborationService:
|
||||
) -> WorkspaceMembership:
|
||||
if role not in {item.value for item in MembershipRole}:
|
||||
raise MembershipPermissionError('Unknown Workspace role')
|
||||
if role == MembershipRole.OWNER.value:
|
||||
raise MembershipPermissionError('Workspace ownership cannot be transferred')
|
||||
|
||||
async def operation(active_session: AsyncSession) -> WorkspaceMembership:
|
||||
await self._require_active_workspace(active_session, workspace_uuid)
|
||||
@@ -617,8 +619,8 @@ class WorkspaceCollaborationService:
|
||||
target_account_uuid,
|
||||
)
|
||||
self._require_can_manage_target(persisted_actor, target, new_role=role)
|
||||
if target.role == MembershipRole.OWNER.value and role != MembershipRole.OWNER.value:
|
||||
await self._require_another_owner(active_session, workspace_uuid, target.account_uuid)
|
||||
if target.role == MembershipRole.OWNER.value:
|
||||
raise LastOwnerError('The Workspace owner cannot be removed or demoted')
|
||||
target.role = role
|
||||
await active_session.flush()
|
||||
return target
|
||||
@@ -644,7 +646,7 @@ class WorkspaceCollaborationService:
|
||||
)
|
||||
self._require_can_manage_target(persisted_actor, target)
|
||||
if target.role == MembershipRole.OWNER.value:
|
||||
await self._require_another_owner(active_session, workspace_uuid, target.account_uuid)
|
||||
raise LastOwnerError('The Workspace owner cannot be removed or demoted')
|
||||
target.status = MembershipStatus.REMOVED.value
|
||||
await active_session.flush()
|
||||
return target
|
||||
@@ -751,26 +753,6 @@ class WorkspaceCollaborationService:
|
||||
raise WorkspaceNotFoundError('Workspace not found')
|
||||
return persisted_actor
|
||||
|
||||
async def _require_another_owner(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
workspace_uuid: str,
|
||||
excluded_account_uuid: str,
|
||||
) -> None:
|
||||
owners = (
|
||||
await session.scalars(
|
||||
sqlalchemy.select(WorkspaceMembership)
|
||||
.where(
|
||||
WorkspaceMembership.workspace_uuid == workspace_uuid,
|
||||
WorkspaceMembership.status == MembershipStatus.ACTIVE.value,
|
||||
WorkspaceMembership.role == MembershipRole.OWNER.value,
|
||||
)
|
||||
.with_for_update()
|
||||
)
|
||||
).all()
|
||||
if not any(owner.account_uuid != excluded_account_uuid for owner in owners):
|
||||
raise LastOwnerError('The last Workspace owner cannot be removed or demoted')
|
||||
|
||||
def _require_actor_workspace(self, actor: WorkspaceMembership, workspace_uuid: str) -> None:
|
||||
if actor.workspace_uuid != workspace_uuid or actor.status != MembershipStatus.ACTIVE.value:
|
||||
raise WorkspaceNotFoundError('Workspace not found')
|
||||
|
||||
Reference in New Issue
Block a user