fix(middleware): fallback to configurable thread_id in thread data middleware (#1237)

Co-authored-by: Exploreunive <Exploreunive@users.noreply.github.com>
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
haoliangxu
2026-03-22 20:14:51 +08:00
committed by GitHub
parent 894875ab1b
commit e6c6770b70
2 changed files with 62 additions and 2 deletions

View File

@@ -2,6 +2,7 @@ from typing import NotRequired, override
from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langgraph.config import get_config
from langgraph.runtime import Runtime
from deerflow.agents.thread_state import ThreadDataState
@@ -71,9 +72,14 @@ class ThreadDataMiddleware(AgentMiddleware[ThreadDataMiddlewareState]):
@override
def before_agent(self, state: ThreadDataMiddlewareState, runtime: Runtime) -> dict | None:
thread_id = runtime.context.get("thread_id")
context = runtime.context or {}
thread_id = context.get("thread_id")
if thread_id is None:
raise ValueError("Thread ID is required in the context")
config = get_config()
thread_id = config.get("configurable", {}).get("thread_id")
if thread_id is None:
raise ValueError("Thread ID is required in runtime context or config.configurable")
if self._lazy_init:
# Lazy initialization: only compute paths, don't create directories

View File

@@ -0,0 +1,54 @@
import pytest
from langgraph.runtime import Runtime
from deerflow.agents.middlewares.thread_data_middleware import ThreadDataMiddleware
class TestThreadDataMiddleware:
def test_before_agent_returns_paths_when_thread_id_present_in_context(self, tmp_path):
middleware = ThreadDataMiddleware(base_dir=str(tmp_path), lazy_init=True)
result = middleware.before_agent(state={}, runtime=Runtime(context={"thread_id": "thread-123"}))
assert result is not None
assert result["thread_data"]["workspace_path"].endswith("threads/thread-123/user-data/workspace")
assert result["thread_data"]["uploads_path"].endswith("threads/thread-123/user-data/uploads")
assert result["thread_data"]["outputs_path"].endswith("threads/thread-123/user-data/outputs")
def test_before_agent_uses_thread_id_from_configurable_when_context_is_none(self, tmp_path, monkeypatch):
middleware = ThreadDataMiddleware(base_dir=str(tmp_path), lazy_init=True)
runtime = Runtime(context=None)
monkeypatch.setattr(
"deerflow.agents.middlewares.thread_data_middleware.get_config",
lambda: {"configurable": {"thread_id": "thread-from-config"}},
)
result = middleware.before_agent(state={}, runtime=runtime)
assert result is not None
assert result["thread_data"]["workspace_path"].endswith("threads/thread-from-config/user-data/workspace")
assert runtime.context is None
def test_before_agent_uses_thread_id_from_configurable_when_context_missing_thread_id(self, tmp_path, monkeypatch):
middleware = ThreadDataMiddleware(base_dir=str(tmp_path), lazy_init=True)
runtime = Runtime(context={})
monkeypatch.setattr(
"deerflow.agents.middlewares.thread_data_middleware.get_config",
lambda: {"configurable": {"thread_id": "thread-from-config"}},
)
result = middleware.before_agent(state={}, runtime=runtime)
assert result is not None
assert result["thread_data"]["uploads_path"].endswith("threads/thread-from-config/user-data/uploads")
assert runtime.context == {}
def test_before_agent_raises_clear_error_when_thread_id_missing_everywhere(self, tmp_path, monkeypatch):
middleware = ThreadDataMiddleware(base_dir=str(tmp_path), lazy_init=True)
monkeypatch.setattr(
"deerflow.agents.middlewares.thread_data_middleware.get_config",
lambda: {"configurable": {}},
)
with pytest.raises(ValueError, match="Thread ID is required in runtime context or config.configurable"):
middleware.before_agent(state={}, runtime=Runtime(context=None))