From 57d36a0977c9229abbe328788ba6bc214bcb2fb5 Mon Sep 17 00:00:00 2001
From: greatmengqi <chenmengqi.0376@bytedance.com>
Date: Sun, 26 Apr 2026 12:42:17 +0800
Subject: [PATCH] test(thread_data_middleware): cover run_id stamping on
 HumanMessage

Two new cases in test_thread_data_middleware.py pin the merge contract
between PR (frozen DeerFlowContext) and release/2.0-rc (HumanMessage
stamping logic):

- test_before_agent_stamps_run_id_and_timestamp_on_last_human_message:
  given DeerFlowContext(run_id="r-stamp"), the trailing HumanMessage
  in state gets additional_kwargs["run_id"] == "r-stamp" and a
  timestamp; id and name are preserved/normalized.

- test_before_agent_stamps_none_run_id_when_context_omits_it:
  run_id is optional on DeerFlowContext, so middleware must stamp
  None gracefully rather than raise.

These were the smoke target for the merge: prior to this commit
the only signal that runtime.context.run_id worked end-to-end was
a manual gateway run.
---
 backend/tests/test_thread_data_middleware.py | 42 +++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/backend/tests/test_thread_data_middleware.py b/backend/tests/test_thread_data_middleware.py
index 4cc289b2d..7336d5c57 100644
--- a/backend/tests/test_thread_data_middleware.py
+++ b/backend/tests/test_thread_data_middleware.py
@@ -10,10 +10,11 @@ def _as_posix(path: str) -> str:
     return path.replace("\\", "/")
 
 
-def _make_context(thread_id: str) -> DeerFlowContext:
+def _make_context(thread_id: str, run_id: str | None = None) -> DeerFlowContext:
     return DeerFlowContext(
         app_config=AppConfig(sandbox=SandboxConfig(use="test")),
         thread_id=thread_id,
+        run_id=run_id,
     )
 
 
@@ -53,3 +54,42 @@ class TestThreadDataMiddleware:
 
         with pytest.raises(ValueError, match="Thread ID is required"):
             middleware.before_agent(state={}, runtime=Runtime(context=_make_context("")))
+
+    def test_before_agent_stamps_run_id_and_timestamp_on_last_human_message(self, tmp_path):
+        """Smoke for the release/2.0-rc + PR merge: run_id from typed
+        DeerFlowContext flows into the trailing HumanMessage's
+        additional_kwargs alongside an ISO-8601 timestamp."""
+        from langchain_core.messages import HumanMessage
+        from langgraph.runtime import Runtime
+
+        middleware = ThreadDataMiddleware(base_dir=str(tmp_path), lazy_init=True)
+        original = HumanMessage(content="hello", id="m-1")
+
+        result = middleware.before_agent(
+            state={"messages": [original]},
+            runtime=Runtime(context=_make_context("t-stamp", run_id="r-stamp")),
+        )
+
+        assert result is not None
+        stamped = result["messages"][-1]
+        assert isinstance(stamped, HumanMessage)
+        assert stamped.id == "m-1"
+        assert stamped.name == "user-input"
+        assert stamped.additional_kwargs["run_id"] == "r-stamp"
+        assert "timestamp" in stamped.additional_kwargs
+
+    def test_before_agent_stamps_none_run_id_when_context_omits_it(self, tmp_path):
+        """run_id is optional: middleware must still stamp (with None) rather than crash."""
+        from langchain_core.messages import HumanMessage
+        from langgraph.runtime import Runtime
+
+        middleware = ThreadDataMiddleware(base_dir=str(tmp_path), lazy_init=True)
+
+        result = middleware.before_agent(
+            state={"messages": [HumanMessage(content="hi", id="m-2")]},
+            runtime=Runtime(context=_make_context("t-no-run")),
+        )
+
+        assert result is not None
+        stamped = result["messages"][-1]
+        assert stamped.additional_kwargs["run_id"] is None