SunnyQjm/dify
1
0 Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-10-04 13:47:55 +08:00
Code Issues Projects Releases Wiki Activity

refactor(api): centralize API key management (#42489)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Byron.wangandautofix-ci[bot] authored and GitHub committed 2026-09-23 08:38:48 +00:00
1 parent 1a09485013
commit 8b7f903cff
36 files changed
+1840 -1077
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

api/controllers/console/agent/roster.py
+30 -38
View File
@@ -18,7 +18,7 @@ from controllers.common.schema import (
)
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.apikey import ApiKeyItem, ApiKeyList, BaseApiKeyListResource, BaseApiKeyResource
from controllers.console.apikey import API_KEY_DELETE_ROLES, API_KEY_EDIT_ROLES, api_key_errors
from controllers.console.app.app import (
APP_LIST_QUERY_ARRAY_FIELDS,
AppListQuery,
@@ -35,6 +35,7 @@ from controllers.console.app.app import (
from controllers.console.app.app import (
UpdateAppPayload as GenericUpdateAppPayload,
)
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import (
RBACPermission,
account_initialization_required,
@@ -49,6 +50,7 @@ from controllers.console.wraps import (
)
from core.agent.publish_visibility import agent_has_workflow_callable_active_snapshot
from enums import DeploymentEdition
from extensions.ext_application_services import application_services
from fields.agent_fields import (
AgentConfigDraftSummaryResponse,
AgentConfigSnapshotDetailResponse,
@@ -63,10 +65,12 @@ from fields.agent_fields import (
AgentRosterListResponse,
AgentStatisticSummaryEnvelopeResponse,
)
from fields.api_key_fields import ApiKeyItem, ApiKeyList
from fields.base import ResponseModel
from libs.datetime_utils import parse_time_range
from libs.helper import dump_response
from libs.login import login_required
from machinery.context import RequestContext
from models import Account
from models.agent import Agent, AgentStatus
from models.agent_config_entities import AgentSoulConfig
@@ -1039,56 +1043,44 @@ class AgentApiStatusApi(Resource):
@console_ns.route("/agent/<uuid:agent_id>/api-keys")
class AgentApiKeyListApi(BaseApiKeyListResource):
resource_type = ApiTokenType.APP
resource_model = App
resource_id_field = "app_id"
token_prefix = "app-"
class AgentApiKeyListApi(Resource):
@console_ns.response(200, "Agent service API keys", console_ns.models[ApiKeyList.__name__])
@rbac_permission_required(RBACCheck(RBACPermission.AGENT_ACCESS_POINT_VIEW, AgentId()))
@with_current_tenant_id
@edit_permission_required
@with_session(write=False)
def get(self, session: Session, tenant_id: str, agent_id: UUID) -> dict[str, object]:
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
return dump_response(ApiKeyList, self._get_api_key_list(str(app_model.id), tenant_id, session=session))
@console_account_admission(
allowed_roles=API_KEY_EDIT_ROLES,
rbac_checks=[RBACCheck(RBACPermission.AGENT_ACCESS_POINT_VIEW, AgentId())],
)
def get(self, request_context: RequestContext, agent_id: UUID) -> dict[str, object]:
with api_key_errors():
keys = application_services().app_api_keys.list_agent_keys(request_context, str(agent_id))
return dump_response(ApiKeyList, {"data": keys})
@console_ns.response(201, "Agent service API key created", console_ns.models[ApiKeyItem.__name__])
@console_ns.response(400, "Maximum keys exceeded")
@with_current_tenant_id
@edit_permission_required
@rbac_permission_required(RBACCheck(RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentId()))
@with_session
def post(self, session: Session, tenant_id: str, agent_id: UUID) -> tuple[dict[str, object], int]:
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
return dump_response(
ApiKeyItem,
self._create_api_key(str(app_model.id), tenant_id, session=session),
), 201
@console_account_admission(
allowed_roles=API_KEY_EDIT_ROLES,
rbac_checks=[RBACCheck(RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentId())],
)
def post(self, request_context: RequestContext, agent_id: UUID) -> tuple[dict[str, object], int]:
with api_key_errors():
key = application_services().app_api_keys.create_agent_key(request_context, str(agent_id))
return dump_response(ApiKeyItem, key), 201
@console_ns.route("/agent/<uuid:agent_id>/api-keys/<uuid:api_key_id>")
class AgentApiKeyApi(BaseApiKeyResource):
resource_type = ApiTokenType.APP
resource_model = App
resource_id_field = "app_id"
class AgentApiKeyApi(Resource):
@console_ns.response(204, "Agent service API key deleted")
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACCheck(RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentId()))
@with_session
@console_account_admission(
allowed_roles=API_KEY_DELETE_ROLES,
rbac_checks=[RBACCheck(RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentId())],
)
def delete(
self,
session: Session,
tenant_id: str,
current_user: Account,
request_context: RequestContext,
agent_id: UUID,
api_key_id: UUID,
) -> tuple[str, int]:
app_model = _resolve_agent_app_model(session, tenant_id=tenant_id, agent_id=agent_id)
self._delete_api_key(str(app_model.id), str(api_key_id), tenant_id, current_user, session=session)
with api_key_errors():
application_services().app_api_keys.delete_agent_key(request_context, str(agent_id), str(api_key_id))
return "", 204
api/controllers/console/apikey.py
+93 -332
View File
@@ -1,390 +1,151 @@
from collections.abc import Iterable, Mapping
from datetime import datetime
"""Console transport adapters for resource-scoped API key management."""
from collections.abc import Generator
from contextlib import contextmanager
from uuid import UUID
import flask_restx
from flask_restx import Resource
from flask_restx._http import HTTPStatus
from pydantic import field_validator
from sqlalchemy import func, or_, select
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden
from configs import dify_config
from controllers.common.rbac import AgentBehindApp, DatasetId, PlainApp, RBACCheck
from controllers.common.schema import register_response_schema_models
from controllers.common.session import with_session
from fields.base import ResponseModel
from libs.helper import dump_response, to_timestamp
from libs.login import login_required
from models import Account
from models.dataset import Dataset
from models.enums import ApiTokenType
from models.model import ApiToken, App
from services.api_token_service import ApiTokenCache
from services.app_service import AppService
from . import console_ns
from .wraps import (
RBACPermission,
account_initialization_required,
edit_permission_required,
rbac_permission_required,
setup_required,
with_current_tenant_id,
with_current_user,
from controllers.console import console_ns
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import RBACPermission
from extensions.ext_application_services import application_services
from fields.api_key_fields import ApiKeyItem, ApiKeyList, build_masked_api_key_list
from libs.helper import dump_response
from machinery.context import RequestContext
from models.account import TenantAccountRole
from services.agent.errors import AgentAccessNotReadyError
from services.app.api_key_service import AppApiKeyNotReadyError
from services.auth.api_key_contracts import (
ApiKeyLimitExceededError,
ApiKeyNotFoundError,
ApiKeyResourceNotFoundError,
)
from services.errors.account import NoPermissionError
from services.knowledge.api_key_service import UnknownDatasetIdsError
class ApiKeyItem(ResponseModel):
id: str
type: str
token: str
# Dataset keys only: the knowledge bases this key is bound to. Empty = the key can
# access every dataset in the tenant (default). App keys are always empty.
dataset_ids: list[str] = []
last_used_at: int | None = None
created_at: int | None = None
@field_validator("last_used_at", "created_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
return to_timestamp(value)
class ApiKeyList(ResponseModel):
data: list[ApiKeyItem]
API_KEY_EDIT_ROLES = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR})
API_KEY_DELETE_ROLES = frozenset({TenantAccountRole.OWNER, TenantAccountRole.ADMIN})
register_response_schema_models(console_ns, ApiKeyItem, ApiKeyList)
def mask_api_token(token: str) -> str:
"""Mask a secret token for list responses.
Reveal-once: the full secret is only returned by the create endpoint. List
endpoints expose just enough (prefix + last 4) to identify a key, never the
full value, so an existing key's secret cannot be retrieved after creation.
"""
if len(token) <= 8:
return "***"
return f"{token[:5]}...{token[-4:]}"
def build_masked_api_key_list(
api_tokens: Iterable[ApiToken],
bindings_by_token: Mapping[str, list[str]] | None = None,
) -> ApiKeyList:
"""Build an ApiKeyList from ORM tokens with their secrets masked.
``bindings_by_token`` maps an api_token id to the dataset ids it is bound to
(from DatasetApiTokenBinding); tokens absent from the map are unbound (empty =
access all). App-key lists omit it entirely.
"""
bindings_by_token = bindings_by_token or {}
items: list[ApiKeyItem] = []
for api_token in api_tokens:
item = ApiKeyItem.model_validate(api_token, from_attributes=True)
item.token = mask_api_token(item.token)
item.dataset_ids = bindings_by_token.get(str(api_token.id), [])
items.append(item)
return ApiKeyList(data=items)
def _get_resource(resource_id, tenant_id, resource_model, *, session: Session):
resource = session.execute(
select(resource_model).filter_by(id=resource_id, tenant_id=tenant_id)
).scalar_one_or_none()
if resource is None:
flask_restx.abort(HTTPStatus.NOT_FOUND, message=f"{resource_model.__name__} not found.")
return resource
class BaseApiKeyListResource(Resource):
method_decorators = [account_initialization_required, login_required, setup_required]
resource_type: ApiTokenType | None = None
resource_model: type | None = None
resource_id_field: str | None = None
token_prefix: str | None = None
max_keys = 10
@with_session(write=False)
def get(self, session: Session, resource_id: str, current_tenant_id: str) -> dict[str, object]:
return dump_response(
ApiKeyList,
self._get_api_key_list(resource_id, current_tenant_id, session=session),
)
def _get_api_key_list(self, resource_id: str, current_tenant_id: str, *, session: Session) -> ApiKeyList:
assert self.resource_id_field is not None, "resource_id_field must be set"
_get_resource(resource_id, current_tenant_id, self.resource_model, session=session)
keys = session.scalars(
select(ApiToken).where(
or_(ApiToken.tenant_id == current_tenant_id, ApiToken.tenant_id.is_(None)),
ApiToken.type == self.resource_type,
getattr(ApiToken, self.resource_id_field) == resource_id,
)
).all()
# App and agent keys keep their existing (unmasked) list behavior; reveal-once
# masking is scoped to dataset keys, which build their list in datasets.py.
return ApiKeyList(data=[ApiKeyItem.model_validate(key, from_attributes=True) for key in keys])
@edit_permission_required
@with_session
def post(self, session: Session, resource_id: str, current_tenant_id: str) -> tuple[dict[str, object], int]:
return dump_response(
ApiKeyItem,
self._create_api_key(resource_id, current_tenant_id, session=session),
), 201
def _create_api_key(self, resource_id: str, current_tenant_id: str, *, session: Session) -> ApiToken:
assert self.resource_id_field is not None, "resource_id_field must be set"
resource = _get_resource(resource_id, current_tenant_id, self.resource_model, session=session)
if isinstance(resource, App):
AppService.ensure_agent_app_access_ready(resource, session=session)
current_key_count: int = (
session.scalar(
select(func.count(ApiToken.id)).where(
or_(ApiToken.tenant_id == current_tenant_id, ApiToken.tenant_id.is_(None)),
ApiToken.type == self.resource_type,
getattr(ApiToken, self.resource_id_field) == resource_id,
)
)
or 0
)
if current_key_count >= self.max_keys:
flask_restx.abort(
HTTPStatus.BAD_REQUEST,
message=f"Cannot create more than {self.max_keys} API keys for this resource type.",
custom="max_keys_exceeded",
)
key = ApiToken.generate_api_key(self.token_prefix or "", 24, session=session)
assert self.resource_type is not None, "resource_type must be set"
api_token = ApiToken()
setattr(api_token, self.resource_id_field, resource_id)
api_token.tenant_id = current_tenant_id
api_token.token = key
api_token.type = self.resource_type
session.add(api_token)
session.commit()
return api_token
class BaseApiKeyResource(Resource):
method_decorators = [account_initialization_required, login_required, setup_required]
resource_type: ApiTokenType | None = None
resource_model: type | None = None
resource_id_field: str | None = None
@with_session
def delete(
self,
session: Session,
resource_id: str,
api_key_id: str,
current_tenant_id: str,
current_user: Account,
) -> tuple[str, int]:
self._delete_api_key(resource_id, api_key_id, current_tenant_id, current_user, session=session)
return "", 204
def _delete_api_key(
self,
resource_id: str,
api_key_id: str,
current_tenant_id: str,
current_user: Account,
*,
session: Session,
) -> None:
assert self.resource_id_field is not None, "resource_id_field must be set"
_get_resource(resource_id, current_tenant_id, self.resource_model, session=session)
if not dify_config.RBAC_ENABLED and not current_user.is_admin_or_owner:
raise Forbidden()
key = session.scalar(
select(ApiToken)
.where(
or_(ApiToken.tenant_id == current_tenant_id, ApiToken.tenant_id.is_(None)),
getattr(ApiToken, self.resource_id_field) == resource_id,
ApiToken.type == self.resource_type,
ApiToken.id == api_key_id,
)
.limit(1)
)
if key is None:
flask_restx.abort(HTTPStatus.NOT_FOUND, message="API key not found")
# Invalidate cache before deleting from database
# Type assertion: key is guaranteed to be non-None here because abort() raises
assert key is not None # nosec - for type checker only
ApiTokenCache.delete(key.token, key.type)
session.delete(key)
session.commit()
@contextmanager
def api_key_errors() -> Generator[None]:
"""Translate application failures without changing the Console error contract."""
try:
yield
except (ApiKeyResourceNotFoundError, ApiKeyNotFoundError) as error:
flask_restx.abort(404, message=str(error))
except ApiKeyLimitExceededError as error:
flask_restx.abort(400, message=str(error), custom="max_keys_exceeded")
except UnknownDatasetIdsError as error:
flask_restx.abort(400, message=str(error))
except NoPermissionError as error:
flask_restx.abort(403, message=str(error))
except AppApiKeyNotReadyError as error:
raise AgentAccessNotReadyError from error
@console_ns.route("/apps/<uuid:resource_id>/api-keys")
class AppApiKeyListResource(BaseApiKeyListResource):
class AppApiKeyListResource(Resource):
@console_ns.doc("get_app_api_keys")
@console_ns.doc(description="Get all API keys for an app")
@console_ns.doc(params={"resource_id": "App ID"})
@console_ns.response(200, "API keys retrieved successfully", console_ns.models[ApiKeyList.__name__])
@with_current_tenant_id
@edit_permission_required
@rbac_permission_required(
RBACCheck(RBACPermission.APP_RELEASE_AND_VERSION, PlainApp("resource_id")),
RBACCheck(RBACPermission.AGENT_ACCESS_POINT_VIEW, AgentBehindApp("resource_id")),
@console_account_admission(
allowed_roles=API_KEY_EDIT_ROLES,
rbac_checks=[
RBACCheck(RBACPermission.APP_RELEASE_AND_VERSION, PlainApp("resource_id")),
RBACCheck(RBACPermission.AGENT_ACCESS_POINT_VIEW, AgentBehindApp("resource_id")),
],
)
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, resource_id: UUID) -> dict[str, object]:
"""Get all API keys for an app"""
return dump_response(
ApiKeyList,
self._get_api_key_list(str(resource_id), current_tenant_id, session=session), # pyrefly: ignore[unnecessary-type-conversion]
)
def get(self, request_context: RequestContext, resource_id: UUID) -> dict[str, object]:
with api_key_errors():
keys = application_services().app_api_keys.list_keys(request_context, str(resource_id))
return dump_response(ApiKeyList, {"data": keys})
@console_ns.doc("create_app_api_key")
@console_ns.doc(description="Create a new API key for an app")
@console_ns.doc(params={"resource_id": "App ID"})
@console_ns.response(201, "API key created successfully", console_ns.models[ApiKeyItem.__name__])
@console_ns.response(400, "Maximum keys exceeded")
@with_current_tenant_id
@edit_permission_required
@rbac_permission_required(
RBACCheck(RBACPermission.APP_RELEASE_AND_VERSION, PlainApp("resource_id")),
RBACCheck(RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentBehindApp("resource_id")),
@console_account_admission(
allowed_roles=API_KEY_EDIT_ROLES,
rbac_checks=[
RBACCheck(RBACPermission.APP_RELEASE_AND_VERSION, PlainApp("resource_id")),
RBACCheck(RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentBehindApp("resource_id")),
],
)
@with_session
def post(self, session: Session, current_tenant_id: str, resource_id: UUID) -> tuple[dict[str, object], int]:
"""Create a new API key for an app"""
return dump_response(
ApiKeyItem,
self._create_api_key(str(resource_id), current_tenant_id, session=session), # pyrefly: ignore[unnecessary-type-conversion]
), 201
resource_type = ApiTokenType.APP
resource_model = App
resource_id_field = "app_id"
token_prefix = "app-"
def post(self, request_context: RequestContext, resource_id: UUID) -> tuple[dict[str, object], int]:
with api_key_errors():
key = application_services().app_api_keys.create_key(request_context, str(resource_id))
return dump_response(ApiKeyItem, key), 201
@console_ns.route("/apps/<uuid:resource_id>/api-keys/<uuid:api_key_id>")
class AppApiKeyResource(BaseApiKeyResource):
class AppApiKeyResource(Resource):
@console_ns.doc("delete_app_api_key")
@console_ns.doc(description="Delete an API key for an app")
@console_ns.doc(params={"resource_id": "App ID", "api_key_id": "API key ID"})
@console_ns.response(204, "API key deleted successfully")
@with_current_user
@with_current_tenant_id
@rbac_permission_required(
RBACCheck(RBACPermission.APP_RELEASE_AND_VERSION, PlainApp("resource_id")),
RBACCheck(RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentBehindApp("resource_id")),
@console_account_admission(
allowed_roles=API_KEY_DELETE_ROLES,
rbac_checks=[
RBACCheck(RBACPermission.APP_RELEASE_AND_VERSION, PlainApp("resource_id")),
RBACCheck(RBACPermission.AGENT_ACCESS_POINT_MANAGE, AgentBehindApp("resource_id")),
],
)
@with_session
def delete(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
resource_id: UUID,
api_key_id: UUID,
) -> tuple[str, int]:
"""Delete an API key for an app"""
self._delete_api_key(
str(resource_id), # pyrefly: ignore[unnecessary-type-conversion]
str(api_key_id),
current_tenant_id,
current_user,
session=session,
)
def delete(self, request_context: RequestContext, resource_id: UUID, api_key_id: UUID) -> tuple[str, int]:
with api_key_errors():
application_services().app_api_keys.delete_key(request_context, str(resource_id), str(api_key_id))
return "", 204
resource_type = ApiTokenType.APP
resource_model = App
resource_id_field = "app_id"
# Dataset service-API keys are also managed at the workspace level (create with a set of
# knowledge bases, list, delete) by DatasetApiKeyApi in
# controllers/console/datasets/datasets.py, using DatasetApiTokenBinding for scoping.
# The per-dataset routes below remain for callers that key an API token to a single dataset.
@console_ns.route("/datasets/<uuid:resource_id>/api-keys")
class DatasetApiKeyListResource(BaseApiKeyListResource):
class DatasetApiKeyListResource(Resource):
@console_ns.doc("get_dataset_api_keys")
@console_ns.doc(description="Get all API keys for a dataset")
@console_ns.doc(params={"resource_id": "Dataset ID"})
@console_ns.response(200, "API keys retrieved successfully", console_ns.models[ApiKeyList.__name__])
@with_current_tenant_id
@edit_permission_required
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, DatasetId("resource_id")))
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str, resource_id: UUID) -> dict[str, object]:
"""Get all API keys for a dataset"""
return dump_response(
ApiKeyList,
self._get_api_key_list(str(resource_id), current_tenant_id, session=session), # pyrefly: ignore[unnecessary-type-conversion]
)
@console_account_admission(
allowed_roles=API_KEY_EDIT_ROLES,
rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, DatasetId("resource_id"))],
)
def get(self, request_context: RequestContext, resource_id: UUID) -> dict[str, object]:
with api_key_errors():
keys = application_services().dataset_api_keys.list_keys(request_context, str(resource_id))
return dump_response(ApiKeyList, build_masked_api_key_list(keys))
@console_ns.doc("create_dataset_api_key")
@console_ns.doc(description="Create a new API key for a dataset")
@console_ns.doc(params={"resource_id": "Dataset ID"})
@console_ns.response(201, "API key created successfully", console_ns.models[ApiKeyItem.__name__])
@console_ns.response(400, "Maximum keys exceeded")
@with_current_tenant_id
@edit_permission_required
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, DatasetId("resource_id")))
@with_session
def post(self, session: Session, current_tenant_id: str, resource_id: UUID) -> tuple[dict[str, object], int]:
"""Create a new API key for a dataset"""
return dump_response(
ApiKeyItem,
self._create_api_key(str(resource_id), current_tenant_id, session=session), # pyrefly: ignore[unnecessary-type-conversion]
), 201
resource_type = ApiTokenType.DATASET
resource_model = Dataset
resource_id_field = "dataset_id"
token_prefix = "ds-"
@console_account_admission(
allowed_roles=API_KEY_EDIT_ROLES,
rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, DatasetId("resource_id"))],
)
def post(self, request_context: RequestContext, resource_id: UUID) -> tuple[dict[str, object], int]:
with api_key_errors():
key = application_services().dataset_api_keys.create_key(request_context, str(resource_id))
return dump_response(ApiKeyItem, key), 201
@console_ns.route("/datasets/<uuid:resource_id>/api-keys/<uuid:api_key_id>")
class DatasetApiKeyResource(BaseApiKeyResource):
class DatasetApiKeyResource(Resource):
@console_ns.doc("delete_dataset_api_key")
@console_ns.doc(description="Delete an API key for a dataset")
@console_ns.doc(params={"resource_id": "Dataset ID", "api_key_id": "API key ID"})
@console_ns.response(204, "API key deleted successfully")
@with_current_user
@with_current_tenant_id
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, DatasetId("resource_id")))
@with_session
def delete(
self,
session: Session,
current_tenant_id: str,
current_user: Account,
resource_id: UUID,
api_key_id: UUID,
) -> tuple[str, int]:
"""Delete an API key for a dataset"""
self._delete_api_key(
str(resource_id), # pyrefly: ignore[unnecessary-type-conversion]
str(api_key_id),
current_tenant_id,
current_user,
session=session,
)
@console_account_admission(
allowed_roles=API_KEY_DELETE_ROLES,
rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, DatasetId("resource_id"))],
)
def delete(self, request_context: RequestContext, resource_id: UUID, api_key_id: UUID) -> tuple[str, int]:
with api_key_errors():
application_services().dataset_api_keys.delete_key(request_context, str(resource_id), str(api_key_id))
return "", 204
resource_type = ApiTokenType.DATASET
resource_model = Dataset
resource_id_field = "dataset_id"
api/controllers/console/datasets/datasets.py
+31 -103
View File
@@ -6,7 +6,7 @@ from uuid import UUID
from flask import request
from flask_restx import Resource
from pydantic import BaseModel, Field, field_validator, model_validator
from sqlalchemy import func, select
from sqlalchemy import select
from sqlalchemy.orm import Session
from werkzeug.exceptions import Forbidden, NotFound
@@ -17,15 +17,15 @@ from controllers.common.rbac import DatasetId, RBACCheck, Workspace, enforce_rba
from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
from controllers.common.session import with_session
from controllers.console import console_ns
from controllers.console.apikey import ApiKeyItem, ApiKeyList, build_masked_api_key_list
from controllers.console.apikey import API_KEY_DELETE_ROLES, api_key_errors
from controllers.console.app.error import ProviderNotInitializeError
from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError
from controllers.console.flask_admission import console_account_admission
from controllers.console.wraps import (
RBACPermission,
account_initialization_required,
cloud_edition_billing_rate_limit_check,
enterprise_license_required,
is_admin_or_owner_required,
model_validate,
rbac_permission_required,
setup_required,
@@ -41,6 +41,8 @@ from core.rag.extractor.entity.datasource_type import DatasourceType
from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
from core.rag.index_processor.constant.index_type import IndexTechniqueType
from core.rag.retrieval.retrieval_methods import RetrievalMethod
from extensions.ext_application_services import application_services
from fields.api_key_fields import ApiKeyItem, ApiKeyList, build_masked_api_key_list
from fields.base import ResponseModel
from fields.dataset_fields import (
DatasetDetailResponse,
@@ -52,12 +54,10 @@ from libs.helper import build_icon_url, dump_response, to_timestamp
from libs.login import login_required
from libs.pagination import clamp_pagination
from libs.url_utils import normalize_api_base_url
from models import Account, ApiToken, App, Dataset, Document, UploadFile
from machinery.context import RequestContext
from models import Account, App, Dataset, Document, UploadFile
from models.dataset import DatasetPermission, DatasetPermissionEnum, DatasetQuery
from models.enums import ApiTokenType
from models.provider_ids import ModelProviderID
from services import dataset_api_key_service
from services.api_token_service import ApiTokenCache
from services.app_service import AppService
from services.dataset_ref_service import DatasetRefService
from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
@@ -1113,121 +1113,49 @@ class DatasetIndexingStatusApi(Resource):
@console_ns.route("/datasets/api-keys")
class DatasetApiKeyApi(Resource):
max_keys = 10
token_prefix = "dataset-"
resource_type = ApiTokenType.DATASET
@console_ns.doc("get_dataset_api_keys")
@console_ns.doc(description="Get dataset API keys")
@console_ns.response(200, "API keys retrieved successfully", console_ns.models[ApiKeyList.__name__])
@setup_required
@login_required
@is_admin_or_owner_required
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace()))
@account_initialization_required
@with_current_tenant_id
@with_session(write=False)
def get(self, session: Session, current_tenant_id: str):
keys = session.scalars(
select(ApiToken).where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id)
).all()
token_ids = [str(key.id) for key in keys]
bindings_by_token = dataset_api_key_service.list_bindings_by_token(session, token_ids)
return dump_response(ApiKeyList, build_masked_api_key_list(keys, bindings_by_token))
@console_account_admission(
allowed_roles=API_KEY_DELETE_ROLES,
rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace())],
)
def get(self, request_context: RequestContext):
with api_key_errors():
keys = application_services().dataset_api_keys.list_workspace_keys(request_context)
return dump_response(ApiKeyList, build_masked_api_key_list(keys))
@console_ns.expect(console_ns.models[DatasetApiKeyCreatePayload.__name__])
@console_ns.response(200, "API key created successfully", console_ns.models[ApiKeyItem.__name__])
@console_ns.response(400, "Maximum keys exceeded")
@setup_required
@login_required
@is_admin_or_owner_required
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace()))
@account_initialization_required
@with_current_tenant_id
@with_session
def post(self, session: Session, current_tenant_id: str):
# Optional list of knowledge bases to scope the key to. Absent/empty => the key
# can access every dataset in the tenant (default). Duplicates are de-duplicated.
@console_account_admission(
allowed_roles=API_KEY_DELETE_ROLES,
rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace())],
)
def post(self, request_context: RequestContext):
# Preserve the optional scope and validation response of the workspace route.
payload = request.get_json(silent=True) or {}
raw_dataset_ids = payload.get("dataset_ids") or []
if not isinstance(raw_dataset_ids, list) or any(not isinstance(item, str) for item in raw_dataset_ids):
console_ns.abort(400, message="dataset_ids must be a list of strings.")
dataset_ids = list(dict.fromkeys(raw_dataset_ids))
if dataset_ids:
unknown = dataset_api_key_service.find_unknown_dataset_ids(session, dataset_ids, current_tenant_id)
if unknown:
console_ns.abort(400, message=f"Unknown knowledge base id(s): {', '.join(unknown)}")
current_key_count = (
session.scalar(
select(func.count(ApiToken.id)).where(
ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id
)
)
or 0
)
if current_key_count >= self.max_keys:
console_ns.abort(
400,
message=f"Cannot create more than {self.max_keys} API keys for this resource type.",
custom="max_keys_exceeded",
)
key = ApiToken.generate_api_key(self.token_prefix, 24, session=session)
api_token = ApiToken()
api_token.tenant_id = current_tenant_id
api_token.token = key
api_token.type = self.resource_type
session.add(api_token)
session.flush()
dataset_api_key_service.bind_datasets(session, api_token.id, dataset_ids)
session.flush()
# Reveal-once: the create response carries the full secret and its bound scope.
item = ApiKeyItem.model_validate(api_token, from_attributes=True)
item.dataset_ids = dataset_ids
return dump_response(ApiKeyItem, item), 200
with api_key_errors():
key = application_services().dataset_api_keys.create_workspace_key(request_context, tuple(raw_dataset_ids))
return dump_response(ApiKeyItem, key), 200
@console_ns.route("/datasets/api-keys/<uuid:api_key_id>")
class DatasetApiDeleteApi(Resource):
resource_type = ApiTokenType.DATASET
@console_ns.doc("delete_dataset_api_key")
@console_ns.doc(description="Delete dataset API key")
@console_ns.doc(params={"api_key_id": "API key ID"})
@console_ns.response(204, "API key deleted successfully")
@setup_required
@login_required
@is_admin_or_owner_required
@rbac_permission_required(RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace()))
@account_initialization_required
@with_current_tenant_id
@with_session
def delete(self, session: Session, current_tenant_id: str, api_key_id: UUID):
api_key_id_str = str(api_key_id)
key = session.scalar(
select(ApiToken)
.where(
ApiToken.tenant_id == current_tenant_id,
ApiToken.type == self.resource_type,
ApiToken.id == api_key_id_str,
)
.limit(1)
)
if key is None:
console_ns.abort(404, message="API key not found")
# Invalidate cache before deleting from database
# Type assertion: key is guaranteed to be non-None here because abort() raises
assert key is not None # nosec - for type checker only
ApiTokenCache.delete(key.token, key.type)
session.delete(key)
@console_account_admission(
allowed_roles=API_KEY_DELETE_ROLES,
rbac_checks=[RBACCheck(RBACPermission.DATASET_API_KEY_MANAGE, Workspace())],
)
def delete(self, request_context: RequestContext, api_key_id: UUID):
with api_key_errors():
application_services().dataset_api_keys.delete_workspace_key(request_context, str(api_key_id))
return "", 204
api/controllers/service_api/wraps.py
+2 -2
View File
@@ -30,7 +30,7 @@ from libs.login import current_user
from models import Account, Tenant, TenantAccountJoin, TenantStatus
from models.dataset import Dataset, RateLimitLog
from models.model import ApiToken, App
from services import dataset_api_key_service
from repositories.knowledge import dataset_api_key_bindings
from services.api_token_service import ApiTokenCache, fetch_token_with_single_flight, record_token_usage
from services.feature_service import FeatureService
@@ -343,7 +343,7 @@ def validate_dataset_token[R](view: Callable[..., R]) -> Callable[..., R]:
# per request (not cached) so scope changes take effect immediately.
# db.session is Flask-SQLAlchemy's scoped_session proxy; cast so the plain-Session
# typed helper accepts it (runtime proxies every Session method through unchanged).
bound_dataset_ids = dataset_api_key_service.get_bound_dataset_ids(cast(Session, db.session), api_token.id)
bound_dataset_ids = dataset_api_key_bindings.get_bound_dataset_ids(cast(Session, db.session), api_token.id)
if bound_dataset_ids and (not dataset_id or str(dataset_id) not in bound_dataset_ids):
raise Forbidden("The API key is not authorized to access this knowledge base.")
api/extensions/application_services/app.py
+14
View File
@@ -0,0 +1,14 @@
"""Compose App API access dependencies."""
from sqlalchemy.orm import Session, sessionmaker
from repositories.app.api_key_repository import AppApiKeyRepository
from services.api_token_service import ApiTokenCache
from services.app.api_key_service import AppApiKeyService
def build_app_api_key_service(*, database_client: sessionmaker[Session]) -> AppApiKeyService:
return AppApiKeyService(
keys=AppApiKeyRepository(session_factory=database_client),
cache=ApiTokenCache,
)
api/extensions/application_services/knowledge.py
+16
View File
@@ -0,0 +1,16 @@
"""Composition of knowledge base application services."""
from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config
from repositories.knowledge.dataset_api_key_repository import DatasetApiKeyRepository
from services.api_token_service import ApiTokenCache
from services.knowledge.api_key_service import DatasetApiKeyService
def build_dataset_api_key_service(*, database_client: sessionmaker[Session]) -> DatasetApiKeyService:
return DatasetApiKeyService(
keys=DatasetApiKeyRepository(session_factory=database_client),
cache=ApiTokenCache,
rbac_enabled=lambda: dify_config.RBAC_ENABLED,
)
api/extensions/ext_application_services.py
+8
View File
@@ -23,6 +23,8 @@ from core.schemas.schema_manager import SchemaManager
from core.tools.tool_file_manager import ToolFileManager
from enums import DeploymentEdition, WebAppAccessMode
from extensions.application_services.agent import AgentAppServices, build_agent_app_services
from extensions.application_services.app import build_app_api_key_service
from extensions.application_services.knowledge import build_dataset_api_key_service
from extensions.ext_redis import RedisClientWrapper, redis_client
from extensions.ext_storage import storage
from libs.datetime_utils import naive_utc_now, utc_now
@@ -143,6 +145,7 @@ from services.account_password_hasher import DefaultAccountPasswordHasher
from services.account_password_service import AccountPasswordService
from services.account_profile_service import AccountProfileService
from services.app.advanced_prompt_template_service import AdvancedPromptTemplateService
from services.app.api_key_service import AppApiKeyService
from services.app_audio_adapters import AppAudioRuntime
from services.app_audio_service import AppAudio
from services.app_definition_query_service import AppDefinitionQueryService
@@ -180,6 +183,7 @@ from services.inner_mail_service import InnerMailService
from services.installed_app_access_service import InstalledAppAccessService
from services.installed_app_generation_adapters import AppGenerateServiceRuntime as InstalledAppGenerateServiceRuntime
from services.installed_app_generation_service import InstalledAppGenerationService
from services.knowledge.api_key_service import DatasetApiKeyService
from services.message_file_preview_service import MessageFilePreviewService
from services.message_suggested_questions_adapters import MessageSuggestedQuestionsRuntime
from services.message_suggested_questions_service import MessageSuggestedQuestions
@@ -298,6 +302,8 @@ class ApplicationServices:
agent_apps: AgentAppServices
advanced_prompt_templates: AdvancedPromptTemplateService
accounts: AccountServices
app_api_keys: AppApiKeyService
dataset_api_keys: DatasetApiKeyService
account_activation: AccountActivationService
app_definitions: AppDefinitionQueryService
app_preview_details: AppPreviewDetails
@@ -717,6 +723,8 @@ def build_application_services(
app_sites=AppSiteService(
sites=AppSiteCommandRepository(session_factory=database_client),
),
app_api_keys=build_app_api_key_service(database_client=database_client),
dataset_api_keys=build_dataset_api_key_service(database_client=database_client),
app_statistics=AppStatisticQueryRepository(session_factory=database_client),
app_tracing_configs=AppTracingConfigService(
configs=SQLAlchemyAppTracingConfigRepository(session_factory=database_client),
api/fields/api_key_fields.py
+61
View File
@@ -0,0 +1,61 @@
"""Response contracts shared by Console API key endpoints."""
from collections.abc import Iterable, Mapping
from datetime import datetime
from pydantic import field_validator
from fields.base import ResponseModel
from libs.helper import to_timestamp
class ApiKeyItem(ResponseModel):
id: str
type: str
token: str
# Dataset keys only: the knowledge bases this key is bound to. Empty = the key can
# access every dataset in the tenant (default). App keys are always empty.
dataset_ids: list[str] = []
last_used_at: int | None = None
created_at: int | None = None
@field_validator("last_used_at", "created_at", mode="before")
@classmethod
def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
return to_timestamp(value)
class ApiKeyList(ResponseModel):
data: list[ApiKeyItem]
def mask_api_token(token: str) -> str:
"""Mask a secret token for list responses.
Reveal-once: the full secret is only returned by the create endpoint. List
endpoints expose just enough (prefix + last 4) to identify a key, never the
full value, so an existing key's secret cannot be retrieved after creation.
"""
if len(token) <= 8:
return "***"
return f"{token[:5]}...{token[-4:]}"
def build_masked_api_key_list(
api_tokens: Iterable[object],
bindings_by_token: Mapping[str, list[str]] | None = None,
) -> ApiKeyList:
"""Build an ApiKeyList from token records with their secrets masked.
``bindings_by_token`` maps an api_token id to the dataset ids it is bound to
(from DatasetApiTokenBinding); tokens absent from the map are unbound (empty =
access all). When omitted, preserve the scope carried by the records.
"""
items: list[ApiKeyItem] = []
for api_token in api_tokens:
item = ApiKeyItem.model_validate(api_token, from_attributes=True)
item.token = mask_api_token(item.token)
if bindings_by_token is not None:
item.dataset_ids = bindings_by_token.get(item.id, [])
items.append(item)
return ApiKeyList(data=items)
api/models/model.py
-3
View File
@@ -2236,9 +2236,6 @@ class ApiToken(Base):
key with binding rows is limited to exactly those datasets). Enforcement lives in
``validate_dataset_token`` (controllers/service_api/wraps.py).
Note: controllers/console/apikey.py assigns the ``*_id`` columns via ``setattr``
keyed on ``resource_id_field``, so renaming ``app_id`` requires updating those
controllers too.
"""
__tablename__ = "api_tokens"
api/openapi/markdown/console-openapi.md
+6 -6
View File
@@ -4807,7 +4807,7 @@ Restore a published workflow version into the draft workflow
| 404 | Workflow not found | |
### [GET] /apps/{resource_id}/api-keys
**Get all API keys for an app**
Get all API keys for an app
#### Parameters
@@ -4822,7 +4822,7 @@ Restore a published workflow version into the draft workflow
| 200 | API keys retrieved successfully | **application/json**: [ApiKeyList](#apikeylist)<br> |
### [POST] /apps/{resource_id}/api-keys
**Create a new API key for an app**
Create a new API key for an app
#### Parameters
@@ -4838,7 +4838,7 @@ Restore a published workflow version into the draft workflow
| 400 | Maximum keys exceeded | |
### [DELETE] /apps/{resource_id}/api-keys/{api_key_id}
**Delete an API key for an app**
Delete an API key for an app
#### Parameters
@@ -6447,7 +6447,7 @@ Check if dataset is in use
| 200 | Dataset use status retrieved successfully | **application/json**: [UsageCheckResponse](#usagecheckresponse)<br> |
### [GET] /datasets/{resource_id}/api-keys
**Get all API keys for a dataset**
Get all API keys for a dataset
#### Parameters
@@ -6462,7 +6462,7 @@ Check if dataset is in use
| 200 | API keys retrieved successfully | **application/json**: [ApiKeyList](#apikeylist)<br> |
### [POST] /datasets/{resource_id}/api-keys
**Create a new API key for a dataset**
Create a new API key for a dataset
#### Parameters
@@ -6478,7 +6478,7 @@ Check if dataset is in use
| 400 | Maximum keys exceeded | |
### [DELETE] /datasets/{resource_id}/api-keys/{api_key_id}
**Delete an API key for a dataset**
Delete an API key for a dataset
#### Parameters
api/repositories/app/api_key_repository.py
+124
View File
@@ -0,0 +1,124 @@
"""App API keys with tenant scoping and repository-owned transactions."""
from typing import override
from sqlalchemy import Select, func, or_, select
from sqlalchemy.orm import Session, sessionmaker
from core.agent.publish_visibility import workflow_callable_active_snapshot_filter
from models.agent import APP_BACKED_AGENT_SOURCES, Agent, AgentScope, AgentStatus
from models.enums import ApiTokenType, AppStatus
from models.model import ApiToken, App, AppMode
from services.app.api_key_service import AppApiKeyAccessState, AppApiKeyStore
from services.auth.api_key_contracts import (
ApiKeyLimitExceededError,
ApiKeyNotFoundError,
ApiKeyRecord,
ApiKeyResourceNotFoundError,
)
class AppApiKeyRepository(AppApiKeyStore):
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
@override
def resolve_agent_app_id(self, workspace_id: str, agent_id: str) -> str:
with self._session_factory() as session:
app_id = session.scalar(
select(App.id)
.join(Agent, Agent.app_id == App.id)
.where(
App.tenant_id == workspace_id,
App.mode == AppMode.AGENT,
App.status == AppStatus.NORMAL,
Agent.tenant_id == workspace_id,
Agent.id == agent_id,
Agent.scope == AgentScope.ROSTER,
Agent.source.in_(APP_BACKED_AGENT_SOURCES),
Agent.status == AgentStatus.ACTIVE,
)
)
if app_id is None:
raise ApiKeyResourceNotFoundError("Agent not found.")
return app_id
@override
def get_access_state(self, workspace_id: str, app_id: str) -> AppApiKeyAccessState:
with self._session_factory() as session:
app = self._get_app(session, workspace_id, app_id)
if app.mode != AppMode.AGENT:
return AppApiKeyAccessState(is_agent_app=False, has_published_snapshot=False)
published = session.scalar(
select(Agent.id)
.where(
Agent.tenant_id == workspace_id,
Agent.app_id == app_id,
Agent.scope == AgentScope.ROSTER,
Agent.source.in_(APP_BACKED_AGENT_SOURCES),
Agent.status == AgentStatus.ACTIVE,
workflow_callable_active_snapshot_filter(),
)
.limit(1)
)
return AppApiKeyAccessState(is_agent_app=True, has_published_snapshot=published is not None)
@override
def list_keys(self, workspace_id: str, app_id: str) -> tuple[ApiKeyRecord, ...]:
with self._session_factory() as session:
self._get_app(session, workspace_id, app_id)
return tuple(self._record(key) for key in session.scalars(self._key_query(workspace_id, app_id)))
@override
def create_key(self, workspace_id: str, app_id: str, *, max_keys: int, prefix: str) -> ApiKeyRecord:
with self._session_factory.begin() as session:
# Serialize count-and-insert for the App, including historical keys.
self._get_app(session, workspace_id, app_id, for_update=True)
count = session.scalar(select(func.count()).select_from(self._key_query(workspace_id, app_id).subquery()))
if count is not None and count >= max_keys:
raise ApiKeyLimitExceededError(max_keys)
key = ApiToken(
tenant_id=workspace_id,
app_id=app_id,
type=ApiTokenType.APP,
token=ApiToken.generate_api_key(prefix, 24, session=session),
)
session.add(key)
session.flush()
return self._record(key)
@override
def delete_key(self, workspace_id: str, app_id: str, key_id: str) -> ApiKeyRecord:
with self._session_factory.begin() as session:
self._get_app(session, workspace_id, app_id)
key = session.scalar(self._key_query(workspace_id, app_id).where(ApiToken.id == key_id))
if key is None:
raise ApiKeyNotFoundError
record = self._record(key)
session.delete(key)
return record
@staticmethod
def _get_app(session: Session, workspace_id: str, app_id: str, *, for_update: bool = False) -> App:
query = select(App).where(App.id == app_id, App.tenant_id == workspace_id)
if for_update:
query = query.with_for_update()
app = session.scalar(query)
if app is None:
raise ApiKeyResourceNotFoundError("App not found.")
return app
@staticmethod
def _key_query(workspace_id: str, app_id: str) -> Select[tuple[ApiToken]]:
# Historical keys have no tenant_id; the owning App is validated first.
return select(ApiToken).where(
or_(ApiToken.tenant_id == workspace_id, ApiToken.tenant_id.is_(None)),
ApiToken.type == ApiTokenType.APP,
ApiToken.app_id == app_id,
)
@staticmethod
def _record(key: ApiToken) -> ApiKeyRecord:
return ApiKeyRecord(
id=key.id, type=key.type, token=key.token, last_used_at=key.last_used_at, created_at=key.created_at
)
api/repositories/knowledge/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Knowledge base persistence and queries."""
api/services/dataset_api_key_service.py → api/repositories/knowledge/dataset_api_key_bindings.py
+13 -15
View File
@@ -6,14 +6,13 @@ A dataset API key is scoped to specific knowledge bases through
- no rows -> the key can access every dataset in its tenant (default / back-compat)
- N rows -> the key is restricted to exactly those datasets
These helpers live in the service layer so controllers and the service-API auth
decorator can share the queries without embedding SQLAlchemy in controller code.
Shared by key management, dataset deletion, and service-API authentication.
The caller owns the transaction boundary (commit/flush).
"""
from collections.abc import Iterable
from sqlalchemy import delete, select
from sqlalchemy import Select, delete, select
from sqlalchemy.orm import Session
from models.dataset import Dataset
@@ -77,18 +76,17 @@ def delete_keys_scoped_only_to(session: Session, dataset_id: str) -> list[str]:
Must run before the dataset row is deleted (so the bindings still exist).
Returns the deleted api_token ids; the caller controls the transaction.
"""
orphan_ids = [
str(token_id)
for token_id in session.scalars(
select(DatasetApiTokenBinding.api_token_id)
.where(DatasetApiTokenBinding.dataset_id == dataset_id)
.where(
DatasetApiTokenBinding.api_token_id.notin_(
select(DatasetApiTokenBinding.api_token_id).where(DatasetApiTokenBinding.dataset_id != dataset_id)
)
)
).all()
]
orphan_ids = [str(token_id) for token_id in session.scalars(token_ids_scoped_only_to(dataset_id)).all()]
if orphan_ids:
session.execute(delete(ApiToken).where(ApiToken.id.in_(orphan_ids)))
return orphan_ids
def token_ids_scoped_only_to(dataset_id: str) -> Select[tuple[str]]:
"""Select keys bound only to one dataset, for management and deletion."""
return select(DatasetApiTokenBinding.api_token_id).where(
DatasetApiTokenBinding.dataset_id == dataset_id,
DatasetApiTokenBinding.api_token_id.notin_(
select(DatasetApiTokenBinding.api_token_id).where(DatasetApiTokenBinding.dataset_id != dataset_id)
),
)
api/repositories/knowledge/dataset_api_key_repository.py
+156
View File
@@ -0,0 +1,156 @@
"""Knowledge base API keys with repository-owned sessions and transactions."""
from typing import override
from sqlalchemy import Select, func, select
from sqlalchemy.orm import Session, sessionmaker
from models.account import Tenant, TenantAccountJoin
from models.dataset import Dataset, DatasetPermission
from models.enums import ApiTokenType
from models.model import ApiToken
from repositories.knowledge import dataset_api_key_bindings
from services.auth.api_key_contracts import (
ApiKeyLimitExceededError,
ApiKeyNotFoundError,
ApiKeyRecord,
ApiKeyResourceNotFoundError,
)
from services.knowledge.api_key_service import DatasetApiKeyStore, UnknownDatasetIdsError
from services.knowledge.dataset_access import DatasetAccess
class DatasetApiKeyRepository(DatasetApiKeyStore):
def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
self._session_factory = session_factory
@override
def get_dataset_access(self, workspace_id: str, dataset_id: str, account_id: str) -> DatasetAccess:
with self._session_factory() as session:
dataset = self._get_dataset(session, workspace_id, dataset_id)
role = session.scalar(
select(TenantAccountJoin.role).where(
TenantAccountJoin.tenant_id == workspace_id, TenantAccountJoin.account_id == account_id
)
)
has_permission = (
session.scalar(
select(DatasetPermission.id)
.where(
DatasetPermission.tenant_id == workspace_id,
DatasetPermission.dataset_id == dataset_id,
DatasetPermission.account_id == account_id,
DatasetPermission.has_permission.is_(True),
)
.limit(1)
)
is not None
)
return DatasetAccess(dataset.permission, dataset.maintainer, role, has_permission)
@override
def list_keys(self, workspace_id: str, dataset_id: str) -> tuple[ApiKeyRecord, ...]:
with self._session_factory() as session:
self._get_dataset(session, workspace_id, dataset_id)
return tuple(
self._record(key, (dataset_id,)) for key in session.scalars(self._key_query(workspace_id, dataset_id))
)
@override
def create_key(self, workspace_id: str, dataset_id: str, *, max_keys: int, prefix: str) -> ApiKeyRecord:
with self._session_factory.begin() as session:
self._lock_workspace(session, workspace_id)
self._get_dataset(session, workspace_id, dataset_id)
return self._create_key(session, workspace_id, (dataset_id,), max_keys=max_keys, prefix=prefix)
@override
def list_workspace_keys(self, workspace_id: str) -> tuple[ApiKeyRecord, ...]:
with self._session_factory() as session:
keys = session.scalars(self._workspace_key_query(workspace_id)).all()
bindings = dataset_api_key_bindings.list_bindings_by_token(session, (key.id for key in keys))
return tuple(self._record(key, tuple(bindings.get(key.id, ()))) for key in keys)
@override
def create_workspace_key(
self, workspace_id: str, dataset_ids: tuple[str, ...], *, max_keys: int, prefix: str
) -> ApiKeyRecord:
with self._session_factory.begin() as session:
self._lock_workspace(session, workspace_id)
unknown = dataset_api_key_bindings.find_unknown_dataset_ids(session, list(dataset_ids), workspace_id)
if unknown:
raise UnknownDatasetIdsError(unknown)
return self._create_key(session, workspace_id, dataset_ids, max_keys=max_keys, prefix=prefix)
@override
def delete_workspace_key(self, workspace_id: str, key_id: str) -> ApiKeyRecord:
with self._session_factory.begin() as session:
key = session.scalar(self._workspace_key_query(workspace_id).where(ApiToken.id == key_id))
if key is None:
raise ApiKeyNotFoundError
bindings = dataset_api_key_bindings.get_bound_dataset_ids(session, key.id)
record = self._record(key, tuple(sorted(bindings)))
session.delete(key)
return record
def _create_key(
self, session: Session, workspace_id: str, dataset_ids: tuple[str, ...], *, max_keys: int, prefix: str
) -> ApiKeyRecord:
count = session.scalar(select(func.count()).select_from(self._workspace_key_query(workspace_id).subquery()))
if count is not None and count >= max_keys:
raise ApiKeyLimitExceededError(max_keys)
key = ApiToken(
tenant_id=workspace_id,
type=ApiTokenType.DATASET,
token=ApiToken.generate_api_key(prefix, 24, session=session),
)
session.add(key)
session.flush()
dataset_api_key_bindings.bind_datasets(session, key.id, dataset_ids)
return self._record(key, dataset_ids)
@staticmethod
def _lock_workspace(session: Session, workspace_id: str) -> None:
# Both creation routes lock the same row before counting all workspace dataset keys.
if session.scalar(select(Tenant.id).where(Tenant.id == workspace_id).with_for_update()) is None:
raise ApiKeyResourceNotFoundError("Workspace not found.")
@override
def delete_key(self, workspace_id: str, dataset_id: str, key_id: str) -> ApiKeyRecord:
with self._session_factory.begin() as session:
self._get_dataset(session, workspace_id, dataset_id)
key = session.scalar(self._key_query(workspace_id, dataset_id).where(ApiToken.id == key_id))
if key is None:
raise ApiKeyNotFoundError
record = self._record(key, (dataset_id,))
session.delete(key)
return record
@staticmethod
def _get_dataset(session: Session, workspace_id: str, dataset_id: str) -> Dataset:
query = select(Dataset).where(Dataset.id == dataset_id, Dataset.tenant_id == workspace_id)
dataset = session.scalar(query)
if dataset is None:
raise ApiKeyResourceNotFoundError("Dataset not found.")
return dataset
@staticmethod
def _workspace_key_query(workspace_id: str) -> Select[tuple[ApiToken]]:
return select(ApiToken).where(ApiToken.tenant_id == workspace_id, ApiToken.type == ApiTokenType.DATASET)
@classmethod
def _key_query(cls, workspace_id: str, dataset_id: str) -> Select[tuple[ApiToken]]:
# A per-dataset permission cannot expose or revoke broader workspace keys.
return cls._workspace_key_query(workspace_id).where(
ApiToken.id.in_(dataset_api_key_bindings.token_ids_scoped_only_to(dataset_id)),
)
@staticmethod
def _record(key: ApiToken, dataset_ids: tuple[str, ...]) -> ApiKeyRecord:
return ApiKeyRecord(
id=key.id,
type=key.type,
token=key.token,
last_used_at=key.last_used_at,
created_at=key.created_at,
dataset_ids=dataset_ids,
)
api/services/app/api_key_service.py
+61
View File
@@ -0,0 +1,61 @@
"""Manage App API access, including Agent routes backed by an App."""
from dataclasses import dataclass
from typing import Protocol
from machinery.context import RequestContext
from services.auth.api_key_contracts import ApiKeyCache, ApiKeyRecord
@dataclass(frozen=True, slots=True)
class AppApiKeyAccessState:
is_agent_app: bool
has_published_snapshot: bool
class AppApiKeyStore(Protocol):
def resolve_agent_app_id(self, workspace_id: str, agent_id: str) -> str: ...
def get_access_state(self, workspace_id: str, app_id: str) -> AppApiKeyAccessState: ...
def list_keys(self, workspace_id: str, app_id: str) -> tuple[ApiKeyRecord, ...]: ...
def create_key(self, workspace_id: str, app_id: str, *, max_keys: int, prefix: str) -> ApiKeyRecord: ...
def delete_key(self, workspace_id: str, app_id: str, key_id: str) -> ApiKeyRecord: ...
class AppApiKeyNotReadyError(Exception):
pass
class AppApiKeyService:
def __init__(self, *, keys: AppApiKeyStore, cache: ApiKeyCache) -> None:
self._keys = keys
self._cache = cache
def list_keys(self, context: RequestContext, app_id: str) -> tuple[ApiKeyRecord, ...]:
return self._keys.list_keys(context.active_workspace_id, app_id)
def create_key(self, context: RequestContext, app_id: str) -> ApiKeyRecord:
access = self._keys.get_access_state(context.active_workspace_id, app_id)
if access.is_agent_app and not access.has_published_snapshot:
raise AppApiKeyNotReadyError
return self._keys.create_key(context.active_workspace_id, app_id, max_keys=10, prefix="app-")
def delete_key(self, context: RequestContext, app_id: str, key_id: str) -> None:
key = self._keys.delete_key(context.active_workspace_id, app_id, key_id)
# External cache I/O follows the repository's committed transaction.
self._cache.delete(key.token, key.type)
def list_agent_keys(self, context: RequestContext, agent_id: str) -> tuple[ApiKeyRecord, ...]:
app_id = self._keys.resolve_agent_app_id(context.active_workspace_id, agent_id)
return self.list_keys(context, app_id)
def create_agent_key(self, context: RequestContext, agent_id: str) -> ApiKeyRecord:
app_id = self._keys.resolve_agent_app_id(context.active_workspace_id, agent_id)
return self.create_key(context, app_id)
def delete_agent_key(self, context: RequestContext, agent_id: str, key_id: str) -> None:
app_id = self._keys.resolve_agent_app_id(context.active_workspace_id, agent_id)
self.delete_key(context, app_id, key_id)
api/services/auth/api_key_contracts.py
+33
View File
@@ -0,0 +1,33 @@
"""Framework-neutral data contracts for API key management."""
from dataclasses import dataclass
from datetime import datetime
from typing import Protocol
@dataclass(frozen=True, slots=True)
class ApiKeyRecord:
id: str
type: str
token: str
last_used_at: datetime | None = None
created_at: datetime | None = None
dataset_ids: tuple[str, ...] = ()
class ApiKeyCache(Protocol):
def delete(self, token: str, scope: str | None = None) -> bool | None: ...
class ApiKeyResourceNotFoundError(Exception):
pass
class ApiKeyNotFoundError(Exception):
def __init__(self) -> None:
super().__init__("API key not found")
class ApiKeyLimitExceededError(Exception):
def __init__(self, max_keys: int) -> None:
super().__init__(f"Cannot create more than {max_keys} API keys for this resource type.")
api/services/dataset_service.py
+25 -20
View File
@@ -64,7 +64,7 @@ from models.model import UploadFile
from models.provider_ids import ModelProviderID
from models.source import DataSourceOauthBinding
from models.workflow import Workflow
from services import dataset_api_key_service
from repositories.knowledge import dataset_api_key_bindings
from services.dataset_ref_service import DatasetRef, DatasetRefService, SegmentRef
from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy
from services.document_indexing_proxy.duplicate_document_indexing_task_proxy import DuplicateDocumentIndexingTaskProxy
@@ -89,6 +89,7 @@ from services.errors.file import FileNotExistsError
from services.external_knowledge_service import ExternalDatasetService
from services.feature_service import FeatureService
from services.file_service import FileService
from services.knowledge.dataset_access import DatasetAccess
from services.rag_pipeline.rag_pipeline import RagPipelineService
from services.tag_service import TagService
from services.vector_service import VectorService
@@ -1363,7 +1364,7 @@ class DatasetService:
# Remove any dataset API key scoped only to this knowledge base, so it cannot
# silently degrade to unrestricted (access-all) once its last binding is gone.
dataset_api_key_service.delete_keys_scoped_only_to(session, str(dataset.id))
dataset_api_key_bindings.delete_keys_scoped_only_to(session, str(dataset.id))
session.delete(dataset)
session.commit()
@@ -1380,25 +1381,29 @@ class DatasetService:
if dataset.tenant_id != user.current_tenant_id:
logger.debug("User %s does not have permission to access dataset %s", user.id, dataset.id)
raise NoPermissionError("You do not have permission to access this dataset.")
if user.current_role != TenantAccountRole.OWNER:
if dataset.permission == DatasetPermissionEnum.ONLY_ME and dataset.maintainer != user.id:
logger.debug("User %s does not have permission to access dataset %s", user.id, dataset.id)
raise NoPermissionError("You do not have permission to access this dataset.")
if dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM:
# For partial team permission, user needs explicit permission or be the maintainer.
if dataset.maintainer != user.id:
user_permission = session.scalar(
select(DatasetPermission)
.where(
DatasetPermission.dataset_id == dataset.id,
DatasetPermission.account_id == user.id,
DatasetPermission.tenant_id == dataset.tenant_id,
)
.limit(1)
has_permission = False
if (
user.current_role != TenantAccountRole.OWNER
and dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM
and dataset.maintainer != user.id
):
has_permission = (
session.scalar(
select(DatasetPermission.id)
.where(
DatasetPermission.dataset_id == dataset.id,
DatasetPermission.account_id == user.id,
DatasetPermission.tenant_id == dataset.tenant_id,
DatasetPermission.has_permission.is_(True),
)
if not user_permission:
logger.debug("User %s does not have permission to access dataset %s", user.id, dataset.id)
raise NoPermissionError("You do not have permission to access this dataset.")
.limit(1)
)
is not None
)
access = DatasetAccess(dataset.permission, dataset.maintainer, user.current_role, has_permission)
if not access.allows(user.id):
logger.debug("User %s does not have permission to access dataset %s", user.id, dataset.id)
raise NoPermissionError("You do not have permission to access this dataset.")
@staticmethod
def check_dataset_operator_permission(
api/services/knowledge/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Knowledge base application services."""
api/services/knowledge/api_key_service.py
+73
View File
@@ -0,0 +1,73 @@
"""Manage workspace knowledge API keys and keys restricted to one knowledge base."""
from collections.abc import Callable
from typing import Protocol
from machinery.context import RequestContext
from services.auth.api_key_contracts import ApiKeyCache, ApiKeyRecord
from services.errors.account import NoPermissionError
from services.knowledge.dataset_access import DatasetAccess
class DatasetApiKeyStore(Protocol):
def get_dataset_access(self, workspace_id: str, dataset_id: str, account_id: str) -> DatasetAccess: ...
def list_keys(self, workspace_id: str, dataset_id: str) -> tuple[ApiKeyRecord, ...]: ...
def create_key(self, workspace_id: str, dataset_id: str, *, max_keys: int, prefix: str) -> ApiKeyRecord: ...
def delete_key(self, workspace_id: str, dataset_id: str, key_id: str) -> ApiKeyRecord: ...
def list_workspace_keys(self, workspace_id: str) -> tuple[ApiKeyRecord, ...]: ...
def create_workspace_key(
self, workspace_id: str, dataset_ids: tuple[str, ...], *, max_keys: int, prefix: str
) -> ApiKeyRecord: ...
def delete_workspace_key(self, workspace_id: str, key_id: str) -> ApiKeyRecord: ...
class UnknownDatasetIdsError(Exception):
def __init__(self, dataset_ids: list[str]) -> None:
super().__init__(f"Unknown knowledge base id(s): {', '.join(dataset_ids)}")
class DatasetApiKeyService:
MAX_KEYS = 10
def __init__(self, *, keys: DatasetApiKeyStore, cache: ApiKeyCache, rbac_enabled: Callable[[], bool]) -> None:
self._keys = keys
self._cache = cache
self._rbac_enabled = rbac_enabled
def list_keys(self, context: RequestContext, dataset_id: str) -> tuple[ApiKeyRecord, ...]:
self._check_access(context, dataset_id)
return self._keys.list_keys(context.active_workspace_id, dataset_id)
def create_key(self, context: RequestContext, dataset_id: str) -> ApiKeyRecord:
self._check_access(context, dataset_id)
return self._keys.create_key(context.active_workspace_id, dataset_id, max_keys=self.MAX_KEYS, prefix="ds-")
def delete_key(self, context: RequestContext, dataset_id: str, key_id: str) -> None:
self._check_access(context, dataset_id)
key = self._keys.delete_key(context.active_workspace_id, dataset_id, key_id)
self._cache.delete(key.token, key.type)
def _check_access(self, context: RequestContext, dataset_id: str) -> None:
# Enterprise RBAC is enforced at admission; community deployments use the dataset ACL.
if not self._rbac_enabled():
access = self._keys.get_dataset_access(context.active_workspace_id, dataset_id, context.account_id)
if not access.allows(context.account_id):
raise NoPermissionError("You do not have permission to access this dataset.")
def list_workspace_keys(self, context: RequestContext) -> tuple[ApiKeyRecord, ...]:
return self._keys.list_workspace_keys(context.active_workspace_id)
def create_workspace_key(self, context: RequestContext, dataset_ids: tuple[str, ...]) -> ApiKeyRecord:
return self._keys.create_workspace_key(
context.active_workspace_id, tuple(dict.fromkeys(dataset_ids)), max_keys=self.MAX_KEYS, prefix="dataset-"
)
def delete_workspace_key(self, context: RequestContext, key_id: str) -> None:
key = self._keys.delete_workspace_key(context.active_workspace_id, key_id)
self._cache.delete(key.token, key.type)
api/services/knowledge/dataset_access.py
+22
View File
@@ -0,0 +1,22 @@
"""Dataset ACL policy independent of persistence and request globals."""
from dataclasses import dataclass
@dataclass(frozen=True, slots=True)
class DatasetAccess:
permission: str
maintainer_id: str | None
role: str | None
has_explicit_permission: bool
def allows(self, account_id: str) -> bool:
if self.role is None:
return False
if self.role == "owner" or self.maintainer_id == account_id:
return True
if self.permission == "only_me":
return False
if self.permission == "partial_members":
return self.has_explicit_permission
return self.permission == "all_team_members"
api/tests/test_containers_integration_tests/controllers/console/test_apikey.py
+67 -25
View File
@@ -2,20 +2,22 @@
from __future__ import annotations
from unittest.mock import MagicMock, patch
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor, TimeoutError
from threading import Event
from uuid import uuid4
import pytest
from flask import Flask
from flask.testing import FlaskClient
from sqlalchemy import delete, select
from sqlalchemy.orm import Session
from sqlalchemy.orm import Session, sessionmaker
from models import Account
from models.account import AccountStatus, TenantAccountRole
from models.account import TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset
from models.enums import ApiTokenType, DataSourceType
from models.model import ApiToken, App, AppMode, DatasetApiTokenBinding
from repositories.knowledge.dataset_api_key_repository import DatasetApiKeyRepository
from services.auth.api_key_contracts import ApiKeyLimitExceededError
from tests.test_containers_integration_tests.controllers.console.helpers import (
authenticate_console_client,
create_console_account_and_tenant,
@@ -58,7 +60,7 @@ def setup_dataset(
@pytest.fixture(autouse=True)
def cleanup_api_tokens(db_session_with_containers: Session):
def cleanup_api_tokens(db_session_with_containers: Session) -> Iterator[None]:
"""Remove API tokens created during each test."""
yield
db_session_with_containers.execute(delete(ApiToken))
@@ -186,28 +188,25 @@ class TestAppApiKeyResource:
def test_delete_forbidden_for_non_admin(
self,
flask_app_with_containers: Flask,
db_session_with_containers: Session,
test_client_with_containers: FlaskClient,
) -> None:
"""A non-admin member cannot delete API keys via the controller permission check."""
from werkzeug.exceptions import Forbidden
"""A non-admin member cannot delete API keys through Console admission."""
account, tenant = create_console_account_and_tenant(db_session_with_containers)
app = create_console_app(db_session_with_containers, tenant.id, account.id, AppMode.CHAT)
membership = db_session_with_containers.scalar(
select(TenantAccountJoin).where(
TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id
)
)
assert membership is not None
membership.role = TenantAccountRole.NORMAL
db_session_with_containers.commit()
headers = authenticate_console_client(test_client_with_containers, account)
from controllers.console.apikey import BaseApiKeyResource
response = test_client_with_containers.delete(f"/console/api/apps/{app.id}/api-keys/{uuid4()}", headers=headers)
resource = BaseApiKeyResource()
resource.resource_type = ApiTokenType.APP
resource.resource_model = MagicMock()
resource.resource_id_field = "app_id"
non_admin = Account(name="Normal User", email="normal@example.com", status=AccountStatus.ACTIVE)
non_admin.id = "normal-user"
non_admin.role = TenantAccountRole.NORMAL
with (
flask_app_with_containers.test_request_context("/"),
patch("controllers.console.apikey._get_resource"),
):
with pytest.raises(Forbidden):
BaseApiKeyResource.delete(resource, "rid", "kid", "tenant-id", non_admin)
assert response.status_code == 403
class TestDatasetApiKeyListResource:
@@ -311,6 +310,49 @@ class TestDatasetApiKeyListResource:
assert listed[scoped_resp.json["id"]] != scoped_resp.json["token"]
def test_dataset_creation_routes_serialize_on_the_same_workspace(
setup_dataset: tuple[FlaskClient, dict[str, str], Dataset],
db_session_with_containers: Session,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""The workspace request must wait until the per-dataset count/insert commits."""
_, _, dataset = setup_dataset
workspace_id, dataset_id = dataset.tenant_id, dataset.id
db_session_with_containers.add_all(
ApiToken(tenant_id=workspace_id, type=ApiTokenType.DATASET, token=f"dataset-{uuid4()}") for _ in range(9)
)
db_session_with_containers.commit()
repository = DatasetApiKeyRepository(session_factory=sessionmaker(bind=db_session_with_containers.get_bind()))
first_counted, release_first, second_started = Event(), Event(), Event()
generate_key = ApiToken.generate_api_key
def pause_first_insert(prefix: str, n: int, *, session: Session) -> str:
if prefix == "ds-":
first_counted.set()
assert release_first.wait(timeout=10), "Timed out waiting to release the first insertion"
return generate_key(prefix, n, session=session)
def create_workspace_key() -> None:
second_started.set()
repository.create_workspace_key(workspace_id, (), max_keys=10, prefix="dataset-")
monkeypatch.setattr(ApiToken, "generate_api_key", pause_first_insert)
with ThreadPoolExecutor(max_workers=2) as pool:
first = pool.submit(repository.create_key, workspace_id, dataset_id, max_keys=10, prefix="ds-")
try:
assert first_counted.wait(timeout=5)
second = pool.submit(create_workspace_key)
assert second_started.wait(timeout=5)
with pytest.raises(TimeoutError):
second.result(timeout=0.5)
finally:
release_first.set()
first.result(timeout=5)
with pytest.raises(ApiKeyLimitExceededError):
second.result(timeout=5)
assert len(repository.list_workspace_keys(workspace_id)) == 10
class TestDatasetApiKeyResource:
"""Tests for DELETE /datasets/api-keys/<api_key_id>."""
api/tests/test_containers_integration_tests/services/test_dataset_service_document.py
+2 -1
View File
@@ -12,6 +12,7 @@ from werkzeug.exceptions import Forbidden, NotFound
from core.rag.index_processor.constant.index_type import IndexStructureType
from extensions.storage.storage_type import StorageType
from models import Account
from models.account import TenantAccountRole
from models.dataset import Dataset, Document
from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus
from models.model import UploadFile
@@ -119,7 +120,7 @@ def current_user_mock():
with patch("services.dataset_service.current_user", create_autospec(Account, instance=True)) as current_user:
current_user.id = str(uuid4())
current_user.current_tenant_id = str(uuid4())
current_user.current_role = None
current_user.current_role = TenantAccountRole.EDITOR
yield current_user
api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py
+3 -67
View File
@@ -2,7 +2,7 @@ from datetime import datetime
from inspect import getsource, unwrap
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import MagicMock, Mock, call
from unittest.mock import MagicMock, Mock
import pytest
from flask import Flask
@@ -24,8 +24,6 @@ from controllers.console.agent.composer import (
)
from controllers.console.agent.roster import (
AgentApiAccessApi,
AgentApiKeyApi,
AgentApiKeyListApi,
AgentApiStatusApi,
AgentApiStatusPayload,
AgentAppApi,
@@ -953,11 +951,10 @@ def test_agent_api_key_count_scopes_tenant_and_keeps_legacy_tokens(sqlite_sessio
assert roster_controller._agent_api_key_count(sqlite_session, app_model) == 2
def test_agent_api_status_and_key_routes_resolve_backing_app(
def test_agent_api_status_resolves_backing_app(
app: Flask, monkeypatch: pytest.MonkeyPatch, unbound_session: Session
) -> None:
agent_id = "00000000-0000-0000-0000-000000000001"
api_key_id = "00000000-0000-0000-0000-000000000002"
app_model = _app_detail_obj(
id="app-1",
tenant_id="tenant-1",
@@ -979,34 +976,6 @@ def test_agent_api_status_and_key_routes_resolve_backing_app(
monkeypatch.setattr(roster_controller, "AppService", FakeAppService)
def fake_get_api_key_list(self, resource_id: str, tenant_id: str, *, session: object):
captured["list_keys"] = {"session": session, "resource_id": resource_id, "tenant_id": tenant_id}
return roster_controller.ApiKeyList(data=[])
def fake_create_api_key(self, resource_id: str, tenant_id: str, *, session: object):
captured["create_key"] = {"session": session, "resource_id": resource_id, "tenant_id": tenant_id}
return ApiToken(id=api_key_id, type="app", token="app-test-token", last_used_at=None, created_at=None)
def fake_delete_api_key(
self,
resource_id: str,
key_id: str,
tenant_id: str,
current_user: object,
*,
session: object,
) -> None:
captured["delete_key"] = {
"session": session,
"resource_id": resource_id,
"api_key_id": key_id,
"tenant_id": tenant_id,
"current_user": current_user,
}
monkeypatch.setattr(AgentApiKeyListApi, "_get_api_key_list", fake_get_api_key_list)
monkeypatch.setattr(AgentApiKeyListApi, "_create_api_key", fake_create_api_key)
monkeypatch.setattr(AgentApiKeyApi, "_delete_api_key", fake_delete_api_key)
with app.test_request_context(
"/console/api/agent/00000000-0000-0000-0000-000000000001/api-enable", json={"enable_api": True}
):
@@ -1015,40 +984,7 @@ def test_agent_api_status_and_key_routes_resolve_backing_app(
)
assert enabled["enabled"] is True
assert captured["enable"] == {"app": app_model, "enable_api": True}
keys = unwrap(AgentApiKeyListApi.get)(AgentApiKeyListApi(), unbound_session, "tenant-1", agent_id)
assert keys == {"data": []}
assert captured["list_keys"] == {
"session": unbound_session,
"resource_id": "app-1",
"tenant_id": "tenant-1",
}
created, status = unwrap(AgentApiKeyListApi.post)(AgentApiKeyListApi(), unbound_session, "tenant-1", agent_id)
assert status == 201
assert created["id"] == api_key_id
assert created["token"] == "app-test-token"
assert captured["create_key"] == {
"session": unbound_session,
"resource_id": "app-1",
"tenant_id": "tenant-1",
}
current_user = _account(privileged=True)
deleted, delete_status = unwrap(AgentApiKeyApi.delete)(
AgentApiKeyApi(), unbound_session, "tenant-1", current_user, agent_id, api_key_id
)
assert (deleted, delete_status) == ("", 204)
assert captured["delete_key"] == {
"session": unbound_session,
"resource_id": "app-1",
"api_key_id": api_key_id,
"tenant_id": "tenant-1",
"current_user": current_user,
}
assert resolve_app.call_args_list == [
call(unbound_session, tenant_id="tenant-1", agent_id=agent_id),
call(unbound_session, tenant_id="tenant-1", agent_id=agent_id),
call(unbound_session, tenant_id="tenant-1", agent_id=agent_id),
call(unbound_session, tenant_id="tenant-1", agent_id=agent_id),
]
resolve_app.assert_called_once_with(unbound_session, tenant_id="tenant-1", agent_id=agent_id)
def test_agent_app_update_allows_empty_role(
api/tests/unit_tests/controllers/console/datasets/test_datasets.py
+78 -131
View File
@@ -8,14 +8,15 @@ from unittest.mock import ANY, MagicMock, PropertyMock, call, patch
import pytest
from flask import Flask
from sqlalchemy import event
from sqlalchemy.orm import Session
from sqlalchemy import event, select
from sqlalchemy.orm import Session, sessionmaker
from werkzeug.exceptions import BadRequest, Forbidden, NotFound
import services
from controllers.common.rbac import DatasetId, Workspace
from controllers.console import console_ns
from controllers.console.app.error import ProviderNotInitializeError
from controllers.console.datasets import datasets as datasets_controller
from controllers.console.datasets.datasets import (
DatasetApi,
DatasetApiBaseUrlApi,
@@ -46,12 +47,14 @@ from core.provider_manager import ProviderManager
from core.rag.datasource.vdb.vector_type import VectorType
from core.rag.index_processor.constant.index_type import IndexStructureType
from core.rag.retrieval.retrieval_methods import RetrievalMethod
from extensions.application_services.knowledge import build_dataset_api_key_service
from extensions.storage.storage_type import StorageType
from fields.dataset_fields import build_dataset_detail_prefetch
from models.account import Account, TenantAccountRole
from machinery.context import RequestContext
from models.account import Account, Tenant, TenantAccountRole
from models.dataset import AppDatasetJoin, Dataset, DatasetPermission, DatasetQuery, Document, DocumentSegment
from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus
from models.model import ApiToken, App, AppMode, IconType, UploadFile
from models.model import ApiToken, App, AppMode, DatasetApiTokenBinding, IconType, UploadFile
from services.dataset_ref_service import DatasetRef
from services.dataset_service import DatasetPermissionService, DatasetService
from services.enterprise import rbac_service as enterprise_rbac_service
@@ -1528,168 +1531,112 @@ class TestDatasetIndexingStatusApi(_UsesSQLiteSession):
class TestDatasetApiKeyApi(_UsesSQLiteSession):
context = RequestContext("request", None, "actor", "tenant-1")
@pytest.fixture(autouse=True)
def key_services(self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], monkeypatch):
tenant = Tenant(name="Workspace")
tenant.id = "tenant-1"
sqlite_session.add(tenant)
sqlite_session.commit()
service = build_dataset_api_key_service(database_client=sqlite_session_factory)
monkeypatch.setattr(
datasets_controller,
"application_services",
lambda: SimpleNamespace(dataset_api_keys=service),
)
def test_get_api_keys_success(self, app: Flask):
api = DatasetApiKeyApi()
method = unwrap(api.get)
mock_key_1 = ApiToken(
id="key-1",
type="dataset",
token="dataset-aaaa1111bbbb",
last_used_at=None,
created_at=None,
self.session.add_all(
[
ApiToken(id="key-1", tenant_id="tenant-1", type="dataset", token="dataset-aaaa1111bbbb"),
ApiToken(id="key-2", tenant_id="tenant-1", type="dataset", token="dataset-cccc2222dddd"),
]
)
mock_key_2 = ApiToken(
id="key-2",
type="dataset",
token="dataset-cccc2222dddd",
last_used_at=None,
created_at=None,
)
session = self.session
mock_key_1.tenant_id = "tenant-1"
mock_key_2.tenant_id = "tenant-1"
session.add_all([mock_key_1, mock_key_2])
session.flush()
self.session.commit()
with app.test_request_context("/"):
response = method(api, session, "tenant-1")
assert "data" in response
assert len(response["data"]) == 2
# reveal-once: the list returns masked tokens, never the full secret
assert response["data"][0]["id"] == "key-1"
assert response["data"][0]["token"] == "datas...bbbb"
assert response["data"][1]["id"] == "key-2"
assert response["data"][1]["token"] == "datas...dddd"
response = unwrap(api.get)(api, self.context)
assert [key["token"] for key in response["data"]] == ["datas...bbbb", "datas...dddd"]
def test_post_create_api_key_success(self, app: Flask):
api = DatasetApiKeyApi()
method = unwrap(api.post)
session = self.session
with (
app.test_request_context("/"),
patch.object(ApiToken, "generate_api_key", return_value="dataset-abc123") as generate_api_key,
):
response, status = method(api, session, "tenant-1")
with app.test_request_context("/"):
response, status = unwrap(api.post)(api, self.context)
assert status == 200
assert isinstance(response, dict)
assert response["token"] == "dataset-abc123"
assert response["token"].startswith("dataset-")
assert response["type"] == "dataset"
assert response["dataset_ids"] == []
assert response["created_at"] is not None
generate_api_key.assert_called_once_with("dataset-", 24, session=session)
assert session.get(ApiToken, response["id"]).token == "dataset-abc123"
stored = self.session.get(ApiToken, response["id"])
assert stored is not None
assert stored.token == response["token"]
def test_post_exceed_max_keys(self, app: Flask):
api = DatasetApiKeyApi()
method = unwrap(api.post)
session = self.session
session.add_all(
[
ApiToken(id=f"key-{index}", tenant_id="tenant-1", type="dataset", token=f"ds-{index}")
for index in range(10)
]
self.session.add_all(
[ApiToken(id=f"key-{i}", tenant_id="tenant-1", type="dataset", token=f"ds-{i}") for i in range(10)]
)
session.flush()
with app.test_request_context("/"):
with pytest.raises(BadRequest) as exc_info:
method(api, session, "tenant-1")
assert exc_info.value.code == 400
self.session.commit()
with app.test_request_context("/"), pytest.raises(BadRequest) as exc_info:
unwrap(api.post)(api, self.context)
assert vars(exc_info.value)["data"] == {
"message": "Cannot create more than 10 API keys for this resource type.",
"custom": "max_keys_exceeded",
}
def test_get_api_keys_include_bound_dataset_ids(self, app: Flask):
def test_scoped_key_persists_bindings_and_lists_masked_scope(self, app: Flask):
api = DatasetApiKeyApi()
method = unwrap(api.get)
mock_key = MagicMock(spec=ApiToken)
mock_key.id = "key-1"
mock_key.type = "dataset"
mock_key.token = "dataset-aaaa1111bbbb"
mock_key.last_used_at = None
mock_key.created_at = None
session = MagicMock()
session.scalars.return_value.all.return_value = [mock_key]
# Two binding rows for the same key -> the masked list surfaces both dataset ids.
session.execute.return_value.all.return_value = [("key-1", "ds-1"), ("key-1", "ds-2")]
with app.test_request_context("/"):
response = method(api, session, "tenant-1")
assert response["data"][0]["dataset_ids"] == ["ds-1", "ds-2"]
def test_post_create_scoped_key_persists_bindings(self, app: Flask):
api = DatasetApiKeyApi()
method = unwrap(api.post)
mock_token = MagicMock()
mock_token.id = "new-key-id"
mock_token.last_used_at = None
mock_token.created_at = datetime.datetime(2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC)
mock_api_token_cls = MagicMock()
mock_api_token_cls.return_value = mock_token
mock_api_token_cls.generate_api_key.return_value = "dataset-abc123"
session = MagicMock()
# Dataset-ownership validation returns exactly the requested ids (all belong to tenant).
session.scalars.return_value.all.return_value = ["ds-1", "ds-2"]
session.scalar.return_value = 1
with (
app.test_request_context("/", json={"dataset_ids": ["ds-1", "ds-2", "ds-1"]}),
patch("controllers.console.datasets.datasets.ApiToken", mock_api_token_cls),
):
response, status = method(api, session, "tenant-1")
self.session.add_all(
[
make_dataset(id="ds-1", tenant_id="tenant-1"),
make_dataset(id="ds-2", tenant_id="tenant-1"),
]
)
self.session.commit()
with app.test_request_context("/", json={"dataset_ids": ["ds-1", "ds-2", "ds-1"]}):
response, status = unwrap(api.post)(api, self.context)
listed = unwrap(api.get)(api, self.context)
assert status == 200
# Duplicates are collapsed and returned in the reveal-once response.
assert response["dataset_ids"] == ["ds-1", "ds-2"]
# One binding row is added per unique dataset id (plus the token itself).
added = [call.args[0] for call in session.add.call_args_list]
binding_dataset_ids = [obj.dataset_id for obj in added if obj.__class__.__name__ == "DatasetApiTokenBinding"]
assert binding_dataset_ids == ["ds-1", "ds-2"]
bindings = self.session.scalars(
select(DatasetApiTokenBinding).where(DatasetApiTokenBinding.api_token_id == response["id"])
).all()
assert {binding.dataset_id for binding in bindings} == {"ds-1", "ds-2"}
assert set(listed["data"][0]["dataset_ids"]) == {"ds-1", "ds-2"}
assert listed["data"][0]["token"] != response["token"]
def test_post_rejects_dataset_ids_from_another_tenant(self, app: Flask):
api = DatasetApiKeyApi()
method = unwrap(api.post)
session = MagicMock()
# None of the requested ids belong to this tenant.
session.scalars.return_value.all.return_value = []
with app.test_request_context("/", json={"dataset_ids": ["foreign-ds"]}):
with pytest.raises(BadRequest) as exc_info:
method(api, session, "tenant-1")
assert exc_info.value.code == 400
self.session.add(make_dataset(id="foreign-ds", tenant_id="foreign"))
self.session.commit()
with app.test_request_context("/", json={"dataset_ids": ["foreign-ds"]}), pytest.raises(BadRequest) as exc_info:
unwrap(api.post)(api, self.context)
assert "Unknown knowledge base id(s)" in vars(exc_info.value)["data"]["message"]
assert self.session.scalar(select(ApiToken)) is None
def test_post_rejects_non_list_dataset_ids(self, app: Flask):
@pytest.mark.parametrize("dataset_ids", ["not-a-list", [1]])
def test_post_rejects_invalid_dataset_ids(self, app: Flask, dataset_ids):
api = DatasetApiKeyApi()
method = unwrap(api.post)
session = MagicMock()
with app.test_request_context("/", json={"dataset_ids": "not-a-list"}):
with pytest.raises(BadRequest) as exc_info:
method(api, session, "tenant-1")
assert exc_info.value.code == 400
with app.test_request_context("/", json={"dataset_ids": dataset_ids}), pytest.raises(BadRequest) as exc_info:
unwrap(api.post)(api, self.context)
assert vars(exc_info.value)["data"]["message"] == "dataset_ids must be a list of strings."
class TestDatasetApiDeleteApi(_UsesSQLiteSession):
def test_delete_success(self, app: Flask):
api = DatasetApiDeleteApi()
method = unwrap(api.delete)
session = self.session
key = ApiToken(id="api-key-id", tenant_id="tenant-1", type="dataset", token="dataset-secret")
session.add(key)
session.flush()
with (
app.test_request_context("/"),
patch("controllers.console.datasets.datasets.ApiTokenCache.delete") as delete_cache,
):
response, status = method(api, session, "tenant-1", "api-key-id")
assert status == 204
assert response == ""
delete_cache.assert_called_once()
session.flush()
assert session.get(ApiToken, "api-key-id") is None
self.session.add(ApiToken(id="api-key-id", tenant_id="tenant-1", type="dataset", token="dataset-secret"))
self.session.commit()
with app.test_request_context("/"), patch("services.api_token_service.ApiTokenCache.delete") as invalidate:
result = unwrap(api.delete)(api, self.context, "api-key-id")
assert result == ("", 204)
self.session.expire_all()
assert self.session.get(ApiToken, "api-key-id") is None
invalidate.assert_called_once_with("dataset-secret", "dataset")
def test_delete_key_not_found(self, app: Flask):
api = DatasetApiDeleteApi()
method = unwrap(api.delete)
session = self.session
with app.test_request_context("/"):
with pytest.raises(NotFound):
method(api, session, "tenant-1", "api-key-id")
with app.test_request_context("/"), pytest.raises(NotFound):
unwrap(api.delete)(api, self.context, "api-key-id")
class TestDatasetEnableApiApi(_UsesSQLiteSession):
api/tests/unit_tests/controllers/console/test_apikey.py
+287 -308
View File
@@ -1,338 +1,317 @@
from __future__ import annotations
"""API key admission, delegation, error and response contracts."""
import inspect
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import cast
from unittest.mock import MagicMock, patch
from unittest.mock import Mock
from uuid import UUID
import pytest
from flask import Flask
from sqlalchemy import event, select
from sqlalchemy.orm import Session
from werkzeug.exceptions import BadRequest, Forbidden, NotFound
from flask_restx import Api
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from controllers.console.agent.roster import AgentApiKeyListApi
from controllers.console.apikey import (
AppApiKeyListResource,
BaseApiKeyListResource,
BaseApiKeyResource,
DatasetApiKeyListResource,
)
from controllers.console.datasets.datasets import DatasetApiKeyApi
from core.rbac import RBACPermission, RBACResourceScope
from controllers.console import apikey, flask_admission, wraps
from controllers.console.agent import roster
from controllers.console.datasets import datasets
from core.rbac import RBACPermission
from enums import DeploymentEdition
from models import Account
from models.account import AccountStatus, TenantAccountRole
from models.enums import ApiTokenType
from models.model import ApiToken, App, AppMode, IconType
from services.agent.errors import AgentAccessNotReadyError
from extensions.application_services.knowledge import build_dataset_api_key_service
from extensions.ext_application_services import ApplicationServices
from libs.login import AccountWithTenant
from machinery.context import RequestContext
from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset
from models.model import ApiToken
from services.app.api_key_service import AppApiKeyNotReadyError
from services.auth.api_key_contracts import (
ApiKeyLimitExceededError,
ApiKeyNotFoundError,
ApiKeyRecord,
ApiKeyResourceNotFoundError,
)
RESOURCE_ID = UUID("00000000-0000-0000-0000-000000000001")
KEY_ID = UUID("00000000-0000-0000-0000-000000000002")
def _make_list_resource() -> BaseApiKeyListResource:
resource = BaseApiKeyListResource()
resource.resource_type = ApiTokenType.APP
resource.resource_model = App
resource.resource_id_field = "app_id"
resource.token_prefix = "app-"
return resource
@dataclass
class RecordingKeys:
calls: list[tuple[str, RequestContext, str, str, str | None]] = field(default_factory=list)
error: Exception | None = None
def _make_key_resource() -> BaseApiKeyResource:
resource = BaseApiKeyResource()
resource.resource_type = ApiTokenType.APP
resource.resource_model = App
resource.resource_id_field = "app_id"
return resource
def _make_account(role: TenantAccountRole) -> Account:
account = Account(
name="Test User",
email=f"{role.value}@example.com",
status=AccountStatus.ACTIVE,
)
account.id = f"{role.value}-user"
account.role = role
return account
def _persist_app(session: Session, *, mode: AppMode = AppMode.CHAT, app_id: str = "app-1") -> App:
app = App(
id=app_id,
tenant_id="tenant-1",
name="API key app",
mode=mode,
icon_type=IconType.EMOJI,
icon="chat",
icon_background="#ffffff",
enable_site=False,
enable_api=True,
)
session.add(app)
session.flush()
return app
def test_list_api_keys_uses_injected_session_and_tenant_id(sqlite_session: Session) -> None:
resource = _make_list_resource()
raw_get = cast(
Callable[[BaseApiKeyListResource, object, str, str], dict[str, object]],
inspect.unwrap(BaseApiKeyListResource.get),
)
session = sqlite_session
_persist_app(session)
api_key = ApiToken(
type=ApiTokenType.APP,
token="app-token",
app_id="app-1",
tenant_id="tenant-1",
)
api_key.id = "key-1"
session.add(api_key)
session.add(
ApiToken(
type=ApiTokenType.APP,
token="foreign-app-token",
app_id="app-1",
tenant_id="tenant-2",
def _record(
self,
operation: str,
context: RequestContext,
kind: str,
resource_id: str,
key_id: str | None = None,
) -> ApiKeyRecord:
self.calls.append((operation, context, kind, resource_id, key_id))
if self.error:
raise self.error
return ApiKeyRecord(
id=str(KEY_ID),
type="dataset" if kind == "dataset" else "app",
token="app-secret-token",
created_at=datetime(2026, 1, 1, tzinfo=UTC),
dataset_ids=(resource_id,) if kind == "dataset" else (),
)
)
legacy_api_key = ApiToken(type=ApiTokenType.APP, token="legacy-app-token", app_id="app-1", tenant_id=None)
session.add(legacy_api_key)
session.commit()
result = raw_get(resource, session, "app-1", "tenant-1")
data = cast(list[dict[str, object]], result["data"])
def list_keys(self, context: RequestContext, app_id: str) -> tuple[ApiKeyRecord, ...]:
return (self._record("list", context, "app", app_id),)
assert {item["token"] for item in data} == {"app-token", "legacy-app-token"}
def create_key(self, context: RequestContext, app_id: str) -> ApiKeyRecord:
return self._record("create", context, "app", app_id)
def delete_key(self, context: RequestContext, app_id: str, key_id: str) -> None:
self._record("delete", context, "app", app_id, key_id)
def list_agent_keys(self, context: RequestContext, agent_id: str) -> tuple[ApiKeyRecord, ...]:
return (self._record("list", context, "agent", agent_id),)
def create_agent_key(self, context: RequestContext, agent_id: str) -> ApiKeyRecord:
return self._record("create", context, "agent", agent_id)
def delete_agent_key(self, context: RequestContext, agent_id: str, key_id: str) -> None:
self._record("delete", context, "agent", agent_id, key_id)
def test_create_api_key_uses_injected_session_and_tenant_id(sqlite_session: Session) -> None:
resource = _make_list_resource()
raw_post = cast(
Callable[[BaseApiKeyListResource, object, str, str], tuple[dict[str, object], int]],
inspect.unwrap(BaseApiKeyListResource.post),
)
session = sqlite_session
_persist_app(session)
session.add_all(
[
ApiToken(type=ApiTokenType.APP, token=f"foreign-token-{index}", app_id="app-1", tenant_id="tenant-2")
for index in range(resource.max_keys)
]
)
session.commit()
commits: list[str] = []
event.listen(session, "after_commit", lambda _session: commits.append("commit"))
@dataclass
class RecordingDatasetKeys:
recorder: RecordingKeys
with patch(
"controllers.console.apikey.ApiToken.generate_api_key", return_value="app-generated-token"
) as generate_api_key:
result, status = raw_post(resource, session, "app-1", "tenant-1")
def list_keys(self, context: RequestContext, dataset_id: str) -> tuple[ApiKeyRecord, ...]:
return (self.recorder._record("list", context, "dataset", dataset_id),)
assert status == 201
assert result["token"] == "app-generated-token"
api_token = session.scalar(select(ApiToken).where(ApiToken.token == "app-generated-token"))
assert api_token is not None
assert api_token.app_id == "app-1"
assert api_token.tenant_id == "tenant-1"
assert api_token.type == ApiTokenType.APP
generate_api_key.assert_called_once_with("app-", 24, session=session)
assert commits == ["commit"]
def create_key(self, context: RequestContext, dataset_id: str) -> ApiKeyRecord:
return self.recorder._record("create", context, "dataset", dataset_id)
def delete_key(self, context: RequestContext, dataset_id: str, key_id: str) -> None:
self.recorder._record("delete", context, "dataset", dataset_id, key_id)
def list_workspace_keys(self, context: RequestContext) -> tuple[ApiKeyRecord, ...]:
return (self.recorder._record("list", context, "dataset", "workspace"),)
def create_workspace_key(self, context: RequestContext, _dataset_ids: tuple[str, ...]) -> ApiKeyRecord:
return self.recorder._record("create", context, "dataset", "workspace")
def delete_workspace_key(self, context: RequestContext, key_id: str) -> None:
self.recorder._record("delete", context, "dataset", "workspace", key_id)
def test_create_api_key_counts_legacy_tokens(sqlite_session: Session) -> None:
resource = _make_list_resource()
_persist_app(sqlite_session)
@dataclass
class ApiKeyTestServices:
app_api_keys: RecordingKeys
dataset_api_keys: RecordingDatasetKeys
type KeysApp = tuple[Flask, RecordingKeys, Account]
@pytest.fixture
def keys_app(monkeypatch: pytest.MonkeyPatch, config_overrides: Callable[..., None]) -> KeysApp:
config_overrides(DEPLOYMENT_EDITION=DeploymentEdition.CLOUD, LOGIN_DISABLED=True, RBAC_ENABLED=False)
account = Account(name="Owner", email="owner@example.com", status=AccountStatus.ACTIVE)
account.id = "actor"
account.role = TenantAccountRole.OWNER
def current_account() -> AccountWithTenant:
return AccountWithTenant(account, "workspace")
monkeypatch.setattr(flask_admission, "current_account_with_tenant", current_account)
monkeypatch.setattr(wraps, "current_account_with_tenant", current_account)
monkeypatch.setattr(flask_admission, "get_request_id", lambda: "request-id")
monkeypatch.setattr(flask_admission, "get_trace_id", lambda: None)
keys = RecordingKeys()
services = cast(ApplicationServices, ApiKeyTestServices(keys, RecordingDatasetKeys(keys)))
monkeypatch.setattr(apikey, "application_services", lambda: services)
monkeypatch.setattr(roster, "application_services", lambda: services)
monkeypatch.setattr(datasets, "application_services", lambda: services)
app = Flask(__name__)
api = Api(app)
api.add_resource(apikey.AppApiKeyListResource, "/app/<uuid:resource_id>")
api.add_resource(apikey.AppApiKeyResource, "/app/<uuid:resource_id>/<uuid:api_key_id>")
api.add_resource(apikey.DatasetApiKeyListResource, "/dataset/<uuid:resource_id>")
api.add_resource(apikey.DatasetApiKeyResource, "/dataset/<uuid:resource_id>/<uuid:api_key_id>")
api.add_resource(roster.AgentApiKeyListApi, "/agent/<uuid:agent_id>")
api.add_resource(roster.AgentApiKeyApi, "/agent/<uuid:agent_id>/<uuid:api_key_id>")
api.add_resource(datasets.DatasetApiKeyApi, "/workspace-keys")
api.add_resource(datasets.DatasetApiDeleteApi, "/workspace-keys/<uuid:api_key_id>")
return app, keys, account
@pytest.mark.parametrize("kind", ["app", "dataset", "agent"])
def test_key_routes_forward_stable_context_and_serialize(keys_app: KeysApp, kind: str) -> None:
app, keys, _ = keys_app
client = app.test_client()
path = f"/{kind}/{RESOURCE_ID}"
headers = {"X-Trace-Id": "trace-id"}
listed = client.get(path, headers=headers)
created = client.post(path, headers=headers)
deleted = client.delete(f"{path}/{KEY_ID}", headers=headers)
assert listed.status_code == 200
assert listed.json is not None
assert listed.json["data"][0]["token"] == ("app-s...oken" if kind == "dataset" else "app-secret-token")
assert listed.json["data"][0]["dataset_ids"] == ([str(RESOURCE_ID)] if kind == "dataset" else [])
assert created.status_code == 201
assert created.json is not None
assert created.json["token"] == "app-secret-token"
assert created.json["created_at"] == 1767225600
assert deleted.status_code == 204
assert deleted.data == b""
context = RequestContext("request-id", "trace-id", "actor", "workspace")
assert keys.calls == [
("list", context, kind, str(RESOURCE_ID), None),
("create", context, kind, str(RESOURCE_ID), None),
("delete", context, kind, str(RESOURCE_ID), str(KEY_ID)),
]
@pytest.mark.parametrize("kind", ["app", "dataset", "agent"])
@pytest.mark.parametrize("role", list(TenantAccountRole))
@pytest.mark.parametrize("method", ["GET", "POST", "DELETE"])
def test_admission_preserves_role_policy(keys_app: KeysApp, kind: str, role: TenantAccountRole, method: str) -> None:
app, keys, account = keys_app
account.role = role
path = f"/{kind}/{RESOURCE_ID}" + (f"/{KEY_ID}" if method == "DELETE" else "")
response = app.test_client().open(path, method=method)
allowed = role in (apikey.API_KEY_DELETE_ROLES if method == "DELETE" else apikey.API_KEY_EDIT_ROLES)
assert response.status_code == ({"GET": 200, "POST": 201, "DELETE": 204}[method] if allowed else 403)
assert bool(keys.calls) is allowed
@pytest.mark.parametrize(
("error", "status", "message"),
[
(ApiKeyResourceNotFoundError("App not found."), 404, "App not found."),
(ApiKeyNotFoundError(), 404, "API key not found"),
(ApiKeyLimitExceededError(10), 400, "Cannot create more than 10 API keys for this resource type."),
(AppApiKeyNotReadyError(), 409, "Publish the Agent before enabling Web App or API access."),
],
)
def test_errors_keep_http_contract(keys_app: KeysApp, error: Exception, status: int, message: str) -> None:
app, keys, _ = keys_app
keys.error = error
response = app.test_client().post(f"/app/{RESOURCE_ID}")
assert response.status_code == status
assert response.json is not None
assert response.json["message"].startswith(message)
if isinstance(error, ApiKeyLimitExceededError):
assert response.json["custom"] == "max_keys_exceeded"
@pytest.mark.parametrize(
("kind", "permission"),
[
("app", RBACPermission.APP_RELEASE_AND_VERSION),
("dataset", RBACPermission.DATASET_API_KEY_MANAGE),
("agent", RBACPermission.AGENT_ACCESS_POINT_VIEW),
],
)
@pytest.mark.parametrize("allowed", [False, True])
def test_rbac_admission_controls_service_access(
keys_app: KeysApp,
config_overrides: Callable[..., None],
monkeypatch: pytest.MonkeyPatch,
kind: str,
permission: RBACPermission,
allowed: bool,
) -> None:
from controllers.common.rbac import checks, locators
app, keys, account = keys_app
account.role = TenantAccountRole.NORMAL
config_overrides(RBAC_ENABLED=True)
monkeypatch.setattr(locators, "agent_binding", lambda *_args, **_kwargs: None)
monkeypatch.setattr(locators.PlainApp, "owner_id", lambda *_args, **_kwargs: None)
monkeypatch.setattr(locators.DatasetId, "owner_id", lambda *_args, **_kwargs: None)
monkeypatch.setattr(locators.AgentId, "owner_id", lambda *_args, **_kwargs: None)
check = Mock(return_value=allowed)
monkeypatch.setattr(checks.RBACService.CheckAccess, "check", check)
response = app.test_client().get(f"/{kind}/{RESOURCE_ID}")
assert response.status_code == (200 if allowed else 403)
assert check.call_args.kwargs["scene"] == permission
assert bool(keys.calls) is allowed
def test_uninitialized_account_is_rejected(keys_app: KeysApp) -> None:
app, keys, account = keys_app
account.status = AccountStatus.UNINITIALIZED
assert app.test_client().post(f"/app/{RESOURCE_ID}").status_code == 400
assert keys.calls == []
@pytest.mark.parametrize("role", list(TenantAccountRole))
@pytest.mark.parametrize("method", ["GET", "POST", "DELETE"])
def test_workspace_key_admission_preserves_admin_role_policy(
keys_app: KeysApp, role: TenantAccountRole, method: str
) -> None:
app, keys, account = keys_app
account.role = role
path = "/workspace-keys" + (f"/{KEY_ID}" if method == "DELETE" else "")
response = app.test_client().open(path, method=method)
allowed = role in apikey.API_KEY_DELETE_ROLES
assert response.status_code == ({"GET": 200, "POST": 200, "DELETE": 204}[method] if allowed else 403)
assert bool(keys.calls) is allowed
@pytest.fixture
def persisted_keys_app(
keys_app: KeysApp,
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
monkeypatch: pytest.MonkeyPatch,
) -> KeysApp:
app, keys, account = keys_app
assert account.role is not None
tenant = Tenant(name="Workspace")
tenant.id = "workspace"
sqlite_session.add_all(
[
ApiToken(type=ApiTokenType.APP, token=f"legacy-token-{index}", app_id="app-1", tenant_id=None)
for index in range(resource.max_keys)
tenant,
TenantAccountJoin(tenant_id="workspace", account_id=account.id, role=account.role),
Dataset(id=str(RESOURCE_ID), tenant_id="workspace", name="Private", created_by="other", maintainer="other"),
]
)
sqlite_session.commit()
with pytest.raises(BadRequest):
resource._create_api_key("app-1", "tenant-1", session=sqlite_session)
service = build_dataset_api_key_service(database_client=sqlite_session_factory)
services = Mock(dataset_api_keys=service)
monkeypatch.setattr(apikey, "application_services", lambda: services)
monkeypatch.setattr(datasets, "application_services", lambda: services)
return app, keys, account
def test_create_agent_api_key_requires_published_access(sqlite_session: Session) -> None:
resource = _make_list_resource()
session = sqlite_session
app = _persist_app(session, mode=AppMode.AGENT)
with patch(
"controllers.console.apikey.AppService.ensure_agent_app_access_ready",
side_effect=AgentAccessNotReadyError(),
) as ensure_access_ready:
with pytest.raises(AgentAccessNotReadyError):
resource._create_api_key("app-1", "tenant-1", session=session)
ensure_access_ready.assert_called_once_with(app, session=session)
assert session.scalar(select(ApiToken)) is None
def test_private_dataset_rejects_editor_without_creating_token(
persisted_keys_app: KeysApp, sqlite_session: Session
) -> None:
app, _, account = persisted_keys_app
account.role = TenantAccountRole.EDITOR
member = sqlite_session.scalar(select(TenantAccountJoin))
assert member is not None
member.role = account.role
sqlite_session.commit()
response = app.test_client().post(f"/dataset/{RESOURCE_ID}")
assert response.status_code == 403
assert sqlite_session.scalar(select(ApiToken)) is None
def test_delete_api_key_rejects_non_admin_account(sqlite_session: Session) -> None:
resource = _make_key_resource()
raw_delete = cast(
Callable[[BaseApiKeyResource, object, str, str, str, Account], tuple[str, int]],
inspect.unwrap(BaseApiKeyResource.delete),
)
session = sqlite_session
_persist_app(session)
with pytest.raises(Forbidden):
raw_delete(
resource,
session,
"app-1",
"key-1",
"tenant-1",
_make_account(TenantAccountRole.NORMAL),
)
def test_delete_api_key_uses_injected_session_user_and_tenant(sqlite_session: Session) -> None:
resource = _make_key_resource()
raw_delete = cast(
Callable[[BaseApiKeyResource, object, str, str, str, Account], tuple[str, int]],
inspect.unwrap(BaseApiKeyResource.delete),
)
session = sqlite_session
_persist_app(session)
api_key = ApiToken(type=ApiTokenType.APP, token="app-token", app_id="app-1", tenant_id=None)
api_key.id = "key-1"
session.add(api_key)
session.commit()
commits: list[str] = []
event.listen(session, "after_commit", lambda _session: commits.append("commit"))
with patch("controllers.console.apikey.ApiTokenCache.delete") as delete_cache:
result, status = raw_delete(
resource,
session,
"app-1",
"key-1",
"tenant-1",
_make_account(TenantAccountRole.OWNER),
)
delete_cache.assert_called_once_with("app-token", ApiTokenType.APP)
assert session.get(ApiToken, "key-1") is None
assert commits == ["commit"]
assert result == ""
assert status == 204
def test_delete_api_key_rejects_foreign_tenant_token(sqlite_session: Session) -> None:
resource = _make_key_resource()
session = sqlite_session
_persist_app(session)
api_key = ApiToken(type=ApiTokenType.APP, token="foreign-token", app_id="app-1", tenant_id="tenant-2")
api_key.id = "key-1"
session.add(api_key)
session.commit()
with patch("controllers.console.apikey.ApiTokenCache.delete") as delete_cache:
with pytest.raises(NotFound):
resource._delete_api_key(
"app-1",
"key-1",
"tenant-1",
_make_account(TenantAccountRole.OWNER),
session=session,
)
delete_cache.assert_not_called()
assert session.get(ApiToken, "key-1") is api_key
def test_api_key_lists_require_matching_rbac_permission(config_overrides: Callable[..., None]) -> None:
config_overrides(
DEPLOYMENT_EDITION=DeploymentEdition.CLOUD,
LOGIN_DISABLED=True,
RBAC_ENABLED=True,
)
app = Flask(__name__)
account = _make_account(TenantAccountRole.OWNER)
api_id = UUID("00000000-0000-0000-0000-000000000001")
cases = [
(
lambda: AppApiKeyListResource().get(resource_id=api_id),
{
"scene": RBACPermission.APP_RELEASE_AND_VERSION,
"resource_type": RBACResourceScope.APP,
"resource_id": str(api_id),
},
),
(
lambda: DatasetApiKeyApi().get(),
{"scene": RBACPermission.DATASET_API_KEY_MANAGE, "resource_type": None, "resource_id": None},
),
(
lambda: DatasetApiKeyListResource().get(resource_id=api_id),
{
"scene": RBACPermission.DATASET_API_KEY_MANAGE,
"resource_type": RBACResourceScope.DATASET,
"resource_id": str(api_id),
},
),
]
with (
app.test_request_context("/"),
patch("controllers.console.wraps.current_account_with_tenant", return_value=(account, "tenant-1")),
patch("controllers.common.wraps.current_account_with_tenant", return_value=(account, "tenant-1")),
patch("controllers.common.rbac.locators.agent_binding", return_value=None),
patch("controllers.common.rbac.locators.PlainApp.owner_id", return_value=None),
patch("controllers.common.rbac.locators.DatasetId.owner_id", return_value=None),
patch.object(BaseApiKeyListResource, "_get_api_key_list") as get_api_key_list,
):
for invoke, expected_kwargs in cases:
with patch(
"controllers.common.rbac.checks.RBACService.CheckAccess.check", return_value=False
) as check_access:
with pytest.raises(Forbidden):
invoke()
check_access.assert_called_once_with(
"tenant-1",
account.id,
scene=expected_kwargs["scene"],
resource_type=expected_kwargs["resource_type"],
resource_id=expected_kwargs["resource_id"],
)
get_api_key_list.assert_not_called()
def test_api_key_lists_reject_legacy_read_only_members(config_overrides: Callable[..., None]) -> None:
config_overrides(
DEPLOYMENT_EDITION=DeploymentEdition.CLOUD,
LOGIN_DISABLED=True,
RBAC_ENABLED=False,
)
app = Flask(__name__)
account = _make_account(TenantAccountRole.NORMAL)
api_id = UUID("00000000-0000-0000-0000-000000000001")
current_user = MagicMock()
current_user._get_current_object.return_value = account
current_user.has_edit_permission = False
with (
app.test_request_context("/"),
patch("libs.login.current_user", current_user),
patch("controllers.console.wraps.current_account_with_tenant", return_value=(account, "tenant-1")),
patch.object(BaseApiKeyListResource, "_get_api_key_list") as get_api_key_list,
):
for invoke in (
lambda: AppApiKeyListResource().get(resource_id=api_id),
lambda: AgentApiKeyListApi().get(agent_id=api_id),
lambda: DatasetApiKeyApi().get(),
lambda: DatasetApiKeyListResource().get(resource_id=api_id),
):
with pytest.raises(Forbidden):
invoke()
get_api_key_list.assert_not_called()
def test_both_http_routes_enforce_the_same_key_limit(persisted_keys_app: KeysApp, sqlite_session: Session) -> None:
app, _, _ = persisted_keys_app
client = app.test_client()
for _ in range(5):
assert client.post("/workspace-keys").status_code == 200
assert client.post(f"/dataset/{RESOURCE_ID}").status_code == 201
for path in ("/workspace-keys", f"/dataset/{RESOURCE_ID}"):
response = client.post(path)
assert response.status_code == 400
assert response.json is not None
assert response.json["custom"] == "max_keys_exceeded"
assert len(sqlite_session.scalars(select(ApiToken)).all()) == 10
api/tests/unit_tests/extensions/test_ext_application_services.py
+33
View File
@@ -36,6 +36,7 @@ from repositories.app_site_command_repository import AppSiteCommandRepository
from repositories.app_statistic_query_repository import AppStatisticQueryRepository
from repositories.app_tracing_config_repository import SQLAlchemyAppTracingConfigRepository
from repositories.human_input_file_upload_repository import SQLAlchemyHumanInputFileUploadRepository
from repositories.installed_app_repository import SQLAlchemyInstalledAppRepository
from repositories.message_file_preview_repository import MessageFilePreviewQueryRepository
from repositories.plugin_file_upload_repository import SQLAlchemyPluginFileUploadOwnerRepository
from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository
@@ -66,6 +67,7 @@ from services.account_oauth_adapters import (
DeploymentOAuthPolicyGateway,
RedisOAuthAccountClaimLock,
)
from services.app.api_key_service import AppApiKeyService
from services.app_generate_service import AppGenerateService
from services.app_preview_query_service import AppPreviewRef, AppPreviewUnavailableError
from services.app_scoped_end_user_query_service import AppScopedEndUserQueryService
@@ -84,7 +86,11 @@ from services.file_service import FileService
from services.human_input_file_upload_service import HumanInputFileUploadService
from services.init_validation_service import InvalidInitializationPasswordError
from services.installed_app_access_service import InstalledAppAccessDeniedError, InstalledAppRef
from services.installed_app_generation_adapters import AppGenerateServiceRuntime as InstalledAppGenerateServiceRuntime
from services.installed_app_generation_service import InstalledAppGenerationService
from services.knowledge.api_key_service import DatasetApiKeyService
from services.message_file_preview_service import MessageFilePreviewService
from services.oauth_device_application_service import OAuthDeviceApplicationService
from services.partner_tenant_binding_service import PartnerTenantBindingService
from services.plugin_file_upload_gateway import ToolFilePluginUploadGateway
from services.plugin_file_upload_service import PluginFileUploadService
@@ -177,6 +183,33 @@ def test_init_app_registers_services_for_the_current_app(
assert isinstance(services.workflow_statistics, WorkflowStatisticQueryService)
def test_build_application_services_preserves_composed_boundaries(
sqlite_session_factory: sessionmaker[Session],
) -> None:
redis = MagicMock(spec=RedisClientWrapper)
services = ext_application_services.build_application_services(
database_client=sqlite_session_factory,
deployment_edition=DeploymentEdition.COMMUNITY,
initialization_password="",
redis=redis,
)
assert isinstance(services.app_api_keys, AppApiKeyService)
assert isinstance(services.dataset_api_keys, DatasetApiKeyService)
assert isinstance(services.oauth_device, OAuthDeviceApplicationService)
assert redis.register_script.call_count == 3
assert isinstance(services.installed_app_generation, InstalledAppGenerationService)
installed_apps = services.installed_app_access._installed_apps
assert isinstance(installed_apps, SQLAlchemyInstalledAppRepository)
assert services.installed_app_generation._usage is installed_apps
assert services.installed_app_generation._app_definitions is services.app_definitions
runtime = services.installed_app_generation._runtime
assert isinstance(runtime, InstalledAppGenerateServiceRuntime)
assert runtime._session_factory is sqlite_session_factory
@pytest.mark.parametrize(
("deployment_edition", "setup_completed"),
[
api/tests/unit_tests/repositories/app/test_api_key_repository.py
+85
View File
@@ -0,0 +1,85 @@
"""Real persistence checks for the resource API key boundary."""
import pytest
from sqlalchemy.orm import Session, sessionmaker
from models.enums import ApiTokenType
from models.model import ApiToken, App, AppMode
from repositories.app.api_key_repository import AppApiKeyRepository
from services.auth.api_key_contracts import (
ApiKeyLimitExceededError,
ApiKeyNotFoundError,
ApiKeyResourceNotFoundError,
)
@pytest.fixture
def repository(sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]) -> AppApiKeyRepository:
sqlite_session.add_all(
[
App(id="app", tenant_id="tenant", name="App", mode=AppMode.CHAT, enable_site=False, enable_api=True),
App(
id="other-app", tenant_id="tenant", name="Other", mode=AppMode.CHAT, enable_site=False, enable_api=True
),
]
)
sqlite_session.commit()
return AppApiKeyRepository(session_factory=sqlite_session_factory)
def test_app_key_list_and_delete_preserve_owner_chain(repository: AppApiKeyRepository, sqlite_session: Session) -> None:
sqlite_session.add_all(
[
ApiToken(id="current", type=ApiTokenType.APP, token="current", tenant_id="tenant", app_id="app"),
ApiToken(id="historical", type=ApiTokenType.APP, token="historical", tenant_id=None, app_id="app"),
ApiToken(id="foreign", type=ApiTokenType.APP, token="foreign", tenant_id="foreign", app_id="app"),
ApiToken(id="other", type=ApiTokenType.APP, token="other", tenant_id="tenant", app_id="other-app"),
ApiToken(id="wrong-type", type=ApiTokenType.DATASET, token="dataset", tenant_id="tenant", app_id="app"),
]
)
sqlite_session.commit()
assert {key.id for key in repository.list_keys("tenant", "app")} == {"current", "historical"}
for key_id in ["foreign", "other", "wrong-type", "missing"]:
with pytest.raises(ApiKeyNotFoundError):
repository.delete_key("tenant", "app", key_id)
deleted = repository.delete_key("tenant", "app", "historical")
assert deleted.token == "historical"
assert sqlite_session.get(ApiToken, "historical") is None
assert sqlite_session.get(ApiToken, "foreign") is not None
@pytest.mark.parametrize("operation", ["list", "create", "delete"])
def test_resource_must_belong_to_workspace(repository: AppApiKeyRepository, operation: str) -> None:
def invoke() -> object:
if operation == "list":
return repository.list_keys("foreign", "app")
if operation == "create":
return repository.create_key("foreign", "app", max_keys=10, prefix="test-")
return repository.delete_key("foreign", "app", "missing")
with pytest.raises(ApiKeyResourceNotFoundError):
invoke()
def test_creation_counts_historical_keys_but_not_foreign_keys(
repository: AppApiKeyRepository,
sqlite_session: Session,
) -> None:
sqlite_session.add_all(
[ApiToken(type=ApiTokenType.APP, token=f"foreign-{i}", tenant_id="foreign", app_id="app") for i in range(10)]
)
sqlite_session.add_all(
[ApiToken(type=ApiTokenType.APP, token=f"historical-{i}", tenant_id=None, app_id="app") for i in range(9)]
)
sqlite_session.commit()
created = repository.create_key("tenant", "app", max_keys=10, prefix="app-")
assert created.token.startswith("app-")
assert len(created.token) == 28
assert created.created_at is not None
with pytest.raises(ApiKeyLimitExceededError):
repository.create_key("tenant", "app", max_keys=10, prefix="app-")
stored = sqlite_session.get(ApiToken, created.id)
assert stored is not None
assert (stored.tenant_id, stored.app_id, stored.type) == ("tenant", "app", ApiTokenType.APP)
assert len(repository.list_keys("tenant", "app")) == 10
api/tests/unit_tests/repositories/knowledge/__init__.py
View File
Whitespace-only changes.
api/tests/unit_tests/services/test_dataset_api_key_service.py → api/tests/unit_tests/repositories/knowledge/test_dataset_api_key_bindings.py
+5 -5
View File
@@ -1,4 +1,4 @@
"""Unit tests for dataset API key scoping helpers."""
"""Persistence tests for dataset API key scoping and deletion."""
import uuid
@@ -7,7 +7,7 @@ from sqlalchemy.orm import Session
from models.enums import ApiTokenType
from models.model import ApiToken, DatasetApiTokenBinding
from services import dataset_api_key_service
from repositories.knowledge import dataset_api_key_bindings
def _dataset_token(session: Session) -> ApiToken:
@@ -28,7 +28,7 @@ class TestDeleteKeysScopedOnlyTo:
sqlite_session.add(DatasetApiTokenBinding(api_token_id=token.id, dataset_id=dataset_id))
sqlite_session.commit()
deleted = dataset_api_key_service.delete_keys_scoped_only_to(sqlite_session, dataset_id)
deleted = dataset_api_key_bindings.delete_keys_scoped_only_to(sqlite_session, dataset_id)
sqlite_session.commit()
assert deleted == [token.id]
@@ -42,7 +42,7 @@ class TestDeleteKeysScopedOnlyTo:
sqlite_session.add(DatasetApiTokenBinding(api_token_id=token.id, dataset_id=other_dataset_id))
sqlite_session.commit()
deleted = dataset_api_key_service.delete_keys_scoped_only_to(sqlite_session, dataset_id)
deleted = dataset_api_key_bindings.delete_keys_scoped_only_to(sqlite_session, dataset_id)
sqlite_session.commit()
# The key is bound elsewhere, so it survives and stays scoped.
@@ -55,7 +55,7 @@ class TestDeleteKeysScopedOnlyTo:
unbound = _dataset_token(sqlite_session)
sqlite_session.commit()
deleted = dataset_api_key_service.delete_keys_scoped_only_to(sqlite_session, dataset_id)
deleted = dataset_api_key_bindings.delete_keys_scoped_only_to(sqlite_session, dataset_id)
sqlite_session.commit()
assert deleted == []
api/tests/unit_tests/repositories/knowledge/test_dataset_api_key_repository.py
+150
View File
@@ -0,0 +1,150 @@
"""Real persistence checks for the resource API key boundary."""
from functools import partial
import pytest
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from models.account import Tenant
from models.dataset import Dataset
from models.enums import ApiTokenType
from models.model import ApiToken, DatasetApiTokenBinding
from repositories.knowledge import dataset_api_key_bindings
from repositories.knowledge.dataset_api_key_repository import DatasetApiKeyRepository
from services.auth.api_key_contracts import (
ApiKeyLimitExceededError,
ApiKeyNotFoundError,
ApiKeyResourceNotFoundError,
)
from services.knowledge.api_key_service import UnknownDatasetIdsError
@pytest.fixture
def repository(sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]) -> DatasetApiKeyRepository:
tenant = Tenant(name="Workspace")
tenant.id = "tenant"
sqlite_session.add(tenant)
sqlite_session.add_all(
[
Dataset(id="dataset", tenant_id="tenant", name="Dataset", created_by="actor"),
Dataset(id="other-dataset", tenant_id="tenant", name="Other dataset", created_by="actor"),
]
)
sqlite_session.commit()
return DatasetApiKeyRepository(session_factory=sqlite_session_factory)
@pytest.mark.parametrize("operation", ["list", "create", "delete"])
def test_resource_must_belong_to_workspace(repository: DatasetApiKeyRepository, operation: str) -> None:
def invoke() -> object:
if operation == "list":
return repository.list_keys("foreign", "dataset")
if operation == "create":
return repository.create_key("foreign", "dataset", max_keys=10, prefix="test-")
return repository.delete_key("foreign", "dataset", "missing")
with pytest.raises(ApiKeyResourceNotFoundError):
invoke()
def test_dataset_keys_use_persisted_bindings_and_cannot_manage_broader_scopes(
repository: DatasetApiKeyRepository,
sqlite_session: Session,
) -> None:
created = repository.create_key("tenant", "dataset", max_keys=10, prefix="ds-")
assert created.dataset_ids == ("dataset",)
assert created.token.startswith("ds-")
assert (
sqlite_session.scalar(
select(DatasetApiTokenBinding.dataset_id).where(DatasetApiTokenBinding.api_token_id == created.id)
)
== "dataset"
)
scopes: list[tuple[str, list[str], str]] = [
("unrestricted", [], "tenant"),
("multiple", ["dataset", "other-dataset"], "tenant"),
("foreign", ["dataset"], "foreign"),
("other", ["other-dataset"], "tenant"),
]
for key_id, dataset_ids, tenant_id in scopes:
sqlite_session.add(ApiToken(id=key_id, type=ApiTokenType.DATASET, token=key_id, tenant_id=tenant_id))
sqlite_session.add_all(DatasetApiTokenBinding(api_token_id=key_id, dataset_id=value) for value in dataset_ids)
sqlite_session.commit()
assert repository.list_keys("tenant", "dataset") == (created,)
for key_id in ["unrestricted", "multiple", "foreign", "other"]:
with pytest.raises(ApiKeyNotFoundError):
repository.delete_key("tenant", "dataset", key_id)
with pytest.raises(ApiKeyLimitExceededError):
repository.create_key("tenant", "dataset", max_keys=1, prefix="ds-")
repository.delete_key("tenant", "dataset", created.id)
assert repository.list_keys("tenant", "dataset") == ()
@pytest.mark.parametrize("workspace", [False, True])
def test_failed_binding_rolls_back_the_new_token(
repository: DatasetApiKeyRepository, sqlite_session: Session, monkeypatch: pytest.MonkeyPatch, workspace: bool
) -> None:
def fail_binding(_session: Session, _api_token_id: str, _dataset_ids: object) -> None:
raise RuntimeError("binding write failed")
monkeypatch.setattr(dataset_api_key_bindings, "bind_datasets", fail_binding)
create_key = (
partial(repository.create_workspace_key, "tenant", ("dataset",), max_keys=10, prefix="dataset-")
if workspace
else partial(repository.create_key, "tenant", "dataset", max_keys=10, prefix="ds-")
)
with pytest.raises(RuntimeError, match="binding write failed"):
create_key()
assert sqlite_session.scalar(select(ApiToken)) is None
@pytest.mark.parametrize("scope", [(), ("other-dataset",), ("dataset", "other-dataset")])
@pytest.mark.parametrize("workspace_creates_last", [False, True])
def test_both_scopes_share_the_workspace_quota(
repository: DatasetApiKeyRepository,
sqlite_session: Session,
scope: tuple[str, ...],
workspace_creates_last: bool,
) -> None:
# Neither another tenant's keys nor app keys consume the dataset quota.
sqlite_session.add_all(
[ApiToken(tenant_id="foreign", type=ApiTokenType.DATASET, token=f"foreign-{i}") for i in range(10)]
+ [ApiToken(tenant_id="tenant", type=ApiTokenType.APP, token=f"app-{i}") for i in range(10)]
)
sqlite_session.commit()
for _ in range(9):
repository.create_workspace_key("tenant", scope, max_keys=10, prefix="dataset-")
if workspace_creates_last:
repository.create_workspace_key("tenant", scope, max_keys=10, prefix="dataset-")
else:
repository.create_key("tenant", "dataset", max_keys=10, prefix="ds-")
with pytest.raises(ApiKeyLimitExceededError):
repository.create_key("tenant", "dataset", max_keys=10, prefix="ds-")
with pytest.raises(ApiKeyLimitExceededError):
repository.create_workspace_key("tenant", (), max_keys=10, prefix="dataset-")
assert len(repository.list_workspace_keys("tenant")) == 10
def test_workspace_operations_preserve_scope_and_tenant(
repository: DatasetApiKeyRepository, sqlite_session: Session
) -> None:
sqlite_session.add(Dataset(id="foreign-dataset", tenant_id="foreign", name="Foreign", created_by="actor"))
sqlite_session.commit()
with pytest.raises(UnknownDatasetIdsError):
repository.create_workspace_key("tenant", ("dataset", "foreign-dataset"), max_keys=10, prefix="dataset-")
assert sqlite_session.scalar(select(ApiToken)) is None
key = repository.create_workspace_key("tenant", ("dataset", "other-dataset"), max_keys=10, prefix="dataset-")
listed = repository.list_workspace_keys("tenant")
assert len(listed) == 1
assert set(listed[0].dataset_ids) == {"dataset", "other-dataset"}
assert repository.list_workspace_keys("foreign") == ()
with pytest.raises(ApiKeyNotFoundError):
repository.delete_workspace_key("foreign", key.id)
repository.delete_workspace_key("tenant", key.id)
assert repository.list_workspace_keys("tenant") == ()
api/tests/unit_tests/services/app/test_api_key_service.py
+204
View File
@@ -0,0 +1,204 @@
from dataclasses import dataclass, field
import pytest
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from extensions.application_services.app import build_app_api_key_service
from machinery.context import RequestContext
from models.agent import Agent, AgentConfigRevision, AgentConfigRevisionOperation, AgentScope, AgentSource, AgentStatus
from models.model import ApiToken, App, AppMode
from repositories.app.api_key_repository import AppApiKeyRepository
from services.app.api_key_service import AppApiKeyNotReadyError, AppApiKeyService
from services.auth.api_key_contracts import ApiKeyNotFoundError, ApiKeyResourceNotFoundError
@dataclass
class ObservingCache:
session_factory: sessionmaker[Session]
calls: list[tuple[str, str | None]] = field(default_factory=list)
def delete(self, token: str, scope: str | None = None) -> bool:
# A fresh connection must see the deletion before any external cache I/O.
with self.session_factory() as session:
assert session.scalar(select(ApiToken).where(ApiToken.token == token)) is None
self.calls.append((token, scope))
return True
def test_deletion_invalidates_cache_after_commit_and_not_on_failure(
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
) -> None:
sqlite_session.add(
App(id="app", tenant_id="tenant", name="App", mode=AppMode.CHAT, enable_site=False, enable_api=True)
)
sqlite_session.commit()
context = RequestContext("request", None, "actor", "tenant")
cache = ObservingCache(sqlite_session_factory)
service = AppApiKeyService(
keys=AppApiKeyRepository(session_factory=sqlite_session_factory),
cache=cache,
)
key = service.create_key(context, "app")
with pytest.raises(ApiKeyNotFoundError):
service.delete_key(context, "app", "missing")
assert cache.calls == []
service.delete_key(context, "app", key.id)
assert cache.calls == [(key.token, key.type)]
assert service.list_keys(context, "app") == ()
def test_composition_builds_working_api_key_service(
sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
sqlite_session.add(
App(id="app", tenant_id="tenant", name="App", mode=AppMode.CHAT, enable_site=False, enable_api=True)
)
sqlite_session.commit()
service = build_app_api_key_service(database_client=sqlite_session_factory)
context = RequestContext("request", None, "actor", "tenant")
key = service.create_key(context, "app")
assert service.list_keys(context, "app") == (key,)
@pytest.fixture
def service(sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]) -> AppApiKeyService:
sqlite_session.add(
App(id="app", tenant_id="tenant", name="App", mode=AppMode.CHAT, enable_site=False, enable_api=True)
)
sqlite_session.commit()
return build_app_api_key_service(database_client=sqlite_session_factory)
def _persist_agent(session: Session) -> Agent:
app = session.get(App, "app")
assert app is not None
app.mode = AppMode.AGENT
agent = Agent(
id="agent",
tenant_id="tenant",
name="Agent",
scope=AgentScope.ROSTER,
source=AgentSource.AGENT_APP,
status=AgentStatus.ACTIVE,
app_id="app",
)
session.add(agent)
session.commit()
return agent
def test_app_access_check_rejects_foreign_workspace(service: AppApiKeyService, sqlite_session: Session) -> None:
context = RequestContext("request", None, "actor", "foreign")
with pytest.raises(ApiKeyResourceNotFoundError, match="App not found"):
service.create_key(context, "app")
assert sqlite_session.scalar(select(ApiToken)) is None
@pytest.mark.parametrize("source", [AgentSource.AGENT_APP, AgentSource.IMPORTED])
def test_agent_requires_publish_visible_version_for_creation(
service: AppApiKeyService,
sqlite_session: Session,
source: AgentSource,
) -> None:
context = RequestContext("request", None, "actor", "tenant")
agent = _persist_agent(sqlite_session)
agent.source = source
sqlite_session.commit()
with pytest.raises(AppApiKeyNotReadyError):
service.create_key(context, "app")
with pytest.raises(AppApiKeyNotReadyError):
service.create_agent_key(context, "agent")
assert sqlite_session.scalar(select(ApiToken)) is None
agent.active_config_snapshot_id = "snapshot"
agent.active_config_is_published = False # New draft edits must not revoke a published version.
sqlite_session.add(
AgentConfigRevision(
tenant_id="tenant",
agent_id=agent.id,
current_snapshot_id="snapshot",
revision=1,
operation=AgentConfigRevisionOperation.PUBLISH_DRAFT,
)
)
sqlite_session.commit()
created = service.create_agent_key(context, "agent")
assert service.list_agent_keys(context, "agent") == (created,)
assert service.list_keys(context, "app") == (created,)
agent.active_config_snapshot_id = None
sqlite_session.commit()
assert service.list_agent_keys(context, "agent") == (created,)
service.delete_agent_key(context, "agent", created.id)
assert service.list_agent_keys(context, "agent") == ()
@pytest.mark.parametrize("invalid_revision", ["tenant", "agent", "snapshot", "unpublished"])
def test_creation_requires_a_publish_revision_for_the_owned_active_snapshot(
service: AppApiKeyService, sqlite_session: Session, invalid_revision: str
) -> None:
context = RequestContext("request", None, "actor", "tenant")
agent = _persist_agent(sqlite_session)
agent.active_config_snapshot_id = "snapshot"
agent.active_config_is_published = True
revision = AgentConfigRevision(
tenant_id="tenant",
agent_id="agent",
current_snapshot_id="snapshot",
revision=1,
operation=AgentConfigRevisionOperation.PUBLISH_DRAFT,
)
match invalid_revision:
case "tenant":
revision.tenant_id = "foreign"
case "agent":
revision.agent_id = "other-agent"
case "snapshot":
revision.current_snapshot_id = "old-snapshot"
case "unpublished":
revision.operation = AgentConfigRevisionOperation.IMPORT_PACKAGE
sqlite_session.add(revision)
sqlite_session.commit()
with pytest.raises(AppApiKeyNotReadyError):
service.create_key(context, "app")
with pytest.raises(AppApiKeyNotReadyError):
service.create_agent_key(context, "agent")
assert sqlite_session.scalar(select(ApiToken)) is None
@pytest.mark.parametrize(
"invalid_state", ["tenant", "archived", "scope", "source", "app-tenant", "missing-app", "app-mode"]
)
def test_agent_resolution_checks_entire_owner_chain(
service: AppApiKeyService, sqlite_session: Session, invalid_state: str
) -> None:
context = RequestContext("request", None, "actor", "tenant")
agent = _persist_agent(sqlite_session)
app = sqlite_session.get(App, "app")
assert app is not None
match invalid_state:
case "tenant":
agent.tenant_id = "foreign"
case "archived":
agent.status = AgentStatus.ARCHIVED
case "scope":
agent.scope = AgentScope.WORKFLOW_ONLY
case "source":
agent.source = AgentSource.ROSTER
case "app-tenant":
app.tenant_id = "foreign"
case "missing-app":
agent.app_id = "missing"
case "app-mode":
app.mode = AppMode.CHAT
sqlite_session.commit()
with pytest.raises(ApiKeyResourceNotFoundError, match="Agent not found"):
service.list_agent_keys(context, "agent")
api/tests/unit_tests/services/knowledge/__init__.py
View File
Whitespace-only changes.
api/tests/unit_tests/services/knowledge/test_api_key_service.py
+154
View File
@@ -0,0 +1,154 @@
"""Knowledge key composition and post-commit cache invalidation."""
from collections.abc import Callable
from functools import partial
import pytest
from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker
from extensions.application_services.knowledge import build_dataset_api_key_service
from machinery.context import RequestContext
from models.account import Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, DatasetPermission
from models.enums import PermissionEnum
from models.model import ApiToken
from services.api_token_service import ApiTokenCache
from services.auth.api_key_contracts import ApiKeyNotFoundError
from services.errors.account import NoPermissionError
from services.knowledge.api_key_service import DatasetApiKeyService
@pytest.fixture
def service(
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
config_overrides: Callable[..., None],
) -> DatasetApiKeyService:
config_overrides(RBAC_ENABLED=False)
tenant = Tenant(name="Workspace")
tenant.id = "tenant"
sqlite_session.add_all(
[
tenant,
TenantAccountJoin(tenant_id="tenant", account_id="actor", role=TenantAccountRole.OWNER),
Dataset(id="dataset", tenant_id="tenant", name="Dataset", created_by="actor", maintainer="maintainer"),
]
)
sqlite_session.commit()
return build_dataset_api_key_service(database_client=sqlite_session_factory)
@pytest.mark.parametrize("workspace", [False, True])
def test_dataset_keys_commit_before_invalidating_cache(
service: DatasetApiKeyService,
sqlite_session_factory: sessionmaker[Session],
monkeypatch: pytest.MonkeyPatch,
workspace: bool,
) -> None:
calls: list[tuple[str, str | None]] = []
def invalidate(token: str, scope: str | None = None) -> bool:
with sqlite_session_factory() as session:
assert session.scalar(select(ApiToken).where(ApiToken.token == token)) is None
calls.append((token, scope))
return True
monkeypatch.setattr(ApiTokenCache, "delete", invalidate)
context = RequestContext("request", None, "actor", "tenant")
key = (
service.create_workspace_key(context, ("dataset", "dataset"))
if workspace
else service.create_key(context, "dataset")
)
assert key.token.startswith("dataset-" if workspace else "ds-")
assert key.dataset_ids == ("dataset",)
assert service.list_keys(context, "dataset") == (key,)
delete_key = (
partial(service.delete_workspace_key, context) if workspace else partial(service.delete_key, context, "dataset")
)
with pytest.raises(ApiKeyNotFoundError):
delete_key("missing")
assert calls == []
delete_key(key.id)
assert calls == [(key.token, key.type)]
assert service.list_keys(context, "dataset") == ()
@pytest.mark.parametrize(
("permission", "role", "maintainer", "grant", "allowed"),
[
(PermissionEnum.ONLY_ME, TenantAccountRole.EDITOR, "other", None, False),
(PermissionEnum.ONLY_ME, TenantAccountRole.ADMIN, "other", None, False),
(PermissionEnum.ONLY_ME, TenantAccountRole.OWNER, "other", None, True),
(PermissionEnum.ONLY_ME, TenantAccountRole.EDITOR, "actor", None, True),
(PermissionEnum.PARTIAL_TEAM, TenantAccountRole.EDITOR, "other", None, False),
(PermissionEnum.PARTIAL_TEAM, TenantAccountRole.EDITOR, "other", "allowed", True),
(PermissionEnum.PARTIAL_TEAM, TenantAccountRole.EDITOR, "actor", None, True),
(PermissionEnum.PARTIAL_TEAM, TenantAccountRole.EDITOR, "other", "revoked", False),
(PermissionEnum.PARTIAL_TEAM, TenantAccountRole.EDITOR, "other", "foreign", False),
(PermissionEnum.ALL_TEAM, TenantAccountRole.EDITOR, "other", None, True),
],
)
def test_dataset_acl_controls_key_management(
service: DatasetApiKeyService,
sqlite_session: Session,
permission: PermissionEnum,
role: TenantAccountRole,
maintainer: str,
grant: str | None,
allowed: bool,
) -> None:
dataset = sqlite_session.get(Dataset, "dataset")
member = sqlite_session.scalar(select(TenantAccountJoin))
assert dataset is not None
assert member is not None
dataset.permission = permission
dataset.maintainer = maintainer
member.role = role
if grant:
sqlite_session.add(
DatasetPermission(
dataset_id="dataset",
account_id="actor",
tenant_id="foreign" if grant == "foreign" else "tenant",
has_permission=grant != "revoked",
)
)
sqlite_session.commit()
context = RequestContext("request", None, "actor", "tenant")
if allowed:
key = service.create_key(context, "dataset")
assert service.list_keys(context, "dataset") == (key,)
else:
with pytest.raises(NoPermissionError):
service.create_key(context, "dataset")
with pytest.raises(NoPermissionError):
service.list_keys(context, "dataset")
with pytest.raises(NoPermissionError):
service.delete_key(context, "dataset", "missing")
assert sqlite_session.scalar(select(ApiToken)) is None
def test_acl_requires_membership_in_the_active_workspace(
service: DatasetApiKeyService, sqlite_session: Session
) -> None:
member = sqlite_session.scalar(select(TenantAccountJoin))
assert member is not None
member.tenant_id = "foreign"
sqlite_session.commit()
with pytest.raises(NoPermissionError):
service.create_key(RequestContext("request", None, "actor", "tenant"), "dataset")
assert sqlite_session.scalar(select(ApiToken)) is None
def test_rbac_admitted_request_uses_rbac_instead_of_community_acl(
service: DatasetApiKeyService, config_overrides: Callable[..., None]
) -> None:
config_overrides(RBAC_ENABLED=True)
key = service.create_key(RequestContext("request", None, "rbac-actor", "tenant"), "dataset")
assert key.dataset_ids == ("dataset",)
api/tests/unit_tests/services/test_app_service.py
+2 -3
View File
@@ -418,9 +418,8 @@ def test_unpublished_agent_app_access_cannot_be_enabled(
commits: list[str] = []
event.listen(sqlite_session, "after_commit", lambda _session: commits.append("commit"))
with patch("services.app_service.agent_has_workflow_callable_active_snapshot", return_value=False):
with pytest.raises(AgentAccessNotReadyError):
update_status(AppService(), app, True, session=sqlite_session)
with pytest.raises(AgentAccessNotReadyError):
update_status(AppService(), app, True, session=sqlite_session)
assert app.enable_site is False
assert app.enable_api is False
packages/contracts/generated/api/console/apps/orpc.gen.ts
-9
View File
@@ -4498,8 +4498,6 @@ export const byAppId2 = {
}
/**
* Delete an API key for an app
*
* Delete an API key for an app
*/
export const delete16 = oc
@@ -4510,7 +4508,6 @@ export const delete16 = oc
operationId: 'deleteAppsByResourceIdApiKeysByApiKeyId',
path: '/apps/{resource_id}/api-keys/{api_key_id}',
successStatus: 204,
summary: 'Delete an API key for an app',
tags: ['console'],
})
.input(z.object({ params: zDeleteAppsByResourceIdApiKeysByApiKeyIdPath }))
@@ -4521,8 +4518,6 @@ export const byApiKeyId = {
}
/**
* Get all API keys for an app
*
* Get all API keys for an app
*/
export const get85 = oc
@@ -4532,15 +4527,12 @@ export const get85 = oc
method: 'GET',
operationId: 'getAppsByResourceIdApiKeys',
path: '/apps/{resource_id}/api-keys',
summary: 'Get all API keys for an app',
tags: ['console'],
})
.input(z.object({ params: zGetAppsByResourceIdApiKeysPath }))
.output(zGetAppsByResourceIdApiKeysResponse)
/**
* Create a new API key for an app
*
* Create a new API key for an app
*/
export const post63 = oc
@@ -4551,7 +4543,6 @@ export const post63 = oc
operationId: 'postAppsByResourceIdApiKeys',
path: '/apps/{resource_id}/api-keys',
successStatus: 201,
summary: 'Create a new API key for an app',
tags: ['console'],
})
.input(z.object({ params: zPostAppsByResourceIdApiKeysPath }))
packages/contracts/generated/api/console/datasets/orpc.gen.ts
-9
View File
@@ -1698,8 +1698,6 @@ export const byDatasetId = {
}
/**
* Delete an API key for a dataset
*
* Delete an API key for a dataset
*/
export const delete10 = oc
@@ -1710,7 +1708,6 @@ export const delete10 = oc
operationId: 'deleteDatasetsByResourceIdApiKeysByApiKeyId',
path: '/datasets/{resource_id}/api-keys/{api_key_id}',
successStatus: 204,
summary: 'Delete an API key for a dataset',
tags: ['console'],
})
.input(z.object({ params: zDeleteDatasetsByResourceIdApiKeysByApiKeyIdPath }))
@@ -1721,8 +1718,6 @@ export const byApiKeyId2 = {
}
/**
* Get all API keys for a dataset
*
* Get all API keys for a dataset
*/
export const get35 = oc
@@ -1732,15 +1727,12 @@ export const get35 = oc
method: 'GET',
operationId: 'getDatasetsByResourceIdApiKeys',
path: '/datasets/{resource_id}/api-keys',
summary: 'Get all API keys for a dataset',
tags: ['console'],
})
.input(z.object({ params: zGetDatasetsByResourceIdApiKeysPath }))
.output(zGetDatasetsByResourceIdApiKeysResponse)
/**
* Create a new API key for a dataset
*
* Create a new API key for a dataset
*/
export const post22 = oc
@@ -1751,7 +1743,6 @@ export const post22 = oc
operationId: 'postDatasetsByResourceIdApiKeys',
path: '/datasets/{resource_id}/api-keys',
successStatus: 201,
summary: 'Create a new API key for a dataset',
tags: ['console'],
})
.input(z.object({ params: zPostDatasetsByResourceIdApiKeysPath }))