feat(api-docs): Improve SDK API documentation (#6626)

Author: khvn26

api/.nvmrc

@@ -0,0 +1 @@
+20

api/Makefile

@@ -12,6 +12,8 @@ POETRY_VERSION ?= 2.2.1
 SAML_REVISION ?= v1.6.6
 RBAC_REVISION ?= v0.13.0
 
+OPENAPI_FORMAT_VERSION ?= 1.23.0
+
 -include .env-local
 -include $(DOTENV_OVERRIDE_FILE)
 
@@ -113,10 +115,18 @@ serve-with-task-processor: TASK_RUN_METHOD=TASK_PROCESSOR
 serve-with-task-processor:
 	make -j2 serve run-task-processor
 
+.PHONY: generate-flagsmith-sdk-openapi
+generate-flagsmith-sdk-openapi: OPENAPI_SERVERS=[{"url": "https://edge.api.flagsmith.com", "description": "Flagsmith Edge API"}]
+generate-flagsmith-sdk-openapi:
+	poetry run python manage.py spectacular | \
+	  npx openapi-format@${OPENAPI_FORMAT_VERSION} /dev/fd/0 \
+	    --filterFile openapi-filter-flagsmith-sdk.yml \
+	    --output ../sdk/openapi.yaml
+
 .PHONY: generate-ld-client-types
 generate-ld-client-types:
 	curl -sSL https://app.launchdarkly.com/api/v2/openapi.json | \
-	npx openapi-format /dev/fd/0 \
+	npx openapi-format@${OPENAPI_FORMAT_VERSION} /dev/fd/0 \
 		--filterFile openapi-filter-launchdarkly.yaml | \
 	datamodel-codegen \
 		--output integrations/launch_darkly/types.py \
@@ -130,7 +140,7 @@ generate-ld-client-types:
 .PHONY: generate-grafana-client-types
 generate-grafana-client-types:
 	curl -sSL https://raw.githubusercontent.com/grafana/grafana/refs/heads/main/public/openapi3.json | \
-	npx openapi-format /dev/fd/0 \
+	npx openapi-format@${OPENAPI_FORMAT_VERSION} /dev/fd/0 \
 		--filterFile openapi-filter-grafana.yaml | \
 	datamodel-codegen \
 		--output integrations/grafana/types.py \
@@ -153,7 +163,7 @@ integrate-private-tests:
 	rm -rf ./flagsmith-saml ./flagsmith-rbac ./flagsmith-workflows ./flagsmith-auth-controller
 
 .PHONY: generate-docs
-generate-docs:
+generate-docs: generate-flagsmith-sdk-openapi
 	poetry run flagsmith docgen metrics > ../docs/docs/administration-and-security/platform-configuration/metrics.md
 
 .PHONY: add-known-sdk-version

api/api/openapi.py

@@ -2,12 +2,14 @@
 
 from drf_spectacular import generators, openapi
 from drf_spectacular.extensions import (
+    OpenApiAuthenticationExtension,
     OpenApiSerializerExtension,
 )
 from drf_spectacular.plumbing import ResolvedComponent, safe_ref
 from drf_spectacular.plumbing import append_meta as append_meta_orig
-from pydantic import BaseModel
+from pydantic import TypeAdapter
 from rest_framework.request import Request
+from typing_extensions import is_typeddict
 
 
 def append_meta(schema: dict[str, Any], meta: dict[str, Any]) -> dict[str, Any]:
@@ -135,48 +137,79 @@ def _update_security_for_mcp(self, schema: dict[str, Any]) -> dict[str, Any]:
         return schema
 
 
-class PydanticSchemaExtension(
-    OpenApiSerializerExtension  # type: ignore[no-untyped-call]
-):
+class TypedDictSchemaExtension(OpenApiSerializerExtension):
     """
     An OpenAPI extension that allows drf-spectacular to generate schema documentation
-    from Pydantic models.
+    from TypedDicts via Pydantic.
 
-    This extension is automatically used when a Pydantic BaseModel subclass is passed
+    This extension is automatically used when a TypedDict subclass is passed
     as a response type in @extend_schema decorators.
     """
 
-    target_class = "pydantic.BaseModel"
-    match_subclasses = True
+    @classmethod
+    def _matches(cls, target: type[Any]) -> bool:
+        return is_typeddict(target)
 
     def get_name(
         self,
         auto_schema: openapi.AutoSchema | None = None,
         direction: Literal["request", "response"] | None = None,
-    ) -> str | None:
-        return self.target.__name__  # type: ignore[no-any-return]
+    ) -> str:
+        name: str = self.target.__name__
+        return name
 
     def map_serializer(
         self,
         auto_schema: openapi.AutoSchema,
         direction: str,
     ) -> dict[str, Any]:
-        model_cls: type[BaseModel] = self.target
-
-        model_json_schema = model_cls.model_json_schema(
+        model_json_schema = TypeAdapter(self.target).json_schema(
             mode="serialization",
-            ref_template="#/components/schemas/{model}",
+            ref_template="#/components/schemas/%s{model}" % self.get_name(),
         )
 
         # Register nested definitions as components
         if "$defs" in model_json_schema:
             for ref_name, schema_kwargs in model_json_schema.pop("$defs").items():
                 component = ResolvedComponent(  # type: ignore[no-untyped-call]
-                    name=ref_name,
+                    name=self.get_name() + ref_name,
                     type=ResolvedComponent.SCHEMA,
                     object=ref_name,
                     schema=schema_kwargs,
                 )
                 auto_schema.registry.register_on_missing(component)
 
         return model_json_schema
+
+
+class EnvironmentKeyAuthenticationExtension(OpenApiAuthenticationExtension):  # type: ignore[no-untyped-call]
+    target_class = "environments.authentication.EnvironmentKeyAuthentication"
+    name = "Environment API Key"
+
+    def get_security_definition(
+        self, auto_schema: openapi.AutoSchema | None = None
+    ) -> dict[str, Any]:
+        return {
+            "type": "apiKey",
+            "in": "header",
+            "name": "X-Environment-Key",
+            "description": "For SDK endpoints. <a href='https://docs.flagsmith.com/clients/rest#public-api-endpoints'>Find out more</a>.",
+        }
+
+
+class MasterAPIKeyAuthenticationExtension(OpenApiAuthenticationExtension):  # type: ignore[no-untyped-call]
+    target_class = "api_keys.authentication.MasterAPIKeyAuthentication"
+    name = "Master API Key"
+
+    def get_security_definition(
+        self, auto_schema: openapi.AutoSchema | None = None
+    ) -> dict[str, Any]:
+        return {
+            "type": "apiKey",
+            "in": "header",
+            "name": "Authorization",
+            "description": (
+                "For Admin API endpoints. "
+                "<a href='https://docs.flagsmith.com/clients/rest#private-api-endpoints'>Find out more</a>."
+            ),
+        }

api/app/settings/common.py

@@ -534,26 +534,7 @@
     "SWAGGER_UI_SETTINGS": {
         "deepLinking": True,
     },
-    "SECURITY": [
-        {"Private": []},
-        {"Public": []},
-    ],
-    "APPEND_COMPONENTS": {
-        "securitySchemes": {
-            "Private": {
-                "type": "apiKey",
-                "in": "header",
-                "name": "Authorization",
-                "description": "For Private Endpoints. <a href='https://docs.flagsmith.com/clients/rest#private-api-endpoints'>Find out more</a>.",
-            },
-            "Public": {
-                "type": "apiKey",
-                "in": "header",
-                "name": "X-Environment-Key",
-                "description": "For Public Endpoints. <a href='https://docs.flagsmith.com/clients/rest#public-api-endpoints'>Find out more</a>.",
-            },
-        },
-    },
+    "SERVERS": env.json("OPENAPI_SERVERS", default=[]),
     "DEFAULT_GENERATOR_CLASS": "api.openapi.SchemaGenerator",
     "EXTENSIONS": [
         "api.openapi",

api/conftest.py

@@ -81,7 +81,6 @@
 )
 from projects.tags.models import Tag
 from segments.models import Condition, Segment, SegmentRule
-from tests.test_helpers import fix_issue_3869
 from tests.types import (
     AdminClientAuthType,
     EnableFeaturesFixture,
@@ -101,10 +100,6 @@ def pytest_addoption(parser: pytest.Parser) -> None:
     )
 
 
-def pytest_sessionstart(session: pytest.Session) -> None:
-    fix_issue_3869()  # type: ignore[no-untyped-call]
-
-
 @pytest.fixture()
 def post_request_mock(mocker: MockerFixture) -> MagicMock:
     def mocked_request(*args, **kwargs) -> None:  # type: ignore[no-untyped-def]

api/environments/identities/serializers.py

@@ -6,14 +6,10 @@
 from rest_framework.exceptions import ValidationError
 
 from environments.identities.models import Identity
-from environments.identities.traits.fields import TraitValueField
 from environments.models import Environment
 from environments.serializers import EnvironmentSerializerFull
 from features.models import FeatureState
-from features.serializers import (
-    FeatureStateSerializerFull,
-    SDKFeatureStateSerializer,
-)
+from features.serializers import FeatureStateSerializerFull
 
 
 class IdentifierOnlyIdentitySerializer(serializers.ModelSerializer):  # type: ignore[type-arg]
@@ -53,16 +49,6 @@ def save(self, **kwargs):  # type: ignore[no-untyped-def]
         return super(IdentitySerializer, self).save(**kwargs)
 
 
-class SDKIdentitiesResponseSerializer(serializers.Serializer):  # type: ignore[type-arg]
-    class _TraitSerializer(serializers.Serializer):  # type: ignore[type-arg]
-        trait_key = serializers.CharField()
-        trait_value = TraitValueField()
-
-    identifier = serializers.CharField()
-    flags = serializers.ListField(child=SDKFeatureStateSerializer())
-    traits = serializers.ListSerializer(child=_TraitSerializer())  # type: ignore[var-annotated]
-
-
 class SDKIdentitiesQuerySerializer(serializers.Serializer):  # type: ignore[type-arg]
     identifier = serializers.CharField(required=True)
     transient = serializers.BooleanField(default=False)

api/environments/identities/views.py

@@ -13,6 +13,7 @@
 from django.views.decorators.cache import cache_page
 from django.views.decorators.vary import vary_on_headers
 from drf_spectacular.utils import extend_schema
+from flagsmith_schemas import api as api_schemas
 from rest_framework import status, viewsets
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
@@ -25,7 +26,6 @@
 from environments.identities.serializers import (
     IdentitySerializer,
     SDKIdentitiesQuerySerializer,
-    SDKIdentitiesResponseSerializer,
 )
 from environments.models import Environment
 from environments.permissions.permissions import NestedEnvironmentPermissions
@@ -151,9 +151,9 @@ class SDKIdentities(SDKAPIView):
     throttle_classes = []
 
     @extend_schema(
-        responses={200: SDKIdentitiesResponseSerializer},
+        responses={200: api_schemas.V1IdentitiesResponse},
         parameters=[SDKIdentitiesQuerySerializer],
-        operation_id="identify_user",
+        operation_id="sdk_v1_get_identities",
     )
     @method_decorator(vary_on_headers(SDK_ENVIRONMENT_KEY_HEADER))
     @method_decorator(
@@ -163,6 +163,9 @@ class SDKIdentities(SDKAPIView):
         )
     )
     def get(self, request):  # type: ignore[no-untyped-def]
+        """
+        Retrieve the flags and traits for an identity.
+        """
         identifier = request.query_params.get("identifier")
         if not identifier:
             return Response(
@@ -234,11 +237,14 @@ def get_serializer_context(self):  # type: ignore[no-untyped-def]
         return context
 
     @extend_schema(
-        request=IdentifyWithTraitsSerializer,
-        responses={200: SDKIdentitiesResponseSerializer},
-        operation_id="identify_user_with_traits",
+        request=api_schemas.V1IdentitiesRequest,
+        responses={200: api_schemas.V1IdentitiesResponse},
+        operation_id="sdk_v1_post_identities",
     )
     def post(self, request):  # type: ignore[no-untyped-def]
+        """
+        Identify a user, set their traits, and retrieve their flags.
+        """
         serializer = self.get_serializer(data=request.data)
         serializer.is_valid(raise_exception=True)
         instance = serializer.save()

api/environments/sdk/schemas.py

@@ -4,6 +4,7 @@
 from django.utils.decorators import method_decorator
 from django.views.decorators.http import condition
 from drf_spectacular.utils import extend_schema
+from flagsmith_schemas.api import V1EnvironmentDocumentResponse
 from rest_framework.request import Request
 from rest_framework.response import Response
 from rest_framework.views import APIView
@@ -14,7 +15,6 @@
 )
 from environments.models import Environment
 from environments.permissions.permissions import EnvironmentKeyPermissions
-from environments.sdk.schemas import SDKEnvironmentDocumentModel
 
 
 def get_last_modified(request: Request) -> datetime | None:
@@ -30,9 +30,16 @@ class SDKEnvironmentAPIView(APIView):
     def get_authenticators(self):  # type: ignore[no-untyped-def]
         return [EnvironmentKeyAuthentication(required_key_prefix="ser.")]
 
-    @extend_schema(responses={200: SDKEnvironmentDocumentModel})
+    @extend_schema(
+        responses={200: V1EnvironmentDocumentResponse},
+        operation_id="sdk_v1_environment_document",
+    )
     @method_decorator(condition(last_modified_func=get_last_modified))
     def get(self, request: Request) -> Response:
+        """
+        Retrieve the environment document.
+        Used by SDKs in local evaluation mode, and Edge Proxy.
+        """
         environment_document = Environment.get_environment_document(
             request.environment.api_key,
         )

api/environments/sdk/views.py

@@ -8,6 +8,7 @@
 from django.db.models import Count, Q, QuerySet
 from django.utils.decorators import method_decorator
 from drf_spectacular.utils import OpenApiParameter, extend_schema
+from flagsmith_schemas.api import V1EnvironmentDocumentResponse
 from rest_framework import mixins, status, viewsets
 from rest_framework.decorators import action
 from rest_framework.exceptions import PermissionDenied, ValidationError
@@ -24,7 +25,6 @@
     EnvironmentPermissions,
     NestedEnvironmentPermissions,
 )
-from environments.sdk.schemas import SDKEnvironmentDocumentModel
 from features.versioning.models import EnvironmentFeatureVersion
 from features.versioning.tasks import (
     disable_v2_versioning,
@@ -288,7 +288,7 @@ def detailed_permissions(
         )
         return Response(serializer.data)
 
-    @extend_schema(responses={200: SDKEnvironmentDocumentModel})
+    @extend_schema(responses={200: V1EnvironmentDocumentResponse})
     @action(detail=True, methods=["GET"], url_path="document")
     def get_document(self, request, api_key: str):  # type: ignore[no-untyped-def]
         environment = (