INTPYTHON-728 Fix handling of metadata in saver (#195)

Author: blink1073

libs/langgraph-checkpoint-mongodb/langgraph/checkpoint/mongodb/aio.py

@@ -357,6 +357,8 @@ async def aput(
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         checkpoint_id = checkpoint["id"]
         type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
+        metadata = metadata.copy()
+        metadata.update(config.get("metadata", {}))
         doc = {
             "parent_checkpoint_id": config["configurable"].get("checkpoint_id"),
             "type": type_,
@@ -368,6 +370,7 @@ async def aput(
             "checkpoint_ns": checkpoint_ns,
             "checkpoint_id": checkpoint_id,
         }
+
         if self.ttl:
             doc["created_at"] = datetime.now()
         # Perform your operations here

libs/langgraph-checkpoint-mongodb/langgraph/checkpoint/mongodb/saver.py

@@ -376,6 +376,8 @@ def put(
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         checkpoint_id = checkpoint["id"]
         type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
+        metadata = metadata.copy()
+        metadata.update(config.get("metadata", {}))
         doc = {
             "parent_checkpoint_id": config["configurable"].get("checkpoint_id"),
             "type": type_,

libs/langgraph-checkpoint-mongodb/tests/integration_tests/test_highlevel_graph.py

@@ -87,7 +87,7 @@ def fanout(state: OverallState) -> list:
         return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
 
     parentgraph = StateGraph(OverallState)
-    parentgraph.add_node("generate_joke", subgraphc)  # type: ignore[arg-type]
+    parentgraph.add_node("generate_joke", subgraphc)
     parentgraph.add_conditional_edges(START, fanout)
     parentgraph.add_edge("generate_joke", END)
     return parentgraph
@@ -176,3 +176,69 @@ async def test_fanout(
         )
         end = time.monotonic()
         print(f"{cname}: {end - start:.4f} seconds")
+
+
+async def test_custom_properties_async(
+    checkpointer_mongodb: MongoDBSaver, checkpointer_mongodb_async: AsyncMongoDBSaver
+) -> None:
+    # Create the state graph
+    state_graph = fanout_to_subgraph()
+
+    # Define configuration with thread ID and assistant ID
+    assistant_id = "456"
+    user_id = "789"
+    config: RunnableConfig = {
+        "configurable": {
+            "thread_id": "123",
+            "assistant_id": assistant_id,
+            "user_id": user_id,
+        }
+    }
+
+    # Compile the state graph with the provided checkpointing mechanism
+    compiled_state_graph = state_graph.compile(checkpointer=checkpointer_mongodb_async)
+
+    # Invoke the compiled state graph with user input
+    await compiled_state_graph.ainvoke(
+        input={"subjects": [], "step": 0},  # type:ignore[arg-type]
+        config=config,
+        stream_mode="values",
+        debug=False,
+    )
+
+    checkpoint_tuple = await checkpointer_mongodb_async.aget_tuple(config)
+    assert checkpoint_tuple is not None
+    assert checkpoint_tuple.metadata["user_id"] == user_id
+    assert checkpoint_tuple.metadata["assistant_id"] == assistant_id
+
+
+def test_custom_properties(checkpointer_mongodb: MongoDBSaver) -> None:
+    # Create the state graph
+    state_graph = fanout_to_subgraph()
+
+    # Define configuration with thread ID and assistant ID
+    assistant_id = "456"
+    user_id = "789"
+    config: RunnableConfig = {
+        "configurable": {
+            "thread_id": "123",
+            "assistant_id": assistant_id,
+            "user_id": user_id,
+        }
+    }
+
+    # Compile the state graph with the provided checkpointing mechanism
+    compiled_state_graph = state_graph.compile(checkpointer=checkpointer_mongodb)
+
+    # Invoke the compiled state graph with user input
+    compiled_state_graph.invoke(
+        input={"subjects": [], "step": 0},  # type:ignore[arg-type]
+        config=config,
+        stream_mode="values",
+        debug=False,
+    )
+
+    checkpoint_tuple = checkpointer_mongodb.get_tuple(config)
+    assert checkpoint_tuple is not None
+    assert checkpoint_tuple.metadata["user_id"] == user_id
+    assert checkpoint_tuple.metadata["assistant_id"] == assistant_id