mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 02:05:45 +00:00
feat(context): record effective memory identity per run (#3556)
* feat(context): record effective memory identity per run * fix(context): address memory identity review feedback --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
co-authored by
Willem Jiang
parent
8da7cbf028
commit
94a34f382d
@@ -584,6 +584,11 @@ The cached value is reused for both the blocking (`runs.wait`) and streaming (`_
|
||||
- Consolidation pass (same LLM invocation as the regular updater, no extra API call): when `consolidation_enabled` is `true` and at least one category holds `consolidation_min_facts` or more facts, `_select_consolidation_candidates` identifies fragmented categories and surfaces at most `consolidation_max_groups_per_cycle` of them (largest first) in the prompt. The LLM decides which groups to merge and proposes a synthesised fact per group. `_apply_updates` enforces guardrails: source IDs must exist and must not overlap across groups, group size is capped at `consolidation_max_sources`, the merged fact's confidence cannot exceed the source maximum, and facts below `fact_confidence_threshold` are not written.
|
||||
- Next interaction injects selected facts + context into `<memory>` tags in the system prompt when `injection_enabled` is true.
|
||||
|
||||
**Run-level memory identity**:
|
||||
- Every Gateway run with an effective hidden memory block hashes the exact `HumanMessage.content`, including the `<memory>` wrapper, and records one `context:memory` event through its run-scoped `RunJournal`. Later runs and checkpoint-based branches reuse the frozen message without reloading memory; goal continuations are deduplicated to one event per run.
|
||||
- A first-run block is trusted only when it comes from `DynamicContextMiddleware`'s current update. A reused block must have existed in the checkpoint before the run, and the Gateway strips dynamic-context markers from untrusted input so a caller cannot forge the identity event by reusing a known message ID.
|
||||
- The production consumer is the existing debug/audit endpoint `GET /api/threads/{thread_id}/runs/{run_id}/events?event_types=context:memory`. Event content has exactly one field, `content_sha256`, which operators use to compare the effective memory identity across runs. The full memory text stays in checkpoint state and is not duplicated into `run_events`.
|
||||
|
||||
**Token counting** (`packages/harness/deerflow/agents/memory/prompt.py`):
|
||||
- `_count_tokens` budgets the injection. In default `tiktoken` mode, the encoding is loaded lazily and cached.
|
||||
- Failed tiktoken loads are cached with a timestamp. During the fixed cooldown (`_TIKTOKEN_RETRY_COOLDOWN_S`, 600s), callers fall back to char estimation immediately instead of re-triggering the blocking BPE download; after the cooldown, transient outages can self-heal without a restart.
|
||||
|
||||
@@ -94,6 +94,7 @@ LLM-powered persistent context retention across conversations:
|
||||
- **Structured storage**: User context (work, personal, top-of-mind), history, and confidence-scored facts
|
||||
- **Debounced updates**: Batches updates to minimize LLM calls (configurable wait time)
|
||||
- **System prompt injection**: Top facts + context injected into agent prompts
|
||||
- **Run-level memory identity**: `GET /api/threads/{thread_id}/runs/{run_id}/events?event_types=context:memory` returns the SHA-256 identity of the effective hidden memory block without copying memory text into the event store
|
||||
- **Storage**: JSON file with mtime-based cache invalidation
|
||||
|
||||
### Tool Ecosystem
|
||||
@@ -121,6 +122,7 @@ FastAPI application providing REST endpoints for frontend integration:
|
||||
| `POST /api/memory/reload` | Force memory reload |
|
||||
| `GET /api/memory/config` | Memory configuration |
|
||||
| `GET /api/memory/status` | Combined config + data |
|
||||
| `GET /api/threads/{id}/runs/{run_id}/events` | Debug/audit events for one run; filter `event_types=context:memory` for effective memory identity |
|
||||
| `POST /api/threads/{id}/uploads` | Upload files (auto-converts PDF/PPT/Excel/Word to Markdown, rejects directory paths, auto-renames duplicate filenames in one request) |
|
||||
| `GET /api/threads/{id}/uploads/list` | List uploaded files |
|
||||
| `DELETE /api/threads/{id}` | Delete DeerFlow-managed local thread data after LangGraph thread deletion; unexpected failures are logged server-side and return a generic 500 detail |
|
||||
|
||||
@@ -29,6 +29,7 @@ from app.gateway.internal_auth import (
|
||||
get_trusted_internal_owner_user_id,
|
||||
)
|
||||
from app.gateway.utils import sanitize_log_param
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, _REMINDER_DATE_KEY
|
||||
from deerflow.config.app_config import get_app_config
|
||||
from deerflow.runtime import (
|
||||
END_SENTINEL,
|
||||
@@ -57,6 +58,13 @@ _TERMINAL_RUN_STATUSES = {
|
||||
RunStatus.interrupted,
|
||||
}
|
||||
|
||||
_SERVER_OWNED_DYNAMIC_CONTEXT_KEYS = frozenset(
|
||||
{
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY,
|
||||
_REMINDER_DATE_KEY,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SSE formatting
|
||||
@@ -120,10 +128,14 @@ def normalize_stream_modes(raw: list[str] | str | None) -> list[str]:
|
||||
|
||||
def _strip_external_message_metadata(message: Any) -> Any:
|
||||
"""Remove server-owned metadata from an untrusted input message."""
|
||||
if not isinstance(message, BaseMessage) or ORIGINAL_USER_CONTENT_KEY not in message.additional_kwargs:
|
||||
if not isinstance(message, BaseMessage):
|
||||
return message
|
||||
additional_kwargs = dict(message.additional_kwargs)
|
||||
additional_kwargs.pop(ORIGINAL_USER_CONTENT_KEY, None)
|
||||
for key in _SERVER_OWNED_DYNAMIC_CONTEXT_KEYS:
|
||||
additional_kwargs.pop(key, None)
|
||||
if additional_kwargs == message.additional_kwargs:
|
||||
return message
|
||||
return message.model_copy(update={"additional_kwargs": additional_kwargs})
|
||||
|
||||
|
||||
@@ -141,10 +153,9 @@ def normalize_input(raw_input: dict[str, Any] | None, *, trusted_internal: bool
|
||||
of bubbling up as a 500. The gateway is a system boundary, so per-entry
|
||||
validation errors are the right shape for clients to retry against.
|
||||
|
||||
``original_user_content`` is server-owned provenance used to undo model-only
|
||||
sanitization at persistence time. External callers cannot supply it; trusted
|
||||
internal channel calls may preserve the value they captured before adding
|
||||
transport or file context.
|
||||
``original_user_content`` and dynamic-context reminder markers are
|
||||
server-owned. External callers cannot supply them; trusted internal channel
|
||||
calls may preserve metadata they added before invoking this boundary.
|
||||
"""
|
||||
if raw_input is None:
|
||||
return {}
|
||||
|
||||
@@ -29,6 +29,7 @@ Date-update format:
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
@@ -39,6 +40,8 @@ from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from deerflow.config.app_config import AppConfig
|
||||
|
||||
@@ -279,7 +282,9 @@ class DynamicContextMiddleware(AgentMiddleware):
|
||||
|
||||
@override
|
||||
def before_agent(self, state, runtime: Runtime) -> dict | None:
|
||||
return self._inject(state)
|
||||
result = self._inject(state)
|
||||
self._record_effective_memory(state, result, runtime)
|
||||
return result
|
||||
|
||||
@override
|
||||
async def abefore_agent(self, state, runtime: Runtime) -> dict | None:
|
||||
@@ -292,16 +297,69 @@ class DynamicContextMiddleware(AgentMiddleware):
|
||||
# Bounded timeout: if startup warm-up failed silently (e.g. network
|
||||
# blip during deploy), the first request's cold tiktoken download can
|
||||
# block for tens of minutes (OS TCP timeout). Time-box injection so
|
||||
# the request degrades gracefully (no memory context) rather than
|
||||
# hanging.
|
||||
# the request degrades gracefully (no new dynamic-context update)
|
||||
# rather than hanging. Frozen context already in state remains active.
|
||||
try:
|
||||
return await asyncio.wait_for(
|
||||
result = await asyncio.wait_for(
|
||||
asyncio.to_thread(self._inject, state),
|
||||
timeout=_INJECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"DynamicContextMiddleware: injection timed out (%.1fs); skipping memory/date injection for this turn",
|
||||
"DynamicContextMiddleware: injection timed out (%.1fs); skipping new memory/date injection for this turn",
|
||||
_INJECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
self._record_effective_memory(state, None, runtime)
|
||||
return None
|
||||
self._record_effective_memory(state, result, runtime)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _effective_memory_message(state, update: dict | None, runtime: Runtime) -> HumanMessage | None:
|
||||
"""Find server-created memory that is effective for this run.
|
||||
|
||||
A first-run block must come from this middleware's update. A reused
|
||||
block must have existed in the checkpoint before the run; the Gateway
|
||||
strips the reminder marker from untrusted input so a caller cannot
|
||||
replace a known checkpoint ID with forged provenance.
|
||||
"""
|
||||
if isinstance(update, dict):
|
||||
update_messages = update.get("messages")
|
||||
if isinstance(update_messages, list):
|
||||
for message in update_messages:
|
||||
if not isinstance(message, HumanMessage):
|
||||
continue
|
||||
message_id = str(message.id or "")
|
||||
if message_id.endswith("__memory") and is_dynamic_context_reminder(message) and isinstance(message.content, str):
|
||||
return message
|
||||
|
||||
context = getattr(runtime, "context", None)
|
||||
raw_pre_existing_ids = context.get(CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY) if isinstance(context, dict) else None
|
||||
if not isinstance(raw_pre_existing_ids, (frozenset, set, list, tuple)):
|
||||
return None
|
||||
pre_existing_ids = {str(message_id) for message_id in raw_pre_existing_ids if message_id}
|
||||
for message in state.get("messages", []):
|
||||
if not isinstance(message, HumanMessage):
|
||||
continue
|
||||
message_id = str(message.id or "")
|
||||
if message_id in pre_existing_ids and message_id.endswith("__memory") and is_dynamic_context_reminder(message) and isinstance(message.content, str):
|
||||
return message
|
||||
return None
|
||||
|
||||
def _record_effective_memory(self, state, update: dict | None, runtime: Runtime) -> None:
|
||||
"""Attach the effective hidden memory block to the current run ledger."""
|
||||
context = getattr(runtime, "context", None)
|
||||
journal = context.get("__run_journal") if isinstance(context, dict) else None
|
||||
if journal is None:
|
||||
return
|
||||
|
||||
message = self._effective_memory_message(state, update, runtime)
|
||||
if message is None:
|
||||
return
|
||||
|
||||
try:
|
||||
journal.record_memory_context(
|
||||
content_sha256=hashlib.sha256(message.content.encode("utf-8")).hexdigest(),
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to record effective memory context", exc_info=True)
|
||||
|
||||
@@ -102,6 +102,7 @@ class RunJournal(BaseCallbackHandler):
|
||||
self._counted_llm_run_ids: set[str] = set()
|
||||
self._counted_external_source_ids: set[str] = set()
|
||||
self._counted_message_llm_run_ids: set[str] = set()
|
||||
self._memory_context_recorded = False
|
||||
|
||||
# Convenience fields
|
||||
self._last_ai_msg: str | None = None
|
||||
@@ -630,6 +631,23 @@ class RunJournal(BaseCallbackHandler):
|
||||
content={"name": name, "hook": hook, "action": action, "changes": changes},
|
||||
)
|
||||
|
||||
def record_memory_context(self, *, content_sha256: str) -> None:
|
||||
"""Record the effective hidden memory block for this run.
|
||||
|
||||
The full block already lives in checkpoint state and may contain user
|
||||
data, so the event stores only its exact SHA-256 identity. Operators
|
||||
consume it through the existing run-events debug API to compare the
|
||||
effective memory used by different runs without copying that content.
|
||||
"""
|
||||
if self._memory_context_recorded:
|
||||
return
|
||||
self._put(
|
||||
event_type="context:memory",
|
||||
category="context",
|
||||
content={"content_sha256": content_sha256},
|
||||
)
|
||||
self._memory_context_recorded = True
|
||||
|
||||
async def flush(self) -> None:
|
||||
"""Force flush remaining buffer. Called in worker's finally block."""
|
||||
if self._pending_flush_tasks:
|
||||
|
||||
@@ -14,6 +14,8 @@ this test fails.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
@@ -22,7 +24,11 @@ from langchain.agents import create_agent
|
||||
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import DynamicContextMiddleware
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import (
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY,
|
||||
DynamicContextMiddleware,
|
||||
)
|
||||
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
@@ -101,24 +107,105 @@ async def test_abefore_agent_returns_same_result_as_before_agent() -> None:
|
||||
|
||||
|
||||
async def test_abefore_agent_returns_none_on_timeout() -> None:
|
||||
"""If _inject() exceeds the timeout, abefore_agent returns None gracefully."""
|
||||
import time
|
||||
|
||||
"""A timed-out worker must not emit a late, phantom context event."""
|
||||
mw = DynamicContextMiddleware()
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
finished = threading.Event()
|
||||
journal = mock.MagicMock()
|
||||
|
||||
def blocking_inject(state):
|
||||
time.sleep(10) # Simulate a blocking call that far exceeds the timeout
|
||||
return {"messages": [HumanMessage(content="should not reach")]}
|
||||
started.set()
|
||||
release.wait(timeout=2)
|
||||
try:
|
||||
return {
|
||||
"messages": [
|
||||
HumanMessage(
|
||||
content="<memory>late context</memory>",
|
||||
id="msg-1__memory",
|
||||
additional_kwargs={
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY: True,
|
||||
},
|
||||
)
|
||||
]
|
||||
}
|
||||
finally:
|
||||
finished.set()
|
||||
|
||||
with (
|
||||
mock.patch.object(mw, "_inject", blocking_inject),
|
||||
mock.patch(
|
||||
"deerflow.agents.middlewares.dynamic_context_middleware._INJECT_TIMEOUT_SECONDS",
|
||||
0.1,
|
||||
0.01,
|
||||
),
|
||||
):
|
||||
state = {"messages": [HumanMessage(content="Hello", id="msg-1")]}
|
||||
runtime = SimpleNamespace(context={})
|
||||
runtime = SimpleNamespace(context={"__run_journal": journal})
|
||||
result = await mw.abefore_agent(state, runtime)
|
||||
|
||||
assert started.is_set()
|
||||
assert result is None
|
||||
release.set()
|
||||
assert await asyncio.to_thread(finished.wait, 1)
|
||||
journal.record_memory_context.assert_not_called()
|
||||
|
||||
|
||||
async def test_abefore_agent_records_checkpointed_memory_on_timeout() -> None:
|
||||
"""A timeout does not hide frozen memory that remains effective for the run."""
|
||||
mw = DynamicContextMiddleware()
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
finished = threading.Event()
|
||||
journal = mock.MagicMock()
|
||||
memory_content = "<memory>checkpoint context</memory>"
|
||||
|
||||
def blocking_inject(state):
|
||||
started.set()
|
||||
release.wait(timeout=2)
|
||||
try:
|
||||
return {
|
||||
"messages": [
|
||||
HumanMessage(
|
||||
content="<memory>late replacement</memory>",
|
||||
id="msg-2__memory",
|
||||
additional_kwargs={_DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
]
|
||||
}
|
||||
finally:
|
||||
finished.set()
|
||||
|
||||
state = {
|
||||
"messages": [
|
||||
HumanMessage(
|
||||
content=memory_content,
|
||||
id="msg-1__memory",
|
||||
additional_kwargs={_DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
]
|
||||
}
|
||||
runtime = SimpleNamespace(
|
||||
context={
|
||||
"__run_journal": journal,
|
||||
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: frozenset({"msg-1__memory"}),
|
||||
}
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch.object(mw, "_inject", blocking_inject),
|
||||
mock.patch(
|
||||
"deerflow.agents.middlewares.dynamic_context_middleware._INJECT_TIMEOUT_SECONDS",
|
||||
0.01,
|
||||
),
|
||||
):
|
||||
result = await mw.abefore_agent(state, runtime)
|
||||
|
||||
recorded_call = journal.record_memory_context.call_args
|
||||
release.set()
|
||||
assert await asyncio.to_thread(finished.wait, 1)
|
||||
assert started.is_set()
|
||||
assert result is None
|
||||
assert recorded_call == mock.call(
|
||||
content_sha256=hashlib.sha256(memory_content.encode("utf-8")).hexdigest(),
|
||||
)
|
||||
journal.record_memory_context.assert_called_once()
|
||||
|
||||
@@ -4,6 +4,7 @@ Verifies that memory and current date are injected as a <system-reminder> into
|
||||
the first HumanMessage exactly once per session (frozen-snapshot pattern).
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
@@ -13,6 +14,7 @@ from deerflow.agents.middlewares.dynamic_context_middleware import (
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY,
|
||||
DynamicContextMiddleware,
|
||||
)
|
||||
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY
|
||||
|
||||
_SYSTEM_REMINDER_TAG = "<system-reminder>"
|
||||
|
||||
@@ -21,8 +23,11 @@ def _make_middleware(**kwargs) -> DynamicContextMiddleware:
|
||||
return DynamicContextMiddleware(**kwargs)
|
||||
|
||||
|
||||
def _fake_runtime():
|
||||
return SimpleNamespace(context={})
|
||||
def _fake_runtime(journal=None, *, pre_existing_message_ids=()):
|
||||
context = {"__run_journal": journal} if journal is not None else {}
|
||||
if pre_existing_message_ids:
|
||||
context[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] = frozenset(pre_existing_message_ids)
|
||||
return SimpleNamespace(context=context)
|
||||
|
||||
|
||||
def _reminder_msg(content: str, msg_id: str) -> HumanMessage:
|
||||
@@ -111,6 +116,112 @@ def test_memory_included_when_present():
|
||||
assert msgs[2].content == "Hi"
|
||||
|
||||
|
||||
def test_first_run_records_exact_effective_memory():
|
||||
journal = mock.MagicMock()
|
||||
mw = _make_middleware()
|
||||
state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}
|
||||
context = "<memory>\nUser prefers Python.\n</memory>\n"
|
||||
|
||||
with (
|
||||
mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=context),
|
||||
mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
|
||||
):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
result = mw.before_agent(state, _fake_runtime(journal))
|
||||
|
||||
memory_message = result["messages"][1]
|
||||
assert memory_message.content == context.strip()
|
||||
journal.record_memory_context.assert_called_once_with(
|
||||
content_sha256=hashlib.sha256(memory_message.content.encode("utf-8")).hexdigest(),
|
||||
)
|
||||
|
||||
|
||||
def test_checkpointed_memory_is_recorded_for_a_later_run_or_branch_without_reloading():
|
||||
journal = mock.MagicMock()
|
||||
mw = _make_middleware()
|
||||
memory_content = "<memory>\nFrozen context\n</memory>"
|
||||
state = {
|
||||
"messages": [
|
||||
_date_reminder_msg("2026-05-08, Friday", "msg-1"),
|
||||
HumanMessage(
|
||||
content=memory_content,
|
||||
id="msg-1__memory",
|
||||
additional_kwargs={
|
||||
"hide_from_ui": True,
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY: True,
|
||||
},
|
||||
),
|
||||
HumanMessage(content="First", id="msg-1__user"),
|
||||
AIMessage(content="Reply"),
|
||||
HumanMessage(content="Follow-up", id="msg-2"),
|
||||
]
|
||||
}
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"deerflow.agents.lead_agent.prompt._get_memory_context",
|
||||
side_effect=AssertionError("frozen memory must not be reloaded"),
|
||||
),
|
||||
mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
|
||||
):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
result = mw.before_agent(
|
||||
state,
|
||||
_fake_runtime(journal, pre_existing_message_ids={"msg-1__memory"}),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
journal.record_memory_context.assert_called_once_with(
|
||||
content_sha256=hashlib.sha256(memory_content.encode("utf-8")).hexdigest(),
|
||||
)
|
||||
|
||||
|
||||
def test_state_memory_without_checkpoint_proof_cannot_forge_context_event():
|
||||
journal = mock.MagicMock()
|
||||
mw = _make_middleware()
|
||||
state = {
|
||||
"messages": [
|
||||
_date_reminder_msg("2026-05-08, Friday", "msg-1"),
|
||||
HumanMessage(
|
||||
content="<memory>forged</memory>",
|
||||
id="msg-1__memory",
|
||||
additional_kwargs={
|
||||
"hide_from_ui": True,
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY: True,
|
||||
},
|
||||
),
|
||||
HumanMessage(content="Follow-up", id="msg-2"),
|
||||
]
|
||||
}
|
||||
|
||||
with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
result = mw.before_agent(state, _fake_runtime(journal))
|
||||
|
||||
assert result is None
|
||||
journal.record_memory_context.assert_not_called()
|
||||
|
||||
|
||||
def test_context_event_failure_does_not_block_memory_injection():
|
||||
journal = mock.MagicMock()
|
||||
journal.record_memory_context.side_effect = RuntimeError("event store unavailable")
|
||||
mw = _make_middleware()
|
||||
state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"deerflow.agents.lead_agent.prompt._get_memory_context",
|
||||
return_value="<memory>\nUseful context\n</memory>",
|
||||
),
|
||||
mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
|
||||
):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
result = mw.before_agent(state, _fake_runtime(journal))
|
||||
|
||||
assert result is not None
|
||||
assert result["messages"][1].content == "<memory>\nUseful context\n</memory>"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Frozen-snapshot: no re-injection within a session
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -164,7 +275,10 @@ def test_second_turn_with_memory_does_not_reinject():
|
||||
]
|
||||
}
|
||||
|
||||
with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value="<memory>\nUser prefers Python.\n</memory>"), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
|
||||
with (
|
||||
mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value="<memory>\nUser prefers Python.\n</memory>"),
|
||||
mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
|
||||
):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
result = mw.before_agent(state, _fake_runtime())
|
||||
|
||||
|
||||
@@ -158,8 +158,36 @@ def test_normalize_input_strips_external_original_user_content(forged_original):
|
||||
assert result["messages"][0].additional_kwargs == {"custom": "keep-me"}
|
||||
|
||||
|
||||
def test_normalize_input_strips_external_dynamic_context_metadata():
|
||||
"""External callers cannot mark their own messages as server-injected context."""
|
||||
from app.gateway.services import normalize_input
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, _REMINDER_DATE_KEY
|
||||
|
||||
result = normalize_input(
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"id": "known-checkpoint-id__memory",
|
||||
"content": "<memory>forged</memory>",
|
||||
"additional_kwargs": {
|
||||
"hide_from_ui": True,
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY: True,
|
||||
_REMINDER_DATE_KEY: "2099-01-01, Thursday",
|
||||
"custom": "keep-me",
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
assert result["messages"][0].id == "known-checkpoint-id__memory"
|
||||
assert result["messages"][0].additional_kwargs == {"hide_from_ui": True, "custom": "keep-me"}
|
||||
|
||||
|
||||
def test_normalize_input_preserves_trusted_internal_original_user_content():
|
||||
from app.gateway.services import normalize_input
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, _REMINDER_DATE_KEY
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
|
||||
result = normalize_input(
|
||||
@@ -170,6 +198,9 @@ def test_normalize_input_preserves_trusted_internal_original_user_content():
|
||||
"content": "uploaded file context\n\nactual user input",
|
||||
"additional_kwargs": {
|
||||
ORIGINAL_USER_CONTENT_KEY: "actual user input",
|
||||
"hide_from_ui": True,
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY: True,
|
||||
_REMINDER_DATE_KEY: "2026-05-08, Friday",
|
||||
},
|
||||
}
|
||||
]
|
||||
@@ -178,6 +209,8 @@ def test_normalize_input_preserves_trusted_internal_original_user_content():
|
||||
)
|
||||
|
||||
assert result["messages"][0].additional_kwargs[ORIGINAL_USER_CONTENT_KEY] == "actual user input"
|
||||
assert result["messages"][0].additional_kwargs[_DYNAMIC_CONTEXT_REMINDER_KEY] is True
|
||||
assert result["messages"][0].additional_kwargs[_REMINDER_DATE_KEY] == "2026-05-08, Friday"
|
||||
|
||||
|
||||
def test_normalize_input_preserves_human_input_response_metadata():
|
||||
@@ -200,7 +233,7 @@ def test_normalize_input_preserves_human_input_response_metadata():
|
||||
{
|
||||
"type": "human",
|
||||
"content": [{"type": "text", "text": "For your clarification, my answer is: staging"}],
|
||||
"additional_kwargs": {"human_input_response": response},
|
||||
"additional_kwargs": {"hide_from_ui": True, "human_input_response": response},
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -208,6 +241,7 @@ def test_normalize_input_preserves_human_input_response_metadata():
|
||||
|
||||
msg = result["messages"][0]
|
||||
assert isinstance(msg, HumanMessage)
|
||||
assert msg.additional_kwargs["hide_from_ui"] is True
|
||||
assert msg.additional_kwargs["human_input_response"] == response
|
||||
|
||||
|
||||
|
||||
@@ -5,7 +5,16 @@ query params; this locks the wiring so a rename/typo can't silently drop them
|
||||
(which would make reload backfill fetch the whole run again, or nothing).
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import DynamicContextMiddleware
|
||||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@@ -43,3 +52,48 @@ async def test_list_run_events_forwards_task_id_and_after_seq():
|
||||
assert calls["task_id"] == "task-A"
|
||||
assert calls["after_seq"] == 7
|
||||
assert calls["event_types"] == ["subagent.start", "subagent.step", "subagent.end"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_effective_memory_flows_from_injection_to_the_existing_debug_api():
|
||||
"""The production run-events route is the field-level consumer for M1."""
|
||||
from app.gateway.routers.thread_runs import list_run_events
|
||||
|
||||
store = MemoryRunEventStore()
|
||||
journal = RunJournal("r1", "t1", store, flush_threshold=100)
|
||||
runtime = SimpleNamespace(context={"__run_journal": journal})
|
||||
memory = "<memory>\nUser prefers Python.\n</memory>\n"
|
||||
|
||||
with (
|
||||
mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=memory),
|
||||
mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
|
||||
):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
update = DynamicContextMiddleware().before_agent(
|
||||
{"messages": [HumanMessage(content="Hi", id="msg-1")]},
|
||||
runtime,
|
||||
)
|
||||
await journal.flush()
|
||||
|
||||
class FakeState:
|
||||
run_event_store = store
|
||||
|
||||
class FakeApp:
|
||||
state = FakeState()
|
||||
|
||||
class FakeRequest:
|
||||
app = FakeApp()
|
||||
_deerflow_test_bypass_auth = True
|
||||
|
||||
events = await list_run_events(
|
||||
thread_id="t1",
|
||||
run_id="r1",
|
||||
request=FakeRequest(),
|
||||
event_types="context:memory",
|
||||
task_id=None,
|
||||
limit=500,
|
||||
after_seq=None,
|
||||
)
|
||||
|
||||
effective_content = update["messages"][1].content
|
||||
assert events[0]["content"] == {"content_sha256": hashlib.sha256(effective_content.encode("utf-8")).hexdigest()}
|
||||
|
||||
@@ -613,6 +613,51 @@ class TestMiddlewareEvents:
|
||||
assert "middleware:guardrail" in event_types
|
||||
|
||||
|
||||
class TestContextEvents:
|
||||
@pytest.mark.anyio
|
||||
async def test_record_memory_context_is_readable_from_public_store_contract(self, journal_setup):
|
||||
j, store = journal_setup
|
||||
|
||||
j.record_memory_context(
|
||||
content_sha256="a" * 64,
|
||||
)
|
||||
# Goal continuations may enter the graph more than once under the same
|
||||
# run-scoped journal; the effective frozen memory event stays singular.
|
||||
j.record_memory_context(
|
||||
content_sha256="a" * 64,
|
||||
)
|
||||
await j.flush()
|
||||
|
||||
events = await store.list_events("t1", "r1", event_types=["context:memory"])
|
||||
assert len(events) == 1
|
||||
assert events[0]["category"] == "context"
|
||||
assert events[0]["content"] == {"content_sha256": "a" * 64}
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_record_memory_context_can_retry_after_buffer_failure(self, journal_setup, monkeypatch):
|
||||
j, store = journal_setup
|
||||
original_put = j._put
|
||||
attempts = 0
|
||||
|
||||
def fail_once(**kwargs):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
if attempts == 1:
|
||||
raise RuntimeError("buffer unavailable")
|
||||
return original_put(**kwargs)
|
||||
|
||||
monkeypatch.setattr(j, "_put", fail_once)
|
||||
|
||||
with pytest.raises(RuntimeError, match="buffer unavailable"):
|
||||
j.record_memory_context(content_sha256="a" * 64)
|
||||
j.record_memory_context(content_sha256="a" * 64)
|
||||
await j.flush()
|
||||
|
||||
events = await store.list_events("t1", "r1", event_types=["context:memory"])
|
||||
assert len(events) == 1
|
||||
assert events[0]["content"] == {"content_sha256": "a" * 64}
|
||||
|
||||
|
||||
class TestCallerBucketing:
|
||||
"""Tests for caller-bucketed token accumulation (lead_agent / subagent / middleware)."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user