fix(flows): align MongoDB state serialization

This commit is contained in:
ViditOstwal
2026-09-18 13:40:56 +05:30
parent 496a3cc81a
commit a453e4f525
4 changed files with 51 additions and 20 deletions

View File

@@ -3,6 +3,7 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from datetime import date, datetime
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
@@ -15,6 +16,20 @@ if TYPE_CHECKING:
_persistence_registry: dict[str, type[FlowPersistence]] = {}
def _json_default(obj: Any) -> Any:
"""Serialize flow-state values that are not JSON-native."""
if isinstance(obj, BaseModel):
try:
return obj.model_dump(mode="json")
except Exception:
return obj.model_dump(mode="python")
if isinstance(obj, (set, tuple)):
return list(obj)
if isinstance(obj, (date, datetime)):
return obj.isoformat()
return str(obj)
class FlowPersistence(BaseModel, ABC):
"""Abstract base class for flow state persistence.

View File

@@ -31,7 +31,7 @@ from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field, PrivateAttr
from crewai.flow.persistence.base import FlowPersistence
from crewai.flow.persistence.base import FlowPersistence, _json_default
if TYPE_CHECKING:
@@ -190,7 +190,7 @@ class MongoDbFlowPersistence(FlowPersistence):
"flow_uuid": flow_uuid,
"method_name": method_name,
"timestamp": datetime.now(timezone.utc).isoformat(),
"state_json": json.dumps(state_dict),
"state_json": json.dumps(state_dict, default=_json_default),
"seq": self._next_sequence(self.states_collection),
}
)
@@ -221,8 +221,8 @@ class MongoDbFlowPersistence(FlowPersistence):
{"flow_uuid": flow_uuid},
{
"flow_uuid": flow_uuid,
"context_json": json.dumps(context.to_dict()),
"state_json": json.dumps(state_dict),
"context_json": json.dumps(context.to_dict(), default=_json_default),
"state_json": json.dumps(state_dict, default=_json_default),
"created_at": datetime.now(timezone.utc).isoformat(),
},
upsert=True,

View File

@@ -3,7 +3,7 @@
from __future__ import annotations
from contextlib import closing
from datetime import date, datetime, timezone
from datetime import datetime, timezone
import json
import os
from pathlib import Path
@@ -15,27 +15,13 @@ from crewai_core.paths import db_storage_path
from pydantic import BaseModel, Field, PrivateAttr, model_validator
from typing_extensions import Self
from crewai.flow.persistence.base import FlowPersistence
from crewai.flow.persistence.base import FlowPersistence, _json_default
if TYPE_CHECKING:
from crewai.flow.async_feedback.types import PendingFeedbackContext
def _json_default(obj: Any) -> Any:
"""Fallback serializer for non-primitive types in JSON dumps."""
if isinstance(obj, BaseModel):
try:
return obj.model_dump(mode="json")
except Exception:
return obj.model_dump(mode="python")
if isinstance(obj, (set, tuple)):
return list(obj)
if isinstance(obj, (date, datetime)):
return obj.isoformat()
return str(obj)
class SQLiteFlowPersistence(FlowPersistence):
"""SQLite-based implementation of flow state persistence.

View File

@@ -176,6 +176,36 @@ def test_basemodel_state_serialized_as_json(
assert loaded["when"].startswith("2026-01-02T03:04:05")
def test_dict_state_serializes_non_json_values(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_patch_client(monkeypatch)
class NestedState(BaseModel):
when: datetime
when = datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc)
persistence = MongoDbFlowPersistence(CONN)
persistence.save_state(
"flow-1",
"s",
{
"when": when,
"tags": {"a", "b"},
"items": (1, 2),
"nested": NestedState(when=when),
},
)
assert persistence.load_state("flow-1") == {
"when": "2026-01-02T03:04:05+00:00",
"tags": ["a", "b"],
"items": [1, 2],
"nested": {"when": "2026-01-02T03:04:05Z"},
}
def test_missing_connection_string_raises(
monkeypatch: pytest.MonkeyPatch,
) -> None: