diff --git a/lib/crewai/src/crewai/__init__.py b/lib/crewai/src/crewai/__init__.py
index 94fef38e2..5835fe7f9 100644
--- a/lib/crewai/src/crewai/__init__.py
+++ b/lib/crewai/src/crewai/__init__.py
@@ -192,15 +192,13 @@ try:
 
     from typing import Annotated
 
-    from pydantic import Discriminator, Tag
+    from pydantic import Field
 
-    from crewai.state.runtime import RuntimeState, _entity_discriminator
+    from crewai.state.runtime import RuntimeState
 
     Entity = Annotated[
-        Annotated[Flow, Tag("flow")]  # type: ignore[type-arg]
-        | Annotated[Crew, Tag("crew")]
-        | Annotated[Agent, Tag("agent")],
-        Discriminator(_entity_discriminator),
+        Flow | Crew | Agent,  # type: ignore[type-arg]
+        Field(discriminator="entity_type"),
     ]
 
     RuntimeState.model_rebuild(
diff --git a/lib/crewai/src/crewai/state/runtime.py b/lib/crewai/src/crewai/state/runtime.py
index a4bc8584c..c5655398e 100644
--- a/lib/crewai/src/crewai/state/runtime.py
+++ b/lib/crewai/src/crewai/state/runtime.py
@@ -15,7 +15,6 @@ from pydantic import (
     ModelWrapValidatorHandler,
     PrivateAttr,
     RootModel,
-    SerializerFunctionWrapHandler,
     model_serializer,
     model_validator,
 )
@@ -84,10 +83,10 @@ class RuntimeState(RootModel):  # type: ignore[type-arg]
         """The execution event record."""
         return self._event_record
 
-    @model_serializer(mode="wrap")
-    def _serialize(self, handler: SerializerFunctionWrapHandler) -> CheckpointPayload:
+    @model_serializer(mode="plain")
+    def _serialize(self) -> dict[str, Any]:
         return {
-            "entities": handler(self),
+            "entities": [e.model_dump() for e in self.root],
             "event_record": self._event_record.model_dump(),
         }