[Bug] Python memory can persist Pemja-incompatible values and fail after state backend restore

[joeyutong]

Search before asking

• I searched in the issues and found nothing similar.

Description

Python-side memory values can currently accept objects that are not stable across the Pemja / Flink state backend boundary. When such values are written to sensory memory or short-term memory and then checkpointed/restored, the restored Java state may contain pemja.core.object.PyObject / PyJObject wrappers rather than a materialized, process-independent payload. After a TaskManager / Python process restart, converting those restored objects back to Python can fail or crash in Pemja native conversion.

There are two related user-facing problems:

1. The memory documentation does not clearly define the Python-side value contract.
   The docs mention primitive types, collections, Java POJOs, general class types, and MemoryObject, but do not say whether Python values must be recursively JSON-like / checkpoint-stable. For example, it is unclear whether uuid.UUID, Pydantic models, custom Python classes, or a dict containing those objects are valid memory values.

2. The built-in tool context path writes UUID objects into sensory memory.
   In python/flink_agents/plan/actions/chat_model_action.py, tool call context uses initial_request_id: UUID / tool_request_event_id: UUID in dictionaries stored in ctx.sensory_memory. For example:

   • _update_tool_call_context(...) stores a dict keyed by initial_request_id.
   • _save_tool_request_event_context(...) stores "initial_request_id": initial_request_id and stores the context under tool_request_event_id.

With Pemja, uuid.UUID does not have a dedicated Python-to-Java conversion branch, so it falls back to a PyObject wrapper. That makes the built-in tool-call flow capable of storing a non-checkpoint-stable payload in Flink state.

Expected behavior:

• The Python memory value contract should be explicit and enforceable.
• Built-in runtime memory payloads should not store raw Python objects such as uuid.UUID.
• Values stored in Flink-backed memory should be safe to checkpoint and restore across TaskManager / Python process restarts.

A reasonable Python-side contract could be to accept only recursively checkpoint-stable values such as:

None | bool | int | float | str | bytes | list[MemoryValue] | dict[str, MemoryValue]

Pydantic values should be materialized before storing, for example:

ctx.sensory_memory.set("x", model.model_dump(mode="json"))

or:

ctx.sensory_memory.set("x", json.loads(model.model_dump_json()))

uuid.UUID should be stored as str(uuid_value) if it needs to be persisted.

How to reproduce

Use a Python Flink Agents streaming job with a real Flink state backend and checkpointing enabled. The important part is to test a real Flink Agent job path, not only direct Flink serializer behavior.

1. Configure a Flink Agent job with keyed state, RocksDB state backend, filesystem checkpointing, and checkpointing enabled.
2. In one action, write a Python dict containing a uuid.UUID to sensory memory:

from uuid import UUID

ctx.sensory_memory.set(
    "restore.payload",
    {
        "request_id": UUID("00000000-0000-0000-0000-000000000001"),
        "metadata": {"input_id": 1},
    },
)

3. Trigger a follow-up action in the same agent run.
4. Wait until a checkpoint containing the memory state completes.
5. Kill the TaskManager and start a replacement TaskManager so the job restores from the checkpoint.
6. In the follow-up action, read ctx.sensory_memory.get("restore.payload").

Observed behavior in a local repro:

• RocksDB state restoration completes.
• The restored action is re-executed after checkpoint restore.
• When the restored memory value is converted back to Python, the JVM crashes in Pemja native code.

Relevant log snippets:

Finished RocksDB incremental recovery in operator ActionExecutionOperator ...
Python awaitable ref is null for action validate_after_tm_restart (likely restored from checkpoint), re-executing from beginning.
SIGSEGV ... pemja_core.cpython-311-darwin.so ... JcpPyObject_FromJObject
JcpPyDict_FromJMap

As a control test, storing a plain string in sensory memory through the same real Flink Agent + RocksDB + checkpoint + TaskManager restart path succeeds. The value written by the old TaskManager is read successfully by the new TaskManager after restore.

This suggests the issue is not that sensory memory is never persisted. The issue is specifically around Python-origin objects that Pemja cannot materialize into checkpoint-stable Java values before they are stored in Flink state.

Version and environment

• Flink Agents: current main branch local checkout
• Python: 3.11
• Pemja: 0.5.5
• Java: 11
• Flink/PyFlink: local mini-cluster using a PyFlink 2.2.0 distribution
• State backend used for the local repro: RocksDB with filesystem checkpoint storage
• Deployment mode: local standalone JobManager / TaskManager, manually killing and restarting the TaskManager to verify checkpoint restore behavior

Are you willing to submit a PR?

• I'm willing to submit a PR!

[weiqingy]

@joeyutong Thanks for the detailed write-up and the repro. I traced both code paths and they line up with what you describe:

• _update_tool_call_context keys the stored dict by a UUID (tool_call_context[initial_request_id]) at chat_model_action.py:77-84.
• _save_tool_request_event_context stores "initial_request_id": initial_request_id (a UUID) keyed under another UUID at chat_model_action.py:96-102.

And FlinkMemoryObject.set() passes the raw Python value straight to the Java MemoryObjectImpl.set() via Pemja (flink_memory_object.py:68), so anything Pemja can't materialize ends up as a PyObject/PyJObject wrapper inside the checkpointed Java state — consistent with the SIGSEGV in JcpPyObject_FromJObject on restore.

One thing that may be worth folding into the scope: the same _save_tool_request_event_context dict also stores "output_schema": output_schema at chat_model_action.py:100, and OutputSchema is a Pydantic BaseModel (api/agents/types.py:25). That has the same non-materialization problem as UUID, so it looks like a second instance of this class of bug in the built-in path — not just the two ID fields.

A small note on the suggested fix: in _update_tool_call_context the UUID is used as a dict key, not just a value, so normalizing would need to cover keys too (e.g. str(initial_request_id)), otherwise the Pemja wrapper just moves from the value side to the key side.

On the broader contract piece — do you have a preference between (a) only documenting that Python memory values must be recursively checkpoint-stable, vs. (b) validating at set() time and rejecting non-stable values with a clear error before they reach state? Option (b) turns a checkpoint-time JVM crash into an early, debuggable failure, but it's a behavior change, so it seems worth getting other opinions before settling on it. WDYT?

[joeyutong]

@weiqingy Sorry for the late reply, and thanks for the detailed follow-up.

I agree with the proposed fix for the built-in tool context path. The framework-owned payloads should be normalized in code, including UUIDs used as dict keys/values and output_schema, so they are safe to checkpoint and restore.

For Python memory values, I think we should update the docs to clarify two things: which Python-side value shapes are compatible, and that this contract is not fully aligned with the Java-side one because Python values need to be materialized across the Pemja boundary first.

I also think adding a Python-side validator with a clear error message is reasonable. It is a behavior change, but at this early stage of the project, I think failing early is better than letting unsupported values reach Flink state and fail during restore.

[weiqingy]

@joeyutong Thanks for confirming — sounds like we're fully aligned. Capturing the agreed scope so it's on record:

1. Normalize the framework-owned tool-context payloads in chat_model_action.py so they're checkpoint-safe — covering the UUID used as a dict key as well as the value side, plus output_schema (the Pydantic OutputSchema instance).
2. Document the Python memory value contract: which value shapes are checkpoint-compatible, and that it intentionally diverges from the Java-side contract because Python values have to be materialized across the Pemja boundary before they reach state.
3. Add a Python-side validator at set() that rejects non-checkpoint-stable values early with a clear error, instead of letting them fail during restore. Agreed that failing early is the better trade-off at this stage.

Since you've left the PR box unchecked, I'm happy to pick this up. My plan is to land the built-in tool-context fix first as a small self-contained change (it's independent of the contract decision), then follow with the validator and the docs. Does that sequencing work for you?

[joeyutong]

@weiqingy Thanks for taking this on, I really appreciate it. The proposed sequencing works for me.

[wenjin272]

Thank you, @joeyutong , for the detailed report, and @weiqingy , for the proposal. Since @weiqingy is willing to take this on and the proposal has been agreed upon, I've assigned the issue to @weiqingy.