Python / AG-UI: Recommended way to JSON-serialize AGUIThreadSnapshot for a durable store?

[Marimo (SuperMarimoBros)]

I'd like to implement a durable agent_framework_ag_ui.AGUIThreadSnapshotStore backed by a database, using agent-framework-core 1.20.0 and agent-framework-ag-ui 1.5.0. My goal is that snapshots round-trip without errors and without silent data loss.

What I observed

• session_state and state pass through agent_framework._serialization.make_json_safe() before reaching the store. In most cases it does not raise; unknown types fall back to str():

  • set({1, 2}) → "{1, 2}"
  • an unregistered dataclass in AgentSession.state → restored as a plain dict

• Some values still are not valid strict JSON:

  • NaN / Infinity
  • and some backends reject strings containing \u0000

• A plain Enum makes make_json_safe() raise RecursionError.

• Failures are logged and swallowed (as designed in Python: bridge AG-UI request state and session continuity #7084), with different outcomes:

  • _safe_serialize_session_continuation_state(): if serializing session_state fails (including the RecursionError above), the snapshot is still saved with session_state=None.
  • ThreadSnapshotSession.save(): if the store rejects the snapshot (e.g. NaN or \u0000), the new snapshot is not saved and the previous one remains.

  In both cases the run completes normally.

Question

Is it the application's responsibility to keep snapshot values strictly JSON-native, converting them before they reach the runner (with custom session_state types registered via agent_framework.register_state_type())? Or is the stringification by make_json_safe() meant to be relied on?

[Kurt Geminic (kgeminicdev)]

From reading the code I'd treat it as the application's job. make_json_safe is a last-resort "json.dumps must not raise" pass, its docstring says as much, and it's lossy by design (set -> str, dataclass -> dict):

agent-framework/python/packages/core/agent_framework/_serialization.py

Lines 733 to 746 in 40763be

	def make_json_safe(obj: Any) -> Any:
	"""Recursively convert an object to a JSON-serializable form.
	Handles dataclasses, Pydantic models, objects with ``to_dict``/``dict``/``__dict__``,
	datetimes, bytes (base64), lists, dicts, and primitives. Falls back to ``str()`` for
	any remaining non-serializable value so that ``json.dumps`` never raises a
	``TypeError``.
	Args:
	obj: Object to make JSON safe.
	Returns:
	A JSON-serializable version of the object.
	"""

The core durable path is stricter. AgentSession.to_dict() warns on unregistered values with "durable session stores will reject it. Call register_state_type()", and the SessionStore codec runs _validate_durable_state_value, which rejects non-finite floats and anything that isn't a JSON scalar/list/mapping:

agent-framework/python/packages/core/agent_framework/_sessions.py

Lines 449 to 466 in 40763be

	def _validate_durable_state_value(value: Any, *, path: str) -> None:
	"""Validate that serialized state can round-trip through the durable codecs."""
	if isinstance(value, float) and not math.isfinite(value):
	raise ValueError(f"Session state value at {path} must be a finite float.")
	if type(value) in _STATE_SCALAR_TYPES:
	return
	if isinstance(value, list):
	for index, item in enumerate(cast(list[Any], value)):
	_validate_durable_state_value(item, path=f"{path}[{index}]")
	return
	if isinstance(value, Mapping):
	for key, item in cast(Mapping[Any, Any], value).items():
	_validate_durable_state_value(item, path=f"{path}.{key}")
	return
	raise TypeError(
	f"Session state value at {path} has unsupported serialized type {type(value).__name__!r}; "
	"call register_state_type() with a codec that returns JSON-compatible values."
	)

The AG-UI snapshot path just doesn't apply that validation. It stringifies instead and logs/swallows failures so a finished run still gets RUN_FINISHED.

What I'd do: register every custom session_state type with register_state_type(cls, type_id="yourapp.x", encoder=..., decoder=...). With that, to_dict emits {"...": ..., "type": "yourapp.x"} and from_dict rebuilds the instance. I checked this round-trips through json.dumps/json.loads on 1.20.0. Convert Enums, sets and NaN/inf yourself before they go into state or Shared State. Then in your store's save(), serialize strictly with json.dumps(dataclasses.asdict(snapshot), allow_nan=False), and reject NUL if your backend can't take it (Postgres jsonb rejects \u0000). The framework will log the failure and keep the previous snapshot, which beats silently storing "{1, 2}".

The Enum RecursionError looks like a bug in make_json_safe to me (it walks vars() into the enum class), probably worth a separate issue. I don't know whether the maintainers consider the str() fallback part of the contract, but nothing in the snapshot code depends on it round-tripping.

[Rodion Kazennov (nefayran)]

Short answer: keep the values JSON-native on the application side. make_json_safe() is a best-effort fallback so that a finished run does not fail on persistence (#7084). It does not promise a lossless round trip, and as you found, part of what it returns is not strict JSON.

The core package holds its own durable session snapshots to a stricter contract, and that contract is a good one to copy. _SessionStatePayload.serialize() runs _serialize_state() and then _validate_durable_state_value(), which rejects non-finite floats and anything that is not a JSON scalar, list or mapping once the registered codecs have run (_sessions.py#L449-L486). For an unregistered type, _serialize_value() already warns that "durable session stores will reject it" (L412-L418).

What I would do in a database-backed AGUIThreadSnapshotStore:

1. Register every custom session-state class at import time with register_state_type(). Pydantic models and classes with to_dict()/from_dict() get default codecs. A plain dataclass needs encoder= and decoder= (for example asdict and lambda d: Prefs(**{k: v for k, v in d.items() if k != "type"})); without them it comes back as a dict, which is what you saw.
2. Convert sets and enums before they reach the runner: a set to a sorted list, an enum to its .value. An enum with a str or int mixin (class Level(str, Enum)) passes through make_json_safe() unchanged because it is an instance of str, but it is restored as a plain string. A plain Enum has no branch of its own in make_json_safe(), falls through to vars() and recurses through the enum class, which is the RecursionError you hit (I reproduced it with a copy of the function at main 40763be).
3. Encode strictly inside the store and let it fail loudly: json.dumps(payload, allow_nan=False, ensure_ascii=False), and refuse or strip "\x00" if the backend is PostgreSQL, whose text and jsonb types cannot hold U+0000. When save() raises, ThreadSnapshotSession.save() logs the exception and "the previous snapshot stays authoritative for hydration" (_snapshot_session.py#L150-L182), so alert on that log record, since the run itself still completes.
4. Add a round-trip test for the store: save, get and compare a representative snapshot, including each registered state type.

[Marimo (SuperMarimoBros)]

Kurt Geminic (@kgeminicdev) Rodion Kazennov (@nefayran) Thank you both, this was really helpful. I'll use the approach you described as the basis for my implementation.

The key point for me is to guarantee the round trip with tests: what the store's get() returns after save() must equal the original snapshot, including registered state types, multiple turns, and resumes after an interrupt.

One note on my design: since save() always receives the full accumulated messages, writing them as a single JSON value means rewriting the whole history on every turn, which gets heavier as the conversation grows. So inside the store I'm planning to keep messages in their own table and append only new ones (by message id), with state and the other fields stored separately. Each part will go through the strict serialization you suggested.

As you suggested, I'll also open a separate issue for the RecursionError that make_json_safe() raises on a plain Enum.

[Marimo (SuperMarimoBros)]

Update: Rodion Kazennov (@nefayran) opened #9194 for the RecursionError on a plain Enum. Thanks!