"""Issue one namespace Capability over an exact, already-authorized Space batch.""" from __future__ import annotations import hashlib from datetime import UTC, datetime from typing import NamedTuple, Protocol import sqlalchemy as sa from sqlalchemy.orm import Session, sessionmaker from models.knowledge_fs import ( KnowledgeFSAuthorizationRevision, KnowledgeFSControlSpace, KnowledgeFSControlSpaceState, ) from services.knowledge_fs.cutover_runtime_gate import KnowledgeFSWorkspaceRuntimeGatePort from services.knowledge_fs.product_remote import KnowledgeFSOperationUnavailableError from services.knowledge_fs_capability import ( CapabilityAuthzRevision, CapabilityIssueRequest, CapabilityResource, IssuedKnowledgeFSCapability, ) MAX_BATCH_SPACE_SUMMARIES = 100 class KnowledgeFSBatchSpaceBinding(NamedTuple): control_space_id: str knowledge_space_id: str class KnowledgeFSIssuedBatchCapability(NamedTuple): token: str expires_at: datetime knowledge_space_ids: tuple[str, ...] trace_id: str class KnowledgeFSBatchCapabilityIssuerPort(Protocol): def issue_interactive( self, *, tenant_id: str, account_id: str, bindings: tuple[KnowledgeFSBatchSpaceBinding, ...], trace_id: str, ) -> KnowledgeFSIssuedBatchCapability: ... class KnowledgeFSCapabilityIssuerPort(Protocol): def issue(self, request: CapabilityIssueRequest) -> IssuedKnowledgeFSCapability: ... class KnowledgeFSBatchCapabilityBroker: """Revalidates exact tenant/registration bindings before signing the scoped ID set.""" def __init__( self, session_maker: sessionmaker[Session], *, cutover_gate: KnowledgeFSWorkspaceRuntimeGatePort, issuer: KnowledgeFSCapabilityIssuerPort | None, ) -> None: self._session_maker = session_maker self._cutover_gate = cutover_gate self._issuer = issuer def issue_interactive( self, *, tenant_id: str, account_id: str, bindings: tuple[KnowledgeFSBatchSpaceBinding, ...], trace_id: str, ) -> KnowledgeFSIssuedBatchCapability: self._cutover_gate.require_capability_v2(tenant_id=tenant_id) if self._issuer is None: raise KnowledgeFSOperationUnavailableError("KnowledgeFS Capability v2 is disabled") _validate_bindings(bindings) control_space_ids = tuple(binding.control_space_id for binding in bindings) with self._session_maker() as session: rows = session.execute( sa.select(KnowledgeFSControlSpace, KnowledgeFSAuthorizationRevision) .join( KnowledgeFSAuthorizationRevision, sa.and_( KnowledgeFSAuthorizationRevision.tenant_id == KnowledgeFSControlSpace.tenant_id, KnowledgeFSAuthorizationRevision.control_space_id == KnowledgeFSControlSpace.id, ), ) .where( KnowledgeFSControlSpace.tenant_id == tenant_id, KnowledgeFSControlSpace.id.in_(control_space_ids), KnowledgeFSControlSpace.state == KnowledgeFSControlSpaceState.ACTIVE, ) ).all() rows_by_control_id = {space.id: (space, revision) for space, revision in rows} if len(rows_by_control_id) != len(bindings): raise KnowledgeFSOperationUnavailableError("KnowledgeFS batch authorization snapshot is incomplete") ordered_revisions: list[KnowledgeFSAuthorizationRevision] = [] for binding in bindings: space, revision = rows_by_control_id[binding.control_space_id] if space.knowledge_space_id != binding.knowledge_space_id: raise KnowledgeFSOperationUnavailableError("KnowledgeFS batch registration changed") ordered_revisions.append(revision) knowledge_space_ids = tuple(binding.knowledge_space_id for binding in bindings) digest = hashlib.sha256("\0".join((*control_space_ids, *knowledge_space_ids)).encode()).hexdigest()[:24] issued = self._issuer.issue( CapabilityIssueRequest( actor=f"dify-account:{account_id}", authz_revision=CapabilityAuthzRevision( membership_epoch=max(revision.membership_epoch for revision in ordered_revisions), space_acl_epoch=max(revision.space_acl_epoch for revision in ordered_revisions), external_access_epoch=max(revision.external_access_epoch for revision in ordered_revisions), credential_revision=None, ), caller_kind="interactive", content_policy_revision=max(revision.content_policy_revision for revision in ordered_revisions), content_scope_ids=knowledge_space_ids, control_space_id=bindings[0].control_space_id, grant_id=f"account:{account_id}:batch:{digest}", namespace_id=tenant_id, operation_id="batchKnowledgeSpaceProductSummaries", principal_id=account_id, resource=CapabilityResource(type="namespace", id=tenant_id), trace_id=trace_id, ) ) return KnowledgeFSIssuedBatchCapability( token=issued.token, expires_at=datetime.fromtimestamp(issued.claims.exp, tz=UTC), knowledge_space_ids=knowledge_space_ids, trace_id=trace_id, ) def _validate_bindings(bindings: tuple[KnowledgeFSBatchSpaceBinding, ...]) -> None: if not bindings or len(bindings) > MAX_BATCH_SPACE_SUMMARIES: raise KnowledgeFSOperationUnavailableError("KnowledgeFS batch size must be between 1 and 100") control_ids = tuple(binding.control_space_id.strip() for binding in bindings) space_ids = tuple(binding.knowledge_space_id.strip() for binding in bindings) if ( any(not identifier for identifier in (*control_ids, *space_ids)) or len(set(control_ids)) != len(control_ids) or len(set(space_ids)) != len(space_ids) ): raise KnowledgeFSOperationUnavailableError("KnowledgeFS batch bindings must be unique and non-empty") __all__ = [ "MAX_BATCH_SPACE_SUMMARIES", "KnowledgeFSBatchCapabilityBroker", "KnowledgeFSBatchCapabilityIssuerPort", "KnowledgeFSBatchSpaceBinding", "KnowledgeFSIssuedBatchCapability", ]