mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
feat(memory): redact queued memory payloads (pii_redaction follow-up) (#5577)
* feat(memory): redact queued memory payloads (pii_redaction follow-up) Closes the remaining vector from #3190 (memory storage): request-scoped redaction leaves originals in thread state, and MemoryMiddleware queued that raw conversation into the durable extraction buffer — so the extraction model's input and the persisted facts carried raw PII even with pii_redaction.enabled. - MemoryMiddleware gains pii_redaction_config (same pattern as summarization / durable-context / title); the payload returned by _resolve_add_args() is redacted at the enqueue boundary before add()/aadd(), covering both after_agent and aafter_agent - Reuses the shared redact_text() seam from #5527 via active_pii_detectors/_Redactor/_redact_content; the buffered payload itself becomes clean, not just the extraction model's view - Wired at the lead-agent assembly site - Tests: queue-payload redaction enabled/disabled, detector toggles, message immutability, token stability across turns * fix(memory): reserve existing placeholder indices across the queued batch Review round 1 on #5577: thread state mixes raw user turns with already-redacted tool results and summaries, so a token in a later message must not collide with a new value in an earlier one. Pre-scan the whole queued batch with redactor.reserve() before allocating. Also extend the two MemoryMiddleware stubs in test_lead_agent_model_resolution.py to accept the new pii_redaction_config kwarg (CI backend-unit-tests red). * fix(memory): redact the compaction-triggered memory_flush_hook queue too Review round 2 on #5577: with memory + summarization enabled, memory_flush_hook() queues list(event.messages_to_summarize) directly into MemoryManager.add_nowait(), bypassing MemoryMiddleware entirely — and compaction removes those messages from state right after, so the after-agent redaction can never repair a raw batch queued there. Extract the batch redaction (reserve pre-scan + rebuild) into a shared redact_queued_messages() used by both MemoryMiddleware's enqueue boundary and the hook; the hook gains an optional pii_redaction_config wired via functools.partial at the summarization factory (subagents use skip_memory_flush=True and are unaffected). Regression covers the hook path enabled/disabled/absent. * test(summarization): unwrap the partial in the factory hook identity check edd916e wrapped memory_flush_hook in functools.partial to carry the pii_redaction config, but missed committing the matching update to test_factory_attaches_memory_flush_hook_by_default — a functools.partial is not equal to its wrapped function, so the identity assertion failed deterministically (caught by static review on #5577). * fix(memory): cover tool-call arguments and lenient text blocks in queued payloads Review round 4 on #5577 — two reproduced gaps still allowed raw PII into memory payloads with pii_redaction.enabled=true: - OpenViking retains the full message object, so parsed tool_calls args and provider-format arguments in additional_kwargs bypassed the content-only redaction; queued payloads now redact tool-call args recursively (and reserve them in the pre-scan). - DeerMem's format_conversation_for_update reads p.get("text") of any dict block, not only type=="text" — _redact_content now redacts every string-valued text field so a lenient consumer cannot bypass it. Both repro shapes are pinned as regressions. * fix(pii): value-derived placeholders for stable capture identity Review round 5 on #5577: sequential allocation is order-dependent — the same unchanged message re-redacted in a later batch can get a different token, which flips OpenViking's content-signature and replays already recorded turns into persistent memory sessions. Placeholders are now a pure function of the matched value (sha256-6 per category): identities stay stable across batches, enqueues, seams, and backends with no shared state and no stored mapping. - _Redactor becomes stateless; the reserve machinery (index pre-scans, _PLACEHOLDER_PATTERN) is removed along with its call sites - redact_texts keeps its signature, now a simple per-text redaction - reserve-semantics tests are superseded by cross-enqueue identity tests (two successive captures where the second batch carries higher existing placeholders); the compiled-graph echo regex matches the new format * fix(pii): 128-bit placeholder digests; redact tool-argument keys Review round 6 on #5577: - [P2] The 24-bit truncation collided distinct identities (reproduced: contact3513@ / contact3727@example.com both rendered [EMAIL_575dbb]). Retain 128 bits of the digest so separate people stay separate at any realistic volume; the compiled-graph echo regex and the AGENTS.md entry follow the new format. - [P2] _redact_strings redacted values but preserved mapping keys, so {"contacts": {"alice@example.com": "manager"}} kept the raw email as a key. Redact string keys along with values; deterministic 128-bit tokens make distinct-key collisions a hash-collision non-event, and a redacted key equal to another key merges the entries. Regression pins the email-keyed arguments shape. * docs(agents): trim PiiRedactionMiddleware entry; drop stale reserve wording The chain grew past the 98304-byte AG002 hard limit again (main growth plus the round-6 format example). Trim the entry to fit, and replace the stale reserve-indices wording with the value-derived semantics that replaced it. * docs(agents): shrink PiiRedactionMiddleware entry below the chain budget main moved (#5609) and the base comparison flipped the nearly-over chain back to an AG002 error (98441 > 98304). Trim the entry to the facts that are not documented elsewhere: seams, allowlist, order, enqueue boundary. * fix(pii): letters-only value-derived tokens; wire the factory SDK path Review round 7 on #5577: - [P2] Hex digests carry digit runs that later digit-anchored detectors re-scan, corrupting minted tokens into nested placeholders (reproduced: user280@example.com -> [EMAIL_…[PHONE_…]…], ~1/140 emails). The 128-bit digest is now encoded over a-z, so a minted token survives the full pinned detector order. - [P2] create_deerflow_agent (the SDK path) constructed MemoryMiddleware and DurableContextMiddleware without pii_redaction_config; the knob is now an explicit optional parameter threaded through _assemble_from_features (the factory reads no config files by design, so SDK callers pass it explicitly). - [Nit] hoist the Command rebuild comprehension out of the per-message loop in _redact_result. - [Nit] drop the dead _iter_tool_calls helper and Iterator import. * test(factory): pin pii_redaction_config wiring; document the parameter Review round 8 on #5577: - [P3] test_create_deerflow_agent.py had zero coverage of the pii_redaction_config parameter; two regressions now assert that create_deerflow_agent(pii_redaction_config=cfg) builds a MemoryMiddleware and a DurableContextMiddleware carrying cfg, and that the default is None (redaction off). - [Nit] document the parameter in the create_deerflow_agent docstring: None leaves redaction off; a middleware= takeover bypasses the assembly that would apply it. * fix(pii): HMAC-scoped placeholder tokens; install model-bound redaction in the SDK chain Review round 10 on #5577: - [P2] sha256(category||value) is an unkeyed, globally deterministic fingerprint: a provider or memory backend can confirm candidate values offline and link identical values across users and threads, contradicting the documented no-re-identification guarantee. Placeholders are now HMAC-derived over a deployment-scoped token_secret (new config field); the docstring documents the linkability scope for both modes. - [P2] _assemble_from_features() never appended PiiRedactionMiddleware, so the SDK knob enabled only the downstream seams while the primary model still received raw content. The middleware is now appended in the same relative position as the lead-agent base chain when the gate is on, with a model-call regression through the assembled SDK chain. * docs(pii): document token_secret; update the placeholder format examples Follow-up to the round-10 review on #5577: - config.example.yaml documents the new token_secret knob: a deployment-scoped secret makes placeholders HMAC-derived and only linkable within one deployment; leaving it unset keeps publicly computable fingerprints (documented trade). - the PiiRedactionConfig class docstring's placeholder example now matches the shipped format (128-bit HMAC digests over base-26 letters). * docs(pii): update the module docstring's placeholder format example The module docstring still showed the sequential [EMAIL_1] example; align it with the shipped value-derived format (128-bit HMAC digests over base-26 letters). * fix(pii): require a non-empty token_secret when redaction is enabled Review round 11 on #5577: with the secret optional, the default enabled configuration minted empty-key HMAC digests — publicly computable, globally linkable fingerprints of the raw values, confirmable by offline guessing for low-entropy identifiers such as phone numbers. - PiiRedactionConfig now rejects enabled=true without a non-empty token_secret (minimum 16 characters); the unkeyed mode is gone entirely. - The request-scoped redactor in _process_request now carries the configured key: it built its redactor without it, so user-message placeholders silently diverged from every other seam for the same value. The empty-key guard in _placeholder_token is the backstop that keeps every path keyed. - Module/config docstrings and config.example.yaml no longer describe the placeholders as irreversible or the secret as optional; the stale reserve-semantics wording and the "memory path not covered" note are corrected. - Regressions: validation errors for missing/blank/short secrets, and the user-message boundary minting the same keyed token as redact_text. * fix(memory): redact original_user_content on queued memory payloads Review round 12 on #5577: UploadsMiddleware preserves the raw user turn in additional_kwargs.original_user_content, but redact_queued_messages rewrote only content, tool_calls, and provider-format tool_calls — the queued copy still carried the raw value, so the memory queue and backend received raw PII with redaction enabled. The shared helper now redacts the provenance string field on the queued copy (same batch redactor, merged with the existing additional_kwargs rewrite, rebuilt via model_copy — thread-state messages stay untouched), covering both the enqueue boundary and the compaction flush hook. Regressions pin the redacted provenance field through both paths and the untouched original. * fix(memory): redact human_input_response on queued memory payloads Review round 13 on #5577: clarification replies keep the raw user answer in additional_kwargs.human_input_response.value (DeerMem reads that mapping via its read_human_input_response mirror), but redact_queued_messages rewrote only content, tool_calls, and original_user_content — the queued copy still carried the raw value into the memory queue and backend. The shared helper now rebuilds that mapping with value redacted (same batch redactor) while preserving the protocol fields (version/kind/source/ request_id/option_id/response_kind), merged with the existing additional_kwargs rewrite and rebuilt via model_copy so thread-state messages stay untouched. Regressions pin the redacted value, preserved protocol fields, the untouched original, and the flush-hook path. * chore(config): bump config_version to 47; generate token_secret on upgrade Review round 13 on #5577: existing deployments can legitimately sit at config_version 46 with pii_redaction.enabled: true and no token_secret; the mandatory-secret validator would fail their startup while config.example.yaml still said 46, so the upgrade checker considered them current and emitted no migration warning. - Bump config_version to 47 so the outdated-config warning fires for v46. - Add a v47 migration to scripts/config-upgrade.sh: when pii_redaction is enabled and token_secret is missing/blank, generate a random deployment- scoped secret (secrets.token_urlsafe) and persist it, so the upgraded deployment starts under the mandatory validation. - Regressions: a v46 enabled config is flagged outdated against the real example; the upgrade run generates a >= 16-char secret and the result passes AppConfig validation; a disabled config gets no secret. - Also switch the two v41 script tests from a hardcoded "bash" to SCRIPT_BASH/skipif, matching the documented Windows constraint and the neighboring tests (hardcoded bash resolves to the WSL stub there). * chore(helm): sync chart config_version to 47 Review round 14 on #5577: check_config_version.sh fails whenever the chart's embedded config_version lags config.example.yaml, and the chart workflow is triggered by config.example.yaml — so the v47 bump left this head failing that check. Update the values.yaml config block and the README example to 47; scripts/check_config_version.sh now passes. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
xiaodu55
GitHub
parent
72933db4e9
commit
64e499be98
@ -34,6 +34,7 @@ if TYPE_CHECKING:
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
|
||||
from deerflow.config.memory_config import MemoryConfig
|
||||
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
||||
from deerflow.subagents.runtime import SubagentRuntime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -78,6 +79,7 @@ def create_deerflow_agent(
|
||||
checkpointer: BaseCheckpointSaver | None = None,
|
||||
name: str = "default",
|
||||
subagent_runtime: SubagentRuntime | None = None,
|
||||
pii_redaction_config: PiiRedactionConfig | None = None,
|
||||
) -> CompiledStateGraph:
|
||||
"""Create a DeerFlow agent from plain Python arguments.
|
||||
|
||||
@ -124,6 +126,11 @@ def create_deerflow_agent(
|
||||
only when the caller needs non-default native-subagent capacity or a
|
||||
caller-managed durable batch worker without Gateway/DeerFlowClient
|
||||
startup. Requires ``features.subagent`` to be enabled.
|
||||
pii_redaction_config:
|
||||
Optional PII redaction policy (#3190). ``None`` leaves redaction off;
|
||||
when enabled, memory-queue and durable-context redaction apply to the
|
||||
assembled middlewares. A ``middleware=`` takeover bypasses the
|
||||
assembly that would apply it.
|
||||
|
||||
Raises
|
||||
------
|
||||
@ -164,6 +171,7 @@ def create_deerflow_agent(
|
||||
plan_mode=plan_mode,
|
||||
extra_middleware=extra_middleware or [],
|
||||
subagent_runtime=subagent_runtime,
|
||||
pii_redaction_config=pii_redaction_config,
|
||||
)
|
||||
# Deduplicate by tool name — user-provided tools take priority.
|
||||
existing_names = {t.name for t in effective_tools}
|
||||
@ -201,6 +209,7 @@ def _assemble_from_features(
|
||||
plan_mode: bool = False,
|
||||
extra_middleware: list[AgentMiddleware] | None = None,
|
||||
subagent_runtime: SubagentRuntime | None = None,
|
||||
pii_redaction_config: PiiRedactionConfig | None = None,
|
||||
) -> tuple[list[AgentMiddleware], list[BaseTool]]:
|
||||
"""Build an ordered middleware chain + extra tools from *feat*.
|
||||
|
||||
@ -269,7 +278,12 @@ def _assemble_from_features(
|
||||
from deerflow.agents.middlewares.durable_context_middleware import DurableContextMiddleware
|
||||
from deerflow.agents.middlewares.system_message_coalescing_middleware import SystemMessageCoalescingMiddleware
|
||||
|
||||
chain.append(DurableContextMiddleware())
|
||||
if pii_redaction_config is not None and pii_redaction_config.enabled:
|
||||
from deerflow.agents.middlewares.pii_redaction_middleware import PiiRedactionMiddleware
|
||||
|
||||
chain.append(PiiRedactionMiddleware(pii_redaction_config))
|
||||
|
||||
chain.append(DurableContextMiddleware(pii_redaction_config=pii_redaction_config))
|
||||
# DurableContext adds its authority contract as a second SystemMessage; strict backends
|
||||
# (vLLM, SGLang, Qwen, Anthropic) reject that, so merge them into one leading message.
|
||||
chain.append(SystemMessageCoalescingMiddleware())
|
||||
@ -318,13 +332,25 @@ def _assemble_from_features(
|
||||
if backend_requires_passive_writes_in_tool_mode(memory_cfg.manager_class):
|
||||
from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware
|
||||
|
||||
chain.append(MemoryMiddleware(agent_name=name, memory_config=memory_cfg))
|
||||
chain.append(
|
||||
MemoryMiddleware(
|
||||
agent_name=name,
|
||||
memory_config=memory_cfg,
|
||||
pii_redaction_config=pii_redaction_config,
|
||||
)
|
||||
)
|
||||
else:
|
||||
if memory_cfg.mode == "tool" and not memory_cfg.enabled:
|
||||
logger.warning("memory.mode is 'tool' but memory.enabled is false; memory tools will not be registered.")
|
||||
from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware
|
||||
|
||||
chain.append(MemoryMiddleware(agent_name=name, memory_config=memory_cfg))
|
||||
chain.append(
|
||||
MemoryMiddleware(
|
||||
agent_name=name,
|
||||
memory_config=memory_cfg,
|
||||
pii_redaction_config=pii_redaction_config,
|
||||
)
|
||||
)
|
||||
|
||||
# --- [10] Vision ---
|
||||
if feat.vision is not False:
|
||||
|
||||
@ -649,11 +649,23 @@ def build_middlewares(
|
||||
from deerflow.agents.memory.manager import backend_requires_passive_writes_in_tool_mode
|
||||
|
||||
if backend_requires_passive_writes_in_tool_mode(resolved_app_config.memory.manager_class):
|
||||
middlewares.append(MemoryMiddleware(agent_name=agent_name, memory_config=resolved_app_config.memory))
|
||||
middlewares.append(
|
||||
MemoryMiddleware(
|
||||
agent_name=agent_name,
|
||||
memory_config=resolved_app_config.memory,
|
||||
pii_redaction_config=getattr(resolved_app_config, "pii_redaction", None),
|
||||
)
|
||||
)
|
||||
else:
|
||||
if resolved_app_config.memory.mode == "tool" and not resolved_app_config.memory.enabled:
|
||||
logger.warning("memory.mode is 'tool' but memory.enabled is false; memory tools will not be registered.")
|
||||
middlewares.append(MemoryMiddleware(agent_name=agent_name, memory_config=resolved_app_config.memory))
|
||||
middlewares.append(
|
||||
MemoryMiddleware(
|
||||
agent_name=agent_name,
|
||||
memory_config=resolved_app_config.memory,
|
||||
pii_redaction_config=getattr(resolved_app_config, "pii_redaction", None),
|
||||
)
|
||||
)
|
||||
|
||||
# Add ViewImageMiddleware only if the current model supports vision.
|
||||
# Use the resolved runtime model_name from make_lead_agent to avoid stale config values.
|
||||
|
||||
@ -3,18 +3,23 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from deerflow.agents.memory import get_memory_manager
|
||||
from deerflow.agents.middlewares.memory_middleware import redact_queued_messages
|
||||
from deerflow.agents.middlewares.summarization_middleware import SummarizationEvent
|
||||
from deerflow.config.memory_config import get_memory_config
|
||||
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
||||
from deerflow.runtime.user_context import resolve_runtime_user_id
|
||||
|
||||
|
||||
def memory_flush_hook(event: SummarizationEvent) -> None:
|
||||
def memory_flush_hook(event: SummarizationEvent, pii_redaction_config: PiiRedactionConfig | None = None) -> None:
|
||||
"""Flush messages about to be summarized into the memory queue.
|
||||
|
||||
Thin, backend-agnostic entry: only the ``enabled`` + ``thread_id`` gate
|
||||
and ``user_id`` resolution live here. The backend (via
|
||||
``manager.add_nowait``) does the filtering, human/AI validation, and
|
||||
correction/reinforcement detection.
|
||||
correction/reinforcement detection. The queued payload is redacted at
|
||||
this boundary (#3190 vector 5): compaction removes these messages from
|
||||
state right after, so the later after-agent redaction cannot repair a
|
||||
raw batch queued here.
|
||||
"""
|
||||
if not get_memory_config().enabled or not event.thread_id:
|
||||
return
|
||||
@ -22,7 +27,7 @@ def memory_flush_hook(event: SummarizationEvent) -> None:
|
||||
user_id = resolve_runtime_user_id(event.runtime)
|
||||
get_memory_manager().add_nowait(
|
||||
event.thread_id,
|
||||
list(event.messages_to_summarize),
|
||||
redact_queued_messages(list(event.messages_to_summarize), pii_redaction_config),
|
||||
agent_name=event.agent_name,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
@ -64,7 +64,7 @@ strict providers reject.
|
||||
ordered by application — the last entry produced the final visible bytes — so an
|
||||
observer classifies raw→visible transforms from facts rather than by sniffing
|
||||
output wording.
|
||||
4. **PiiRedactionMiddleware** - *(optional, `pii_redaction.enabled`, default off, #3190)* Rewrites PII in genuine user messages (`wrap_model_call`, request-scoped, raw text stays in thread state) and remote-content tool results (`wrap_tool_call`, ToolResultSanitizationMiddleware's allowlist incl. `Command.update.messages`) to irreversible placeholders (`[EMAIL_1]`, …). Deterministic regex detectors only, no new dependencies; national IDs run before the card detector so a checksum-valid resident ID whose digits also pass Luhn is never mislabeled `[CREDIT_CARD_n]`. One redactor per result keeps numbering continuous; existing summary/message tokens reserve indices before new values; no persistent mapping links compacted identities. Innermost Layer-1 wrapper, so budget-externalized copies hold redacted text; compaction, reinjected summaries, and title-model inputs are redacted; title fields are redacted before truncation. Memory extraction is a follow-up slice.
|
||||
4. **PiiRedactionMiddleware** - *(optional, `pii_redaction.enabled`, default off, #3190)* Rewrites PII in genuine user messages (`wrap_model_call`) and remote-content tool results (`wrap_tool_call`, the tool-result allowlist) to irreversible value-derived placeholders. Deterministic regex detectors only. Innermost Layer-1 wrapper; compaction input, summaries, title input and queued memory payloads are redacted via `redact_text`.
|
||||
|
||||
5. **ThreadDataMiddleware** - Creates per-thread directories under the user's isolation scope (`backend/.deer-flow/users/{user_id}/threads/{thread_id}/user-data/{workspace,uploads,outputs}`); resolves identity via `resolve_runtime_user_id(runtime)`, including Gateway runtime context and standalone LangGraph Server auth, then falls back to the request ContextVar / `"default"`
|
||||
6. **UploadsMiddleware** - Tracks and injects newly uploaded files into conversation (lead agent only); upload existence checks use the same runtime-resolved user bucket as thread-data creation
|
||||
|
||||
@ -2,6 +2,7 @@
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, override
|
||||
|
||||
from langchain.agents import AgentState
|
||||
@ -9,10 +10,14 @@ from langchain.agents.middleware import AgentMiddleware
|
||||
from langgraph.config import get_config
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from deerflow.agents.human_input import HUMAN_INPUT_RESPONSE_KEY
|
||||
from deerflow.agents.memory import get_memory_manager
|
||||
from deerflow.agents.middlewares.pii_redaction_middleware import _make_redactor, _redact_content, _Redactor
|
||||
from deerflow.config.memory_config import get_memory_config
|
||||
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
||||
from deerflow.runtime.user_context import resolve_runtime_user_id
|
||||
from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, resolve_trace_id
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from deerflow.config.memory_config import MemoryConfig
|
||||
@ -20,6 +25,90 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def redact_queued_messages(messages: list, pii_redaction_config: PiiRedactionConfig | None) -> list:
|
||||
"""Redact a conversation payload queued for memory extraction (#3190 vector 5).
|
||||
|
||||
Shared by MemoryMiddleware's enqueue boundary and the compaction-triggered
|
||||
``memory_flush_hook``: placeholders are value-derived, so the same identity
|
||||
renders the same token across the batch, across enqueues, and across
|
||||
seams; message objects are rebuilt rather than mutated. Structured
|
||||
tool-call arguments are covered too — backends like OpenViking retain the full
|
||||
message object, including ``tool_calls`` and provider-format arguments in
|
||||
``additional_kwargs``. Text-bearing provenance metadata is covered as
|
||||
well: UploadsMiddleware preserves the raw user turn in
|
||||
``original_user_content``, and clarification replies keep the raw answer
|
||||
in ``human_input_response.value`` — DeerMem reads that mapping — so a
|
||||
redacted ``content`` alone would still leak either one to the memory
|
||||
backend.
|
||||
"""
|
||||
redactor = _make_redactor(pii_redaction_config)
|
||||
|
||||
redacted = []
|
||||
for message in messages:
|
||||
updates: dict = {}
|
||||
content, changed = _redact_content(message.content, redactor)
|
||||
if changed:
|
||||
updates["content"] = content
|
||||
tool_calls = getattr(message, "tool_calls", None)
|
||||
if tool_calls:
|
||||
rewritten = [_redact_tool_call(call, redactor) for call in tool_calls]
|
||||
if rewritten != tool_calls:
|
||||
updates["tool_calls"] = rewritten
|
||||
additional_kwargs = getattr(message, "additional_kwargs", None)
|
||||
if isinstance(additional_kwargs, dict):
|
||||
new_additional_kwargs = additional_kwargs
|
||||
if additional_kwargs.get("tool_calls"):
|
||||
rewritten = _redact_strings(additional_kwargs["tool_calls"], redactor)
|
||||
if rewritten != additional_kwargs["tool_calls"]:
|
||||
new_additional_kwargs = {**new_additional_kwargs, "tool_calls": rewritten}
|
||||
original_content = additional_kwargs.get(ORIGINAL_USER_CONTENT_KEY)
|
||||
if isinstance(original_content, str) and original_content:
|
||||
redacted_original = redactor.redact(original_content)
|
||||
if redacted_original != original_content:
|
||||
new_additional_kwargs = {**new_additional_kwargs, ORIGINAL_USER_CONTENT_KEY: redacted_original}
|
||||
human_input = additional_kwargs.get(HUMAN_INPUT_RESPONSE_KEY)
|
||||
if isinstance(human_input, Mapping) and isinstance(human_input.get("value"), str) and human_input["value"]:
|
||||
redacted_value = redactor.redact(human_input["value"])
|
||||
if redacted_value != human_input["value"]:
|
||||
new_additional_kwargs = {**new_additional_kwargs, HUMAN_INPUT_RESPONSE_KEY: {**human_input, "value": redacted_value}}
|
||||
if new_additional_kwargs is not additional_kwargs:
|
||||
updates["additional_kwargs"] = new_additional_kwargs
|
||||
redacted.append(message.model_copy(update=updates) if updates else message)
|
||||
return redacted
|
||||
|
||||
|
||||
def _redact_tool_call(call: dict, redactor: _Redactor) -> dict:
|
||||
"""Redact PII inside one parsed tool-call dict (args values recursively)."""
|
||||
rewritten = dict(call)
|
||||
args = rewritten.get("args")
|
||||
if args is not None:
|
||||
rewritten["args"] = _redact_strings(args, redactor)
|
||||
return rewritten
|
||||
|
||||
|
||||
def _redact_strings(value: object, redactor: _Redactor) -> object:
|
||||
"""Redact every string leaf in a JSON-like *value* tree, keys included.
|
||||
|
||||
Tool arguments can carry user data in mapping keys (e.g.
|
||||
``{"contacts": {"alice@example.com": "manager"}}``), so string keys are
|
||||
redacted along with values. Two distinct raw keys redacting to the same
|
||||
token would require a sha256 collision; if a redacted key does equal
|
||||
another key already present in the rebuilt dict, the entries merge —
|
||||
accepted, since at 128 bits that only happens for identical identities.
|
||||
"""
|
||||
if isinstance(value, str):
|
||||
return redactor.redact(value)
|
||||
if isinstance(value, list):
|
||||
return [_redact_strings(item, redactor) for item in value]
|
||||
if isinstance(value, dict):
|
||||
rebuilt = {}
|
||||
for key, item in value.items():
|
||||
new_key = redactor.redact(key) if isinstance(key, str) else key
|
||||
rebuilt[new_key] = _redact_strings(item, redactor)
|
||||
return rebuilt
|
||||
return value
|
||||
|
||||
|
||||
class MemoryMiddlewareState(AgentState):
|
||||
"""Compatible with the `ThreadState` schema."""
|
||||
|
||||
@ -38,17 +127,29 @@ class MemoryMiddleware(AgentMiddleware[MemoryMiddlewareState]):
|
||||
|
||||
state_schema = MemoryMiddlewareState
|
||||
|
||||
def __init__(self, agent_name: str | None = None, *, memory_config: "MemoryConfig | None" = None):
|
||||
def __init__(
|
||||
self,
|
||||
agent_name: str | None = None,
|
||||
*,
|
||||
memory_config: "MemoryConfig | None" = None,
|
||||
pii_redaction_config: PiiRedactionConfig | None = None,
|
||||
):
|
||||
"""Initialize the MemoryMiddleware.
|
||||
|
||||
Args:
|
||||
agent_name: If provided, memory is stored per-agent. If None, uses global memory.
|
||||
memory_config: Explicit memory config. When omitted, legacy global
|
||||
config fallback is used.
|
||||
pii_redaction_config: When enabled, the queued conversation payload
|
||||
is redacted at the enqueue boundary (#3190 vector 5) — request-scoped
|
||||
redaction leaves originals in thread state, and the buffered payload
|
||||
is durable, so the extraction model's input and the persisted facts
|
||||
would otherwise carry raw PII.
|
||||
"""
|
||||
super().__init__()
|
||||
self._agent_name = agent_name
|
||||
self._memory_config = memory_config
|
||||
self._pii_redaction_config = pii_redaction_config
|
||||
|
||||
def _resolve_add_args(self, state: MemoryMiddlewareState, runtime: Runtime) -> tuple[str, list, str, str] | None:
|
||||
"""Resolve one write request without invoking the manager."""
|
||||
@ -83,7 +184,11 @@ class MemoryMiddleware(AgentMiddleware[MemoryMiddlewareState]):
|
||||
runtime_context = runtime.context if isinstance(runtime.context, dict) else {}
|
||||
trace_id = resolve_trace_id(runtime_context.get(DEERFLOW_TRACE_METADATA_KEY))
|
||||
|
||||
return thread_id, messages, user_id, trace_id
|
||||
return thread_id, self._redact_queued_messages(messages), user_id, trace_id
|
||||
|
||||
def _redact_queued_messages(self, messages: list) -> list:
|
||||
"""Redact the conversation payload queued for extraction (#3190 vector 5)."""
|
||||
return redact_queued_messages(messages, self._pii_redaction_config)
|
||||
|
||||
@override
|
||||
def after_agent(self, state: MemoryMiddlewareState, runtime: Runtime) -> dict | None:
|
||||
|
||||
@ -2,33 +2,38 @@
|
||||
|
||||
Detects personally identifiable information in the two untrusted-content entry
|
||||
points — genuine user messages and remote-content tool results — and rewrites
|
||||
it to irreversible placeholders (``[EMAIL_1]`` …) before it reaches the model.
|
||||
it to keyed, value-derived placeholders (128-bit HMAC digests over base-26
|
||||
letters, keyed by the deployment-scoped ``token_secret``) before it reaches
|
||||
the model.
|
||||
Complements the structural guardrails: ``InputSanitizationMiddleware``
|
||||
neutralizes injection tags in user input and ``ToolResultSanitizationMiddleware``
|
||||
does the same for remote tool results; neither inspects *content* for PII.
|
||||
|
||||
v1 is deterministic-only: fixed regex detectors with checksum validation where
|
||||
the identifier format defines one (Luhn for card numbers, mod-11 for CN resident
|
||||
IDs and CPF), no model calls, no new dependencies. Redaction is irreversible —
|
||||
no mapping table is stored, so there is nothing to protect and no
|
||||
re-identification path.
|
||||
IDs and CPF), no model calls, no new dependencies. Placeholders are HMAC-keyed
|
||||
with the deployment-scoped ``token_secret`` (required whenever redaction is
|
||||
enabled) and no mapping table is stored, so tokens cannot be re-derived
|
||||
offline — neither from the repository nor from an observed token.
|
||||
|
||||
Scope model (mirrors the structural guardrails):
|
||||
|
||||
* the user-message rewrite is request-scoped — thread state keeps the raw text,
|
||||
so the UI still shows the original message and the whole conversation is
|
||||
re-redacted on every model call. Existing summary/message placeholders reserve
|
||||
their indices before new values are assigned, avoiding collisions after
|
||||
compaction. Without a persisted mapping, a repeated raw value cannot be
|
||||
linked to a placeholder whose source was compacted away;
|
||||
re-redacted on every model call. Placeholders are value-derived
|
||||
(a 128-bit HMAC of the value under the deployment secret): the same raw
|
||||
value always renders the same token, so identities stay stable across
|
||||
turns, compaction, enqueues, and downstream content-signature
|
||||
deduplication — without any stored mapping;
|
||||
* tool-result redaction runs at the tool boundary (``wrap_tool_call``) with the
|
||||
same allowlist as ``ToolResultSanitizationMiddleware`` (first-party web tools
|
||||
by name, MCP tools via their ``deerflow_mcp`` tag), so redacted text is what
|
||||
enters model context in the first place; placeholders restart per result;
|
||||
enters model context in the first place;
|
||||
* subagents are covered because ``build_subagent_runtime_middlewares`` reuses
|
||||
this base;
|
||||
* NOT covered in v1: the memory-extraction path (follow-up slice per the issue
|
||||
discussion). Allowlisted tool results are redacted before budget externalization.
|
||||
* The memory-enqueue path is covered by the follow-up slice (#5577):
|
||||
``redact_queued_messages`` (memory middleware) applies the same configured
|
||||
policy to the extraction payloads queued for the memory backend.
|
||||
* The compaction and durable-context seams run outside ``wrap_model_call``;
|
||||
:func:`redact_text` is the shared entry point they call, wired from
|
||||
SummarizationMiddleware (compaction input) and DurableContextMiddleware
|
||||
@ -45,6 +50,8 @@ gates did not claim.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
@ -202,55 +209,74 @@ def redact_text(text: str | None, config: PiiRedactionConfig | None) -> str | No
|
||||
"""
|
||||
if config is None or not config.enabled or not isinstance(text, str) or not text:
|
||||
return text
|
||||
return _Redactor(active_pii_detectors(config)).redact(text)
|
||||
return _make_redactor(config).redact(text)
|
||||
|
||||
|
||||
def redact_texts(texts: Sequence[str], config: PiiRedactionConfig | None) -> list[str]:
|
||||
"""Redact related text fields with one allocation scope, before truncation."""
|
||||
redactor = _Redactor(active_pii_detectors(config))
|
||||
for text in texts:
|
||||
redactor.reserve(text)
|
||||
"""Redact related text fields, before truncation."""
|
||||
redactor = _make_redactor(config)
|
||||
return [redactor.redact(text) for text in texts]
|
||||
|
||||
|
||||
# Bound the numeric field before int() conversion; generated indices are tiny.
|
||||
_PLACEHOLDER_PATTERN = re.compile(r"\[(EMAIL|API_KEY|NATIONAL_ID|CREDIT_CARD|PHONE)_([1-9][0-9]{0,19})\]")
|
||||
def _make_redactor(config: PiiRedactionConfig | None) -> _Redactor:
|
||||
"""Build the scan redactor: active detectors + the configured token key."""
|
||||
if config is None:
|
||||
return _Redactor((), None)
|
||||
return _Redactor(active_pii_detectors(config), config.token_secret)
|
||||
|
||||
|
||||
def _placeholder_token(category: str, value: str, token_key: str) -> str:
|
||||
"""Value-derived, deterministic placeholder (128-bit HMAC digest, base-26 letters).
|
||||
|
||||
Three properties the token format must uphold:
|
||||
|
||||
* **Value-derived** — sequential allocation is order-dependent: the same
|
||||
unchanged message re-redacted in a later batch can get a different
|
||||
token, which breaks downstream content-signature deduplication (e.g.
|
||||
OpenViking capture) and cross-turn identity. A pure function of the
|
||||
matched value keeps the token stable across batches, seams, and
|
||||
enqueues, with no shared state and no stored mapping;
|
||||
* **Letters-only** — the digest is encoded over ``a-z`` instead of hex, so
|
||||
a minted token contains no digit runs and can never satisfy the
|
||||
digit-anchored detectors that run later in the pinned order (no nested
|
||||
``[EMAIL_…[PHONE_…]…`` corruption);
|
||||
* **Deployment-scoped linkability** — the digest is an HMAC over the
|
||||
deployment-scoped ``token_secret``: tokens are linkable only within one
|
||||
deployment and cannot be re-derived offline without the secret. The
|
||||
config model rejects enabling redaction without a non-empty secret, and
|
||||
the empty-key guard below is the backstop that keeps every code path on
|
||||
keyed digests (an unkeyed digest would be a publicly computable
|
||||
fingerprint, linkable across deployments and confirmable by guessing
|
||||
for low-entropy values). 128 bits keep distinct identities
|
||||
collision-free at any realistic volume.
|
||||
"""
|
||||
if not token_key:
|
||||
raise ValueError("placeholder tokens require the configured non-empty token_secret")
|
||||
key_bytes = token_key.encode("utf-8")
|
||||
digest = hmac.new(key_bytes, f"{category}\x00{value}".encode(), hashlib.sha256).digest()
|
||||
alphabet = "abcdefghijklmnopqrstuvwxyz"
|
||||
number = int.from_bytes(digest[:16], "big")
|
||||
letters = []
|
||||
for _ in range(27):
|
||||
number, rem = divmod(number, 26)
|
||||
letters.append(alphabet[rem])
|
||||
return f"[{category.upper()}_{''.join(letters)}]" # noqa: FS003
|
||||
|
||||
|
||||
class _Redactor:
|
||||
"""Per-scan redaction state: one stable placeholder per distinct value.
|
||||
"""Applies the active detectors, in pinned order, to one piece of text.
|
||||
|
||||
A single instance covers one scan (one model request, or one tool result),
|
||||
so identical raw values in that scan render the same placeholder. Existing
|
||||
tokens reserve indices but cannot recover compacted raw-value identities. Placeholders are irreversible — the mapping lives only as
|
||||
long as this instance.
|
||||
Stateless by design: the placeholder is a pure function of the matched
|
||||
value (see :func:`_placeholder_token`), so identical values render the
|
||||
same token within a scan, across scans, across enqueues, and across every
|
||||
seam that shares this module.
|
||||
"""
|
||||
|
||||
def __init__(self, detectors: Sequence[_Detector]) -> None:
|
||||
def __init__(self, detectors: Sequence[_Detector], token_key: str | None = None) -> None:
|
||||
self._detectors = detectors
|
||||
self._tokens: dict[tuple[str, str], str] = {}
|
||||
self._counts: dict[str, int] = {}
|
||||
|
||||
def reserve(self, content: object) -> None:
|
||||
"""Reserve visible placeholders before assigning any new values.
|
||||
|
||||
Only placeholder counters survive compaction, never raw-value mappings.
|
||||
Pre-scan all fields so a token in a later block cannot collide with a
|
||||
new value in an earlier one.
|
||||
"""
|
||||
if isinstance(content, str):
|
||||
for match in _PLACEHOLDER_PATTERN.finditer(content):
|
||||
name, index = match.groups()
|
||||
self._counts[name.lower()] = max(self._counts.get(name.lower(), 0), int(index))
|
||||
elif isinstance(content, list):
|
||||
for block in content:
|
||||
if isinstance(block, str):
|
||||
self.reserve(block)
|
||||
elif isinstance(block, dict) and block.get("type") == "text":
|
||||
self.reserve(block.get("text"))
|
||||
self._token_key = token_key
|
||||
|
||||
def redact(self, text: str) -> str:
|
||||
self.reserve(text)
|
||||
for detector in self._detectors:
|
||||
text = detector.pattern.sub(self._replacer(detector), text)
|
||||
return text
|
||||
@ -260,13 +286,7 @@ class _Redactor:
|
||||
value = match.group(0)
|
||||
if detector.validator is not None and not detector.validator(value):
|
||||
return value
|
||||
key = (detector.name, value)
|
||||
token = self._tokens.get(key)
|
||||
if token is None:
|
||||
self._counts[detector.name] = self._counts.get(detector.name, 0) + 1
|
||||
token = f"[{detector.name.upper()}_{self._counts[detector.name]}]"
|
||||
self._tokens[key] = token
|
||||
return token
|
||||
return _placeholder_token(detector.name, value, self._token_key)
|
||||
|
||||
return replace
|
||||
|
||||
@ -275,10 +295,12 @@ def _redact_content(content: object, redactor: _Redactor) -> tuple[object, bool]
|
||||
"""Redact *content*, preserving its shape. Returns ``(content, changed)``.
|
||||
|
||||
Handles the two shapes message content takes — plain ``str`` and a list of
|
||||
content blocks. Non-text blocks (images, etc.) pass through untouched.
|
||||
The input is never mutated.
|
||||
content blocks. Blocks pass through untouched except for any string-valued
|
||||
``text`` field they carry: lenient downstream consumers (e.g. DeerMem's
|
||||
``format_conversation_for_update``) read ``p.get("text")`` regardless of
|
||||
the block type, so a non-text block must not smuggle raw PII past the
|
||||
helper. The input is never mutated.
|
||||
"""
|
||||
redactor.reserve(content)
|
||||
if isinstance(content, str):
|
||||
redacted = redactor.redact(content)
|
||||
return redacted, redacted != content
|
||||
@ -291,7 +313,7 @@ def _redact_content(content: object, redactor: _Redactor) -> tuple[object, bool]
|
||||
redacted = redactor.redact(block)
|
||||
changed = changed or redacted != block
|
||||
new_content.append(redacted)
|
||||
elif isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str):
|
||||
elif isinstance(block, dict) and isinstance(block.get("text"), str):
|
||||
redacted = redactor.redact(block["text"])
|
||||
if redacted != block["text"]:
|
||||
new_content.append({**block, "text": redacted})
|
||||
@ -314,6 +336,7 @@ class PiiRedactionMiddleware(AgentMiddleware[AgentState]):
|
||||
|
||||
def __init__(self, config: PiiRedactionConfig) -> None:
|
||||
self._detectors = active_pii_detectors(config)
|
||||
self._token_key = config.token_secret if config else None
|
||||
|
||||
def release_policy_parameters(self) -> dict[str, object]:
|
||||
"""Declare the behaviour-affecting settings (middleware module guide)."""
|
||||
@ -325,13 +348,10 @@ class PiiRedactionMiddleware(AgentMiddleware[AgentState]):
|
||||
# -- model-call boundary: genuine user messages ---------------------------
|
||||
|
||||
def _process_request(self, request: ModelRequest) -> ModelRequest:
|
||||
redactor = _Redactor(self._detectors)
|
||||
redactor = _Redactor(self._detectors, self._token_key)
|
||||
messages = list(request.messages)
|
||||
state = getattr(request, "state", None) or {}
|
||||
summary = state.get("summary_text")
|
||||
redactor.reserve(summary)
|
||||
for message in messages:
|
||||
redactor.reserve(message.content)
|
||||
redacted_summary = redactor.redact(summary) if isinstance(summary, str) else summary
|
||||
changed = False
|
||||
for index, msg in enumerate(messages):
|
||||
@ -407,16 +427,13 @@ class PiiRedactionMiddleware(AgentMiddleware[AgentState]):
|
||||
spans the whole result, so placeholder numbering stays continuous
|
||||
across every ToolMessage the result carries.
|
||||
"""
|
||||
redactor = _Redactor(self._detectors)
|
||||
redactor = _Redactor(self._detectors, self._token_key)
|
||||
if isinstance(result, ToolMessage):
|
||||
return self._redact_tool_message(result, redactor)
|
||||
update = getattr(result, "update", None)
|
||||
if isinstance(update, dict):
|
||||
messages = update.get("messages")
|
||||
if isinstance(messages, list) and any(isinstance(m, ToolMessage) for m in messages):
|
||||
for message in messages:
|
||||
if isinstance(message, ToolMessage):
|
||||
redactor.reserve(message.content)
|
||||
new_messages = [self._redact_tool_message(m, redactor) if isinstance(m, ToolMessage) else m for m in messages]
|
||||
if new_messages != messages:
|
||||
return dc_replace(result, update={**update, "messages": new_messages})
|
||||
|
||||
@ -1023,9 +1023,11 @@ def create_summarization_middleware(
|
||||
|
||||
hooks: list[BeforeSummarizationHook] = []
|
||||
if resolved_app_config.memory.enabled and not skip_memory_flush:
|
||||
from functools import partial
|
||||
|
||||
from deerflow.agents.memory.summarization_hook import memory_flush_hook
|
||||
|
||||
hooks.append(memory_flush_hook)
|
||||
hooks.append(partial(memory_flush_hook, pii_redaction_config=getattr(resolved_app_config, "pii_redaction", None)))
|
||||
|
||||
return DeerFlowSummarizationMiddleware(
|
||||
**kwargs,
|
||||
|
||||
@ -1,6 +1,10 @@
|
||||
"""Configuration for the PII redaction middleware (issue #3190)."""
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
# Minimum accepted token_secret length: a short secret makes the keyed digest
|
||||
# brute-forceable offline for low-entropy identifiers (e.g. phone numbers).
|
||||
MIN_TOKEN_SECRET_LENGTH = 16
|
||||
|
||||
|
||||
class PiiRedactionConfig(BaseModel):
|
||||
@ -8,7 +12,9 @@ class PiiRedactionConfig(BaseModel):
|
||||
|
||||
Default-off. When enabled, personally identifiable information found in
|
||||
genuine user messages and remote-content tool results is rewritten to
|
||||
irreversible placeholders (``[EMAIL_1]`` …) before it reaches the model.
|
||||
keyed, value-derived placeholders (128-bit HMAC digests encoded over
|
||||
base-26 letters, keyed by the required ``token_secret``) before it reaches
|
||||
the model.
|
||||
Each detector can be toggled independently for deployments that only need
|
||||
a subset (e.g. credentials but not phone numbers).
|
||||
"""
|
||||
@ -37,3 +43,32 @@ class PiiRedactionConfig(BaseModel):
|
||||
default=True,
|
||||
description="Redact national IDs (CN resident ID and CPF with checksum validation, formatted CUIT/RFC)",
|
||||
)
|
||||
token_secret: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Deployment-scoped secret used as the HMAC key for placeholder "
|
||||
"digests; required (non-empty, at least 16 characters) whenever "
|
||||
"enabled is true. Tokens are linkable only within this deployment "
|
||||
"and cannot be re-derived offline without the secret — an unkeyed "
|
||||
"digest would instead be a publicly computable, globally linkable "
|
||||
"fingerprint of the raw values."
|
||||
),
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _token_secret_required_when_enabled(self) -> "PiiRedactionConfig":
|
||||
"""Reject enabling redaction without a usable key (review round 11 on #5577).
|
||||
|
||||
An unkeyed digest is a publicly computable fingerprint: linkable
|
||||
across deployments and confirmable by offline guessing for low-entropy
|
||||
values such as phone numbers, which defeats the module's
|
||||
re-identification guarantee for exactly the values redaction exists to
|
||||
protect.
|
||||
"""
|
||||
if not self.enabled:
|
||||
return self
|
||||
if self.token_secret is None or not self.token_secret.strip():
|
||||
raise ValueError("pii_redaction.enabled=true requires a non-empty token_secret: without a deployment-scoped HMAC key, placeholders are publicly computable fingerprints linkable across deployments")
|
||||
if len(self.token_secret) < MIN_TOKEN_SECRET_LENGTH:
|
||||
raise ValueError(f"token_secret must be at least {MIN_TOKEN_SECRET_LENGTH} characters; use a long random value")
|
||||
return self
|
||||
|
||||
@ -200,6 +200,7 @@ def test_version_26_config_upgrades_to_checkpoint_channel_mode(tmp_path, caplog)
|
||||
assert upgraded["verification"]["judge_model_name"] is None
|
||||
|
||||
|
||||
@pytest.mark.skipif(SCRIPT_BASH is None, reason="repo shell-script tests need Git Bash on Windows")
|
||||
def test_version_41_config_moves_legacy_ragflow_settings_to_tool(tmp_path):
|
||||
"""The v46 migration keeps provider settings on the RAGFlow tool entry."""
|
||||
import subprocess
|
||||
@ -234,7 +235,7 @@ def test_version_41_config_moves_legacy_ragflow_settings_to_tool(tmp_path):
|
||||
|
||||
env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)}
|
||||
result = subprocess.run(
|
||||
["bash", str(repo_root / "scripts" / "config-upgrade.sh")],
|
||||
[SCRIPT_BASH, str(repo_root / "scripts" / "config-upgrade.sh")],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
@ -254,6 +255,7 @@ def test_version_41_config_moves_legacy_ragflow_settings_to_tool(tmp_path):
|
||||
assert tool["api_key"] == "$CURRENT_RAGFLOW_API_KEY"
|
||||
|
||||
|
||||
@pytest.mark.skipif(SCRIPT_BASH is None, reason="repo shell-script tests need Git Bash on Windows")
|
||||
def test_version_41_tools_only_ragflow_config_enables_knowledge_capability(tmp_path):
|
||||
"""Tools-only legacy configs must not be disabled by the new capability gate."""
|
||||
import subprocess
|
||||
@ -282,7 +284,7 @@ def test_version_41_tools_only_ragflow_config_enables_knowledge_capability(tmp_p
|
||||
|
||||
env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)}
|
||||
result = subprocess.run(
|
||||
["bash", str(repo_root / "scripts" / "config-upgrade.sh")],
|
||||
[SCRIPT_BASH, str(repo_root / "scripts" / "config-upgrade.sh")],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
@ -517,3 +519,109 @@ def test_config_upgrade_adds_security_fail_closed_preserving_user_values():
|
||||
assert user["skill_evolution"]["enabled"] is True
|
||||
assert user["skill_evolution"]["moderation_model_name"] == "custom-moderation-model"
|
||||
assert user["config_version"] == example["config_version"]
|
||||
|
||||
|
||||
def test_version_46_pii_enabled_config_reported_outdated_against_example(caplog):
|
||||
"""token_secret became mandatory for enabled redaction in v47; a v46
|
||||
deployment with redaction on must be flagged outdated so the upgrade
|
||||
warning fires instead of a raw startup validation error with no guidance."""
|
||||
example = _load_repo_example()
|
||||
example_version = example["config_version"]
|
||||
assert example_version >= 47, "config.example.yaml must be bumped past 46 for the mandatory token_secret"
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
config_path = _make_config_files(
|
||||
Path(tmpdir),
|
||||
user_config={"config_version": 46, "pii_redaction": {"enabled": True}},
|
||||
example_config=example,
|
||||
)
|
||||
with caplog.at_level(logging.WARNING, logger="deerflow.config.app_config"):
|
||||
AppConfig._check_config_version({"config_version": 46}, config_path)
|
||||
assert "outdated" in caplog.text
|
||||
assert "(version 46)" in caplog.text
|
||||
assert f"version is {example_version}" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.skipif(SCRIPT_BASH is None, reason="repo shell-script tests need Git Bash on Windows")
|
||||
def test_version_46_pii_enabled_config_upgrade_generates_token_secret(tmp_path):
|
||||
"""Upgrading a v46 config with redaction enabled persists a generated
|
||||
token_secret, leaving the deployment startable under the v47 mandatory
|
||||
validation instead of failing startup after the version bump."""
|
||||
import subprocess
|
||||
|
||||
repo_root = Path(__file__).resolve().parents[2]
|
||||
example_src = repo_root / "config.example.yaml"
|
||||
expected_version = yaml.safe_load(example_src.read_text(encoding="utf-8"))["config_version"]
|
||||
assert expected_version >= 47
|
||||
|
||||
config_path = tmp_path / "config.yaml"
|
||||
(tmp_path / "config.example.yaml").write_text(example_src.read_text(encoding="utf-8"), encoding="utf-8")
|
||||
config_path.write_text(
|
||||
yaml.dump(
|
||||
{
|
||||
"config_version": 46,
|
||||
"sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"},
|
||||
"pii_redaction": {"enabled": True},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)}
|
||||
result = subprocess.run(
|
||||
[SCRIPT_BASH, str(repo_root / "scripts" / "config-upgrade.sh")],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
assert "token_secret generated" in result.stdout
|
||||
|
||||
upgraded = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||||
assert upgraded["config_version"] == expected_version
|
||||
assert upgraded["pii_redaction"]["enabled"] is True
|
||||
secret = upgraded["pii_redaction"]["token_secret"]
|
||||
assert isinstance(secret, str) and len(secret) >= 16
|
||||
|
||||
# The upgraded deployment passes the now-mandatory validation.
|
||||
app_config = AppConfig.model_validate(upgraded)
|
||||
assert app_config.pii_redaction.token_secret == secret
|
||||
|
||||
|
||||
@pytest.mark.skipif(SCRIPT_BASH is None, reason="repo shell-script tests need Git Bash on Windows")
|
||||
def test_version_46_pii_disabled_config_upgrade_skips_token_secret(tmp_path):
|
||||
"""The migration must not invent a secret when redaction is off."""
|
||||
import subprocess
|
||||
|
||||
repo_root = Path(__file__).resolve().parents[2]
|
||||
example_src = repo_root / "config.example.yaml"
|
||||
expected_version = yaml.safe_load(example_src.read_text(encoding="utf-8"))["config_version"]
|
||||
assert expected_version >= 47
|
||||
|
||||
config_path = tmp_path / "config.yaml"
|
||||
(tmp_path / "config.example.yaml").write_text(example_src.read_text(encoding="utf-8"), encoding="utf-8")
|
||||
config_path.write_text(
|
||||
yaml.dump(
|
||||
{
|
||||
"config_version": 46,
|
||||
"sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"},
|
||||
"pii_redaction": {"enabled": False},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
env = {**os.environ, "DEER_FLOW_CONFIG_PATH": str(config_path)}
|
||||
result = subprocess.run(
|
||||
[SCRIPT_BASH, str(repo_root / "scripts" / "config-upgrade.sh")],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
upgraded = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||||
assert upgraded["config_version"] == expected_version
|
||||
assert "token_secret" not in upgraded["pii_redaction"]
|
||||
AppConfig.model_validate(upgraded)
|
||||
|
||||
@ -16,6 +16,7 @@ from deerflow.agents.features import Next, Prev, RuntimeFeatures
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware
|
||||
from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware
|
||||
from deerflow.agents.thread_state import DeltaThreadState, ThreadState
|
||||
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
||||
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
||||
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
||||
from deerflow.subagents import SubagentRuntime
|
||||
@ -170,6 +171,49 @@ def test_features_mode(mock_create_agent):
|
||||
assert "ClarificationMiddleware" in mw_types
|
||||
|
||||
|
||||
@patch("deerflow.agents.factory.create_agent")
|
||||
def test_factory_threads_pii_redaction_config(mock_create_agent):
|
||||
"""#5577 review: the SDK-path wiring must not silently drop the knob —
|
||||
both MemoryMiddleware and DurableContextMiddleware carry the config."""
|
||||
mock_create_agent.return_value = MagicMock()
|
||||
pii = PiiRedactionConfig(enabled=True, token_secret="sdk-path-pii-redaction-secret")
|
||||
|
||||
create_deerflow_agent(_make_mock_model(), features=RuntimeFeatures(memory=True), pii_redaction_config=pii)
|
||||
|
||||
middleware = mock_create_agent.call_args[1]["middleware"]
|
||||
memory_mw = next(m for m in middleware if type(m).__name__ == "MemoryMiddleware")
|
||||
durable_mw = next(m for m in middleware if type(m).__name__ == "DurableContextMiddleware")
|
||||
pii_mw = next(m for m in middleware if type(m).__name__ == "PiiRedactionMiddleware")
|
||||
assert memory_mw._pii_redaction_config is pii
|
||||
assert durable_mw._pii_redaction_config is pii
|
||||
|
||||
# Model-call regression: the assembled SDK chain actually redacts user
|
||||
# content at the model boundary (round-10 review).
|
||||
class _Req:
|
||||
def __init__(self, messages):
|
||||
self.messages = list(messages)
|
||||
|
||||
def override(self, **kwargs):
|
||||
copy = object.__new__(type(self))
|
||||
copy.messages = kwargs.get("messages", self.messages)
|
||||
return copy
|
||||
|
||||
captured = {}
|
||||
pii_mw.wrap_model_call(_Req([HumanMessage("reach alice@example.com")]), lambda req: captured.update(messages=req.messages) or "r")
|
||||
assert "alice@example.com" not in str(captured["messages"][0].content)
|
||||
|
||||
|
||||
@patch("deerflow.agents.factory.create_agent")
|
||||
def test_factory_defaults_to_redaction_off(mock_create_agent):
|
||||
mock_create_agent.return_value = MagicMock()
|
||||
|
||||
create_deerflow_agent(_make_mock_model(), features=RuntimeFeatures(memory=True))
|
||||
|
||||
middleware = mock_create_agent.call_args[1]["middleware"]
|
||||
memory_mw = next(m for m in middleware if type(m).__name__ == "MemoryMiddleware")
|
||||
assert memory_mw._pii_redaction_config is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Middleware full takeover
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@ -859,7 +859,7 @@ def test_build_middlewares_passes_explicit_app_config_to_shared_factory(monkeypa
|
||||
monkeypatch.setattr(
|
||||
lead_agent_module,
|
||||
"MemoryMiddleware",
|
||||
lambda agent_name=None, *, memory_config: captured.setdefault("memory_config", memory_config) or "memory-middleware",
|
||||
lambda agent_name=None, *, memory_config, pii_redaction_config=None: captured.setdefault("memory_config", memory_config) or "memory-middleware",
|
||||
)
|
||||
|
||||
middlewares = lead_agent_module.build_middlewares(
|
||||
@ -900,7 +900,7 @@ def test_build_middlewares_passes_run_model_name_to_summarization(monkeypatch):
|
||||
)
|
||||
monkeypatch.setattr(lead_agent_module, "_create_todo_list_middleware", lambda is_plan_mode: None)
|
||||
monkeypatch.setattr(lead_agent_module, "TitleMiddleware", lambda *, app_config, extensions: "title-middleware")
|
||||
monkeypatch.setattr(lead_agent_module, "MemoryMiddleware", lambda agent_name=None, *, memory_config: "memory-middleware")
|
||||
monkeypatch.setattr(lead_agent_module, "MemoryMiddleware", lambda agent_name=None, *, memory_config, pii_redaction_config=None: "memory-middleware")
|
||||
|
||||
lead_agent_module.build_middlewares(
|
||||
{"configurable": {"is_plan_mode": False, "subagent_enabled": False}},
|
||||
|
||||
283
backend/tests/test_memory_pii_redaction.py
Normal file
283
backend/tests/test_memory_pii_redaction.py
Normal file
@ -0,0 +1,283 @@
|
||||
"""Memory-queue PII redaction (#3190 vector 5, follow-up to #5527)."""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from deerflow.agents.human_input import HUMAN_INPUT_RESPONSE_KEY
|
||||
from deerflow.agents.middlewares import memory_middleware as memory_middleware_module
|
||||
from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware
|
||||
from deerflow.agents.middlewares.pii_redaction_middleware import redact_text
|
||||
from deerflow.config.memory_config import MemoryConfig
|
||||
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
|
||||
_TOKEN_SECRET = "unit-test-deployment-secret-0123456789"
|
||||
|
||||
EMAIL_TOKEN = redact_text("alice@example.com", PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
|
||||
|
||||
def _middleware(pii_config):
|
||||
manager = MagicMock()
|
||||
mw = MemoryMiddleware(
|
||||
agent_name="researcher",
|
||||
memory_config=MemoryConfig(enabled=True),
|
||||
pii_redaction_config=pii_config,
|
||||
)
|
||||
return mw, manager
|
||||
|
||||
|
||||
def _run(mw, manager, monkeypatch, messages):
|
||||
monkeypatch.setattr(memory_middleware_module, "get_memory_manager", lambda: manager)
|
||||
runtime = Runtime(context={"thread_id": "thread-123", "user_id": "runtime-user"})
|
||||
mw.after_agent({"messages": messages}, runtime)
|
||||
return manager.add.call_args
|
||||
|
||||
|
||||
def test_queue_payload_redacted_when_enabled(monkeypatch):
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
call = _run(mw, manager, monkeypatch, [HumanMessage("reach alice@example.com"), AIMessage("noted")])
|
||||
queued = call.args[1]
|
||||
assert EMAIL_TOKEN in queued[0].content and "alice@example.com" not in queued[0].content
|
||||
assert queued[1].content == "noted"
|
||||
|
||||
|
||||
def test_queue_payload_untouched_without_config(monkeypatch):
|
||||
mw, manager = _middleware(None)
|
||||
call = _run(mw, manager, monkeypatch, [HumanMessage("reach alice@example.com")])
|
||||
assert "alice@example.com" in call.args[1][0].content
|
||||
|
||||
|
||||
def test_queue_payload_untouched_when_disabled(monkeypatch):
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=False))
|
||||
call = _run(mw, manager, monkeypatch, [HumanMessage("reach alice@example.com")])
|
||||
assert "alice@example.com" in call.args[1][0].content
|
||||
|
||||
|
||||
def test_detector_toggles_respected(monkeypatch):
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, redact_email=False, token_secret=_TOKEN_SECRET))
|
||||
call = _run(mw, manager, monkeypatch, [HumanMessage("reach alice@example.com")])
|
||||
assert "alice@example.com" in call.args[1][0].content
|
||||
|
||||
|
||||
def test_original_messages_not_mutated(monkeypatch):
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
original = HumanMessage("reach alice@example.com")
|
||||
_run(mw, manager, monkeypatch, [original])
|
||||
assert original.content == "reach alice@example.com"
|
||||
|
||||
|
||||
def test_same_value_shares_token_across_turns(monkeypatch):
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
call = _run(
|
||||
mw,
|
||||
manager,
|
||||
monkeypatch,
|
||||
[HumanMessage("alice@example.com"), AIMessage("got alice@example.com")],
|
||||
)
|
||||
queued = call.args[1]
|
||||
assert queued[0].content == EMAIL_TOKEN
|
||||
assert EMAIL_TOKEN in queued[1].content
|
||||
|
||||
|
||||
def test_existing_placeholder_in_later_message_keeps_identity_stable(monkeypatch):
|
||||
# Round-5 review on #5577: sequential numbering made an unchanged
|
||||
# message's token depend on batch contents (a higher existing placeholder
|
||||
# in a later batch would shift it, breaking OpenViking capture dedup).
|
||||
# Value-derived tokens keep the unchanged message's identity stable.
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
alice_token = redact_text("alice@example.com", PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
messages = [
|
||||
HumanMessage("Bob's email is bob@example.com"),
|
||||
AIMessage(f"Alice's email is {alice_token} and Bob's email is [EMAIL_2]"),
|
||||
]
|
||||
first = _run(mw, manager, monkeypatch, messages).args[1]
|
||||
second = _run(mw, manager, monkeypatch, messages).args[1]
|
||||
bob_token = redact_text("bob@example.com", PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
assert first[0].content == f"Bob's email is {bob_token}"
|
||||
assert second[0].content == first[0].content
|
||||
assert f"Alice's email is {alice_token}" in first[1].content
|
||||
|
||||
|
||||
def _flush_hook_call(monkeypatch, pii_config, messages):
|
||||
import deerflow.agents.memory.summarization_hook as hook_module
|
||||
from deerflow.agents.middlewares.summarization_middleware import SummarizationEvent
|
||||
|
||||
manager = MagicMock()
|
||||
monkeypatch.setattr(hook_module, "get_memory_manager", lambda: manager)
|
||||
monkeypatch.setattr(hook_module, "resolve_runtime_user_id", lambda runtime: "u1")
|
||||
event = SummarizationEvent(
|
||||
messages_to_summarize=tuple(messages),
|
||||
preserved_messages=(),
|
||||
thread_id="thread-123",
|
||||
agent_name="researcher",
|
||||
runtime=None,
|
||||
)
|
||||
hook_module.memory_flush_hook(event, pii_redaction_config=pii_config)
|
||||
return manager.add_nowait.call_args.args[1]
|
||||
|
||||
|
||||
def test_compaction_flush_hook_redacts_queued_payload(monkeypatch):
|
||||
# Review round 2 on #5577: the compaction-triggered memory_flush_hook
|
||||
# queues messages that compaction is about to remove from state, so the
|
||||
# after-agent redaction can never repair a raw batch queued here.
|
||||
queued = _flush_hook_call(
|
||||
monkeypatch,
|
||||
PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET),
|
||||
[HumanMessage("reach alice@example.com"), AIMessage("noted")],
|
||||
)
|
||||
assert EMAIL_TOKEN in queued[0].content and "alice@example.com" not in queued[0].content
|
||||
assert queued[1].content == "noted"
|
||||
|
||||
|
||||
def test_compaction_flush_hook_untouched_without_config(monkeypatch):
|
||||
queued = _flush_hook_call(monkeypatch, None, [HumanMessage("reach alice@example.com")])
|
||||
assert "alice@example.com" in queued[0].content
|
||||
|
||||
|
||||
def test_compaction_flush_hook_untouched_when_disabled(monkeypatch):
|
||||
queued = _flush_hook_call(monkeypatch, PiiRedactionConfig(enabled=False), [HumanMessage("reach alice@example.com")])
|
||||
assert "alice@example.com" in queued[0].content
|
||||
|
||||
|
||||
def test_tool_call_args_redacted(monkeypatch):
|
||||
# OpenViking-style retention keeps the full message object, including
|
||||
# parsed tool_calls and provider-format arguments in additional_kwargs.
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
ai = AIMessage(
|
||||
content="",
|
||||
tool_calls=[{"name": "memory_search", "args": {"query": "alice@example.com"}, "id": "call_1"}],
|
||||
additional_kwargs={
|
||||
"tool_calls": [
|
||||
{"function": {"name": "memory_search", "arguments": '{"query": "alice@example.com"}'}},
|
||||
],
|
||||
},
|
||||
)
|
||||
call = _run(mw, manager, monkeypatch, [HumanMessage("hi"), ai])
|
||||
queued = call.args[1][1]
|
||||
assert "alice@example.com" not in str(queued.tool_calls)
|
||||
assert EMAIL_TOKEN in str(queued.tool_calls)
|
||||
assert "alice@example.com" not in str(queued.additional_kwargs.get("tool_calls"))
|
||||
assert EMAIL_TOKEN in str(queued.additional_kwargs.get("tool_calls"))
|
||||
|
||||
|
||||
def test_non_text_block_text_field_redacted(monkeypatch):
|
||||
# DeerMem's format_conversation_for_update reads the "text" value of any
|
||||
# dict block, not only type=="text" — the helper must follow.
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
call = _run(mw, manager, monkeypatch, [HumanMessage([{"type": "custom_card", "text": "alice@example.com"}])])
|
||||
block = call.args[1][0].content[0]
|
||||
assert block["text"] == EMAIL_TOKEN
|
||||
|
||||
|
||||
def test_tool_argument_keys_redacted(monkeypatch):
|
||||
# Review round 6: tool arguments can carry user data in mapping keys.
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
ai = AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{"name": "memory_search", "args": {"contacts": {"alice@example.com": "manager"}}, "id": "call_1"},
|
||||
],
|
||||
)
|
||||
call = _run(mw, manager, monkeypatch, [HumanMessage("hi"), ai])
|
||||
args = call.args[1][1].tool_calls[0]["args"]
|
||||
assert args["contacts"] == {EMAIL_TOKEN: "manager"}
|
||||
assert "alice@example.com" not in str(args)
|
||||
|
||||
|
||||
def test_distinct_identities_keep_distinct_tokens():
|
||||
# Review round 6: a 24-bit truncation collided distinct identities; at
|
||||
# 128 bits the review's collision pair stays distinct, including when
|
||||
# both appear together in one message.
|
||||
cfg = PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)
|
||||
a = redact_text("contact3513@example.com", cfg)
|
||||
b = redact_text("contact3727@example.com", cfg)
|
||||
assert a != b
|
||||
both = redact_text("contact3513@example.com and contact3727@example.com", cfg)
|
||||
assert a in both and b in both
|
||||
|
||||
|
||||
def test_original_user_content_provenance_redacted(monkeypatch):
|
||||
# Review round 12 on #5577: UploadsMiddleware preserves the raw user turn
|
||||
# in additional_kwargs.original_user_content, so a redacted content alone
|
||||
# still leaks the raw value to the memory backend. The queued copy must
|
||||
# redact the provenance field too, and the thread-state message must stay
|
||||
# untouched.
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
original = HumanMessage("reach alice@example.com", additional_kwargs={ORIGINAL_USER_CONTENT_KEY: "reach alice@example.com"})
|
||||
_run(mw, manager, monkeypatch, [original])
|
||||
queued = manager.add.call_args.args[1][0]
|
||||
assert EMAIL_TOKEN in queued.content and "alice@example.com" not in queued.content
|
||||
assert EMAIL_TOKEN in queued.additional_kwargs[ORIGINAL_USER_CONTENT_KEY]
|
||||
assert "alice@example.com" not in queued.additional_kwargs[ORIGINAL_USER_CONTENT_KEY]
|
||||
assert original.additional_kwargs[ORIGINAL_USER_CONTENT_KEY] == "reach alice@example.com"
|
||||
assert original.content == "reach alice@example.com"
|
||||
|
||||
|
||||
def test_original_user_content_redacted_in_compaction_flush_hook(monkeypatch):
|
||||
# Same provenance leak through the compaction-triggered hook, which shares
|
||||
# redact_queued_messages: one fix must cover both enqueues.
|
||||
queued = _flush_hook_call(
|
||||
monkeypatch,
|
||||
PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET),
|
||||
[HumanMessage("reach alice@example.com", additional_kwargs={ORIGINAL_USER_CONTENT_KEY: "reach alice@example.com"})],
|
||||
)
|
||||
assert EMAIL_TOKEN in queued[0].additional_kwargs[ORIGINAL_USER_CONTENT_KEY]
|
||||
assert "alice@example.com" not in queued[0].additional_kwargs[ORIGINAL_USER_CONTENT_KEY]
|
||||
|
||||
|
||||
def _human_input_response(value: str) -> dict:
|
||||
return {
|
||||
"version": 1,
|
||||
"kind": "human_input_response",
|
||||
"source": "clarification",
|
||||
"request_id": "req-1",
|
||||
"response_kind": "text",
|
||||
"value": value,
|
||||
}
|
||||
|
||||
|
||||
def test_human_input_response_value_redacted(monkeypatch):
|
||||
# Review round 13 on #5577: clarification replies keep the raw user answer
|
||||
# in additional_kwargs.human_input_response.value (DeerMem reads that
|
||||
# mapping), so the queued copy must redact the value while preserving the
|
||||
# protocol fields. The same message also carries original_user_content —
|
||||
# both rewrites must survive in one merged additional_kwargs update.
|
||||
mw, manager = _middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
original = HumanMessage(
|
||||
"my email is alice@example.com",
|
||||
additional_kwargs={
|
||||
HUMAN_INPUT_RESPONSE_KEY: _human_input_response("my email is alice@example.com"),
|
||||
ORIGINAL_USER_CONTENT_KEY: "my email is alice@example.com",
|
||||
},
|
||||
)
|
||||
_run(mw, manager, monkeypatch, [original])
|
||||
queued = manager.add.call_args.args[1][0]
|
||||
assert EMAIL_TOKEN in queued.content and "alice@example.com" not in queued.content
|
||||
|
||||
queued_response = queued.additional_kwargs[HUMAN_INPUT_RESPONSE_KEY]
|
||||
assert EMAIL_TOKEN in queued_response["value"] and "alice@example.com" not in queued_response["value"]
|
||||
assert queued_response["version"] == 1
|
||||
assert queued_response["kind"] == "human_input_response"
|
||||
assert queued_response["source"] == "clarification"
|
||||
assert queued_response["request_id"] == "req-1"
|
||||
assert queued_response["response_kind"] == "text"
|
||||
assert EMAIL_TOKEN in queued.additional_kwargs[ORIGINAL_USER_CONTENT_KEY]
|
||||
assert "alice@example.com" not in str(queued.additional_kwargs)
|
||||
|
||||
assert original.additional_kwargs[HUMAN_INPUT_RESPONSE_KEY]["value"] == "my email is alice@example.com"
|
||||
assert original.additional_kwargs[ORIGINAL_USER_CONTENT_KEY] == "my email is alice@example.com"
|
||||
assert original.content == "my email is alice@example.com"
|
||||
|
||||
|
||||
def test_human_input_response_redacted_in_compaction_flush_hook(monkeypatch):
|
||||
# Same leak through the compaction-triggered hook, which shares
|
||||
# redact_queued_messages: one fix must cover both enqueues.
|
||||
queued = _flush_hook_call(
|
||||
monkeypatch,
|
||||
PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET),
|
||||
[HumanMessage("my email is alice@example.com", additional_kwargs={HUMAN_INPUT_RESPONSE_KEY: _human_input_response("my email is alice@example.com")})],
|
||||
)
|
||||
assert EMAIL_TOKEN in queued[0].additional_kwargs[HUMAN_INPUT_RESPONSE_KEY]["value"]
|
||||
assert "alice@example.com" not in str(queued[0].additional_kwargs[HUMAN_INPUT_RESPONSE_KEY])
|
||||
@ -1,9 +1,9 @@
|
||||
"""Tests for PiiRedactionMiddleware (issue #3190).
|
||||
|
||||
Verifies deterministic detector coverage (including checksum gates), stable
|
||||
placeholder numbering across a conversation, that the rewrite is request-scoped
|
||||
without mutating the original request or messages, the tool-boundary allowlist,
|
||||
and the pinned detector registry.
|
||||
Verifies deterministic detector coverage (including checksum gates),
|
||||
value-derived placeholder identity (stable across turns, batches, and seams),
|
||||
that the rewrite is request-scoped without mutating the original request or
|
||||
messages, the tool-boundary allowlist, and the pinned detector registry.
|
||||
"""
|
||||
|
||||
import re
|
||||
@ -16,7 +16,7 @@ from langchain.agents import create_agent
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage, get_buffer_string
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
from langgraph.types import Command
|
||||
from pydantic import Field
|
||||
from pydantic import Field, ValidationError
|
||||
|
||||
from deerflow.agents.middlewares.pii_redaction_middleware import (
|
||||
_DETECTORS,
|
||||
@ -26,9 +26,31 @@ from deerflow.agents.middlewares.pii_redaction_middleware import (
|
||||
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
||||
from deerflow.tools.mcp_metadata import MCP_TOOL_METADATA_KEY
|
||||
|
||||
_TOKEN_SECRET = "unit-test-deployment-secret-0123456789"
|
||||
|
||||
_PII_CFG = PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)
|
||||
|
||||
EMAIL_ALICE = redact_text("alice@example.com", _PII_CFG)
|
||||
EMAIL_BOB_COM = redact_text("bob@example.com", _PII_CFG)
|
||||
EMAIL_BOB_ORG = redact_text("bob@example.org", _PII_CFG)
|
||||
EMAIL_CAROL = redact_text("carol@example.com", _PII_CFG)
|
||||
EMAIL_CHARLIE = redact_text("charlie@example.net", _PII_CFG)
|
||||
PHONE_INTL = redact_text("+86 138 0013 8000", _PII_CFG)
|
||||
PHONE_CN = redact_text("13800138000", _PII_CFG)
|
||||
PHONE_US = redact_text("(212) 555-0123", _PII_CFG)
|
||||
PHONE_US2 = redact_text("+1 415 555 2671", _PII_CFG)
|
||||
KEY_SK = redact_text("sk-proj4aaaaaaaaaaaaaaaaaaaaaaaaaaaa", _PII_CFG)
|
||||
KEY_AWS = redact_text("AKIAIOSFODNN7EXAMPLE", _PII_CFG)
|
||||
CARD_VISA = redact_text("4111 1111 1111 1111", _PII_CFG)
|
||||
ID_X = redact_text("11010519491231002X", _PII_CFG)
|
||||
ID_L150 = redact_text("110105194912310150", _PII_CFG)
|
||||
ID_W239 = redact_text("110105197506150239", _PII_CFG)
|
||||
ID_CUIT = redact_text("20-12345678-6", _PII_CFG)
|
||||
ID_CPF = redact_text("529.982.247-25", _PII_CFG)
|
||||
|
||||
|
||||
def _make_middleware(**config_overrides) -> PiiRedactionMiddleware:
|
||||
return PiiRedactionMiddleware(PiiRedactionConfig(enabled=True, **config_overrides))
|
||||
return PiiRedactionMiddleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET, **config_overrides))
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
@ -72,7 +94,7 @@ class TestDetectors:
|
||||
# national-id detector must claim it before the credit-card detector
|
||||
# (review finding on #5527, reproduced at 0a2a9d0).
|
||||
messages, _ = _run_model_call(_make_middleware(), [HumanMessage("id 110105194912310150")])
|
||||
assert "id [NATIONAL_ID_1]" in messages[0].content
|
||||
assert f"id {ID_L150}" in messages[0].content
|
||||
|
||||
def test_email_redacted(self):
|
||||
result = _make_middleware()._detectors[0].pattern.sub("X", "ping me at alice@example.com today")
|
||||
@ -84,7 +106,7 @@ class TestDetectors:
|
||||
middleware,
|
||||
[HumanMessage("from alice@example.com to bob@example.org")],
|
||||
)
|
||||
assert "from [EMAIL_1] to [EMAIL_2]" in messages[0].content
|
||||
assert f"from {EMAIL_ALICE} to {EMAIL_BOB_ORG}" in messages[0].content
|
||||
|
||||
def test_same_email_shares_placeholder(self):
|
||||
middleware = _make_middleware()
|
||||
@ -95,29 +117,29 @@ class TestDetectors:
|
||||
HumanMessage("reply to alice@example.com"),
|
||||
],
|
||||
)
|
||||
assert messages[0].content == "[EMAIL_1] here"
|
||||
assert messages[1].content == "reply to [EMAIL_1]"
|
||||
assert messages[0].content == f"{EMAIL_ALICE} here"
|
||||
assert messages[1].content == f"reply to {EMAIL_ALICE}"
|
||||
|
||||
def test_openai_style_api_key_redacted(self):
|
||||
messages, _ = _run_model_call(
|
||||
_make_middleware(),
|
||||
[HumanMessage("key: sk-proj4aaaaaaaaaaaaaaaaaaaaaaaaaaaa")],
|
||||
)
|
||||
assert "[API_KEY_1]" in messages[0].content
|
||||
assert KEY_SK in messages[0].content
|
||||
|
||||
def test_aws_access_key_redacted(self):
|
||||
messages, _ = _run_model_call(
|
||||
_make_middleware(),
|
||||
[HumanMessage("use AKIAIOSFODNN7EXAMPLE please")],
|
||||
)
|
||||
assert "[API_KEY_1]" in messages[0].content
|
||||
assert KEY_AWS in messages[0].content
|
||||
|
||||
def test_credit_card_luhn_valid_redacted(self):
|
||||
messages, _ = _run_model_call(
|
||||
_make_middleware(),
|
||||
[HumanMessage("card 4111 1111 1111 1111 on file")],
|
||||
)
|
||||
assert "card [CREDIT_CARD_1] on file" in messages[0].content
|
||||
assert f"card {CARD_VISA} on file" in messages[0].content
|
||||
|
||||
def test_credit_card_luhn_invalid_untouched(self):
|
||||
original = "card 1234 5678 9012 3456 on file"
|
||||
@ -134,22 +156,22 @@ class TestDetectors:
|
||||
_make_middleware(),
|
||||
[HumanMessage("call +86 138 0013 8000 now")],
|
||||
)
|
||||
assert "call [PHONE_1] now" in messages[0].content
|
||||
assert f"call {PHONE_INTL} now" in messages[0].content
|
||||
|
||||
def test_cn_mobile_redacted(self):
|
||||
messages, _ = _run_model_call(_make_middleware(), [HumanMessage("phone 13800138000")])
|
||||
assert "phone [PHONE_1]" in messages[0].content
|
||||
assert f"phone {PHONE_CN}" in messages[0].content
|
||||
|
||||
def test_us_phone_redacted(self):
|
||||
messages, _ = _run_model_call(_make_middleware(), [HumanMessage("dial (212) 555-0123")])
|
||||
assert "dial [PHONE_1]" in messages[0].content
|
||||
assert f"dial {PHONE_US}" in messages[0].content
|
||||
|
||||
def test_cn_resident_id_valid_redacted(self):
|
||||
messages, _ = _run_model_call(
|
||||
_make_middleware(),
|
||||
[HumanMessage("id 11010519491231002X")],
|
||||
)
|
||||
assert "id [NATIONAL_ID_1]" in messages[0].content
|
||||
assert f"id {ID_X}" in messages[0].content
|
||||
|
||||
def test_cn_resident_id_invalid_checksum_untouched(self):
|
||||
original = "id 110105194912310020"
|
||||
@ -160,11 +182,11 @@ class TestDetectors:
|
||||
# Review vector on #5527: numeric-check-digit resident ID whose digits
|
||||
# also pass Luhn must render NATIONAL_ID, not CREDIT_CARD.
|
||||
messages, _ = _run_model_call(_make_middleware(), [HumanMessage("id 110105197506150239")])
|
||||
assert "id [NATIONAL_ID_1]" in messages[0].content
|
||||
assert f"id {ID_W239}" in messages[0].content
|
||||
|
||||
def test_cuit_valid_form_redacted(self):
|
||||
messages, _ = _run_model_call(_make_middleware(), [HumanMessage("CUIT 20-12345678-6")])
|
||||
assert "CUIT [NATIONAL_ID_1]" in messages[0].content
|
||||
assert f"CUIT {ID_CUIT}" in messages[0].content
|
||||
|
||||
def test_cuit_wrong_digit_count_untouched(self):
|
||||
original = "CUIT 20-1234567890-6"
|
||||
@ -179,7 +201,7 @@ class TestDetectors:
|
||||
[HumanMessage("身份证11010519491231002X 手机号13800138000 信用卡4111 1111 1111 1111")],
|
||||
)
|
||||
content = messages[0].content
|
||||
assert "[NATIONAL_ID_1]" in content and "[PHONE_1]" in content and "[CREDIT_CARD_1]" in content
|
||||
assert ID_X in content and PHONE_CN in content and CARD_VISA in content
|
||||
assert "11010519491231002X" not in content and "13800138000" not in content and "4111" not in content
|
||||
|
||||
def test_international_phone_does_not_consume_next_line(self):
|
||||
@ -187,14 +209,14 @@ class TestDetectors:
|
||||
_make_middleware(),
|
||||
[HumanMessage("Call +1 415 555 2671\n20260918")],
|
||||
)
|
||||
assert messages[0].content == "Call [PHONE_1]\n20260918"
|
||||
assert messages[0].content == f"Call {PHONE_US2}\n20260918"
|
||||
|
||||
def test_cpf_valid_redacted(self):
|
||||
messages, _ = _run_model_call(
|
||||
_make_middleware(),
|
||||
[HumanMessage("cpf 529.982.247-25")],
|
||||
)
|
||||
assert "cpf [NATIONAL_ID_1]" in messages[0].content
|
||||
assert f"cpf {ID_CPF}" in messages[0].content
|
||||
|
||||
def test_cpf_invalid_untouched(self):
|
||||
original = "cpf 529.982.247-11"
|
||||
@ -213,12 +235,12 @@ class TestModelCallBoundary:
|
||||
_make_middleware(),
|
||||
[HumanMessage("my email is alice@example.com")],
|
||||
)
|
||||
assert messages[0].content == "my email is [EMAIL_1]"
|
||||
assert messages[0].content == f"my email is {EMAIL_ALICE}"
|
||||
|
||||
def test_original_request_not_mutated(self):
|
||||
original = HumanMessage("my email is alice@example.com")
|
||||
messages, request = _run_model_call(_make_middleware(), [original])
|
||||
assert messages[0].content == "my email is [EMAIL_1]"
|
||||
assert messages[0].content == f"my email is {EMAIL_ALICE}"
|
||||
assert request.messages[0].content == "my email is alice@example.com"
|
||||
|
||||
def test_additional_kwargs_preserved(self):
|
||||
@ -245,14 +267,14 @@ class TestModelCallBoundary:
|
||||
HumanMessage("then bob@example.org"),
|
||||
],
|
||||
)
|
||||
assert "first [EMAIL_1]" in messages[0].content
|
||||
assert "then [EMAIL_2]" in messages[2].content
|
||||
assert f"first {EMAIL_ALICE}" in messages[0].content
|
||||
assert f"then {EMAIL_BOB_ORG}" in messages[2].content
|
||||
|
||||
def test_redaction_deterministic_across_calls(self):
|
||||
middleware = _make_middleware()
|
||||
messages_a, _ = _run_model_call(middleware, [HumanMessage("alice@example.com")])
|
||||
messages_b, _ = _run_model_call(middleware, [HumanMessage("alice@example.com")])
|
||||
assert messages_a[0].content == messages_b[0].content == "[EMAIL_1]"
|
||||
assert messages_a[0].content == messages_b[0].content == EMAIL_ALICE
|
||||
|
||||
def test_disabled_detector_untouched(self):
|
||||
messages, _ = _run_model_call(
|
||||
@ -271,10 +293,10 @@ class TestModelCallBoundary:
|
||||
]
|
||||
)
|
||||
messages, _ = _run_model_call(_make_middleware(), [original])
|
||||
assert messages[0].content[0] == "reach me at [EMAIL_1]"
|
||||
assert messages[0].content[0] == f"reach me at {EMAIL_ALICE}"
|
||||
# LangChain rebuilds content blocks on construction, so compare by value.
|
||||
assert messages[0].content[1] == image_block
|
||||
assert messages[0].content[2] == "or [EMAIL_2]"
|
||||
assert messages[0].content[2] == f"or {EMAIL_BOB_ORG}"
|
||||
# The original message object is untouched.
|
||||
assert original.content[0] == "reach me at alice@example.com"
|
||||
|
||||
@ -292,7 +314,7 @@ class TestToolBoundary:
|
||||
name="web_fetch",
|
||||
)
|
||||
final = _run_tool_call(_make_middleware(), "web_fetch", result)
|
||||
assert final.content == "page says contact [EMAIL_1]"
|
||||
assert final.content == f"page says contact {EMAIL_ALICE}"
|
||||
transforms = final.additional_kwargs["deerflow_tool_transforms"]
|
||||
assert transforms[-1]["kind"] == "pii_redaction"
|
||||
assert transforms[-1]["by"] == "PiiRedactionMiddleware"
|
||||
@ -314,7 +336,7 @@ class TestToolBoundary:
|
||||
)
|
||||
tool = SimpleNamespace(metadata={MCP_TOOL_METADATA_KEY: True})
|
||||
final = _run_tool_call(_make_middleware(), "fetch_url", result, tool=tool)
|
||||
assert final.content == "[EMAIL_1]"
|
||||
assert final.content == EMAIL_ALICE
|
||||
|
||||
def test_command_result_passthrough(self):
|
||||
result = Command(update={"events": ["alice@example.com"]})
|
||||
@ -327,7 +349,7 @@ class TestToolBoundary:
|
||||
final = _run_tool_call(_make_middleware(), "web_fetch", result)
|
||||
assert isinstance(final, Command)
|
||||
new_message = final.update["messages"][0]
|
||||
assert new_message.content == "page says [EMAIL_1]"
|
||||
assert new_message.content == f"page says {EMAIL_ALICE}"
|
||||
assert new_message.additional_kwargs["deerflow_tool_transforms"][-1]["kind"] == "pii_redaction"
|
||||
# The original Command and its message are untouched.
|
||||
assert tool_message.content == "page says alice@example.com"
|
||||
@ -346,7 +368,7 @@ class TestToolBoundary:
|
||||
response_metadata={"latency_ms": 12},
|
||||
)
|
||||
final = _run_tool_call(_make_middleware(), "web_fetch", result)
|
||||
assert final.content == "[EMAIL_1]"
|
||||
assert final.content == EMAIL_ALICE
|
||||
assert final.artifact == {"rows": 3}
|
||||
assert final.response_metadata == {"latency_ms": 12}
|
||||
assert final.status == "success"
|
||||
@ -368,8 +390,8 @@ class TestToolBoundary:
|
||||
result = Command(update={"messages": [first, second]})
|
||||
final = _run_tool_call(_make_middleware(), "web_fetch", result)
|
||||
messages = final.update["messages"]
|
||||
assert messages[0].content == "[EMAIL_1]"
|
||||
assert messages[1].content == "then [EMAIL_2] and [EMAIL_1]"
|
||||
assert messages[0].content == EMAIL_ALICE
|
||||
assert messages[1].content == f"then {EMAIL_BOB_ORG} and {EMAIL_ALICE}"
|
||||
|
||||
def test_placeholder_restarts_per_result(self):
|
||||
middleware = _make_middleware()
|
||||
@ -383,8 +405,8 @@ class TestToolBoundary:
|
||||
"web_fetch",
|
||||
ToolMessage(content="bob@example.com", tool_call_id="c2", name="web_fetch"),
|
||||
)
|
||||
assert first.content == "[EMAIL_1]"
|
||||
assert second.content == "[EMAIL_1]"
|
||||
assert first.content == EMAIL_ALICE
|
||||
assert second.content == EMAIL_BOB_COM
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -406,7 +428,7 @@ class TestReleasePolicy:
|
||||
"config",
|
||||
[
|
||||
PiiRedactionConfig(enabled=False),
|
||||
PiiRedactionConfig(enabled=True),
|
||||
PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET),
|
||||
],
|
||||
)
|
||||
def test_config_defaults_are_consistent(config):
|
||||
@ -441,7 +463,7 @@ class TestChainWiring:
|
||||
from deerflow.agents.middlewares.tool_result_sanitization_middleware import ToolResultSanitizationMiddleware
|
||||
|
||||
middlewares = build_lead_runtime_middlewares(
|
||||
app_config=_wiring_app_config(pii_redaction=PiiRedactionConfig(enabled=True)),
|
||||
app_config=_wiring_app_config(pii_redaction=PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)),
|
||||
)
|
||||
types = [type(m) for m in middlewares]
|
||||
assert PiiRedactionMiddleware in types
|
||||
@ -452,7 +474,7 @@ class TestChainWiring:
|
||||
from deerflow.agents.middlewares.tool_error_handling_middleware import build_subagent_runtime_middlewares
|
||||
|
||||
middlewares = build_subagent_runtime_middlewares(
|
||||
app_config=_wiring_app_config(pii_redaction=PiiRedactionConfig(enabled=True)),
|
||||
app_config=_wiring_app_config(pii_redaction=PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)),
|
||||
)
|
||||
assert PiiRedactionMiddleware in [type(m) for m in middlewares]
|
||||
|
||||
@ -470,10 +492,10 @@ class TestRedactTextSharedSeam:
|
||||
assert redact_text("alice@example.com", PiiRedactionConfig(enabled=False)) == "alice@example.com"
|
||||
|
||||
def test_enabled_config_redacts(self):
|
||||
assert redact_text("call alice@example.com", PiiRedactionConfig(enabled=True)) == "call [EMAIL_1]"
|
||||
assert redact_text("call alice@example.com", PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)) == f"call {EMAIL_ALICE}"
|
||||
|
||||
def test_non_string_passthrough(self):
|
||||
assert redact_text(None, PiiRedactionConfig(enabled=True)) is None
|
||||
assert redact_text(None, PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)) is None
|
||||
|
||||
|
||||
class _StateRequest:
|
||||
@ -497,13 +519,13 @@ class TestDurableContextReinjection:
|
||||
return DurableContextMiddleware(pii_redaction_config=config)
|
||||
|
||||
def test_reinjected_summary_redacted(self):
|
||||
mw = self._make_dc(PiiRedactionConfig(enabled=True))
|
||||
mw = self._make_dc(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
request = _StateRequest({"summary_text": "summary of alice@example.com"}, [HumanMessage("hi")])
|
||||
final = mw._inject(request)
|
||||
# insert_after_leading_system_messages puts the injected pair up front:
|
||||
# [authority SystemMessage, durable-context data block, original…].
|
||||
block = final.messages[1].content
|
||||
assert "[EMAIL_1]" in block and "alice@example.com" not in block
|
||||
assert EMAIL_ALICE in block and "alice@example.com" not in block
|
||||
|
||||
def test_reinjected_summary_untouched_without_config(self):
|
||||
mw = self._make_dc(None)
|
||||
@ -512,7 +534,7 @@ class TestDurableContextReinjection:
|
||||
assert "alice@example.com" in final.messages[1].content
|
||||
|
||||
def test_policy_declares_pii_gate(self):
|
||||
enabled = self._make_dc(PiiRedactionConfig(enabled=True)).release_policy_parameters()
|
||||
enabled = self._make_dc(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)).release_policy_parameters()
|
||||
disabled = self._make_dc(None).release_policy_parameters()
|
||||
assert enabled["pii_redaction_enabled"] is True
|
||||
assert disabled["pii_redaction_enabled"] is False
|
||||
@ -535,10 +557,10 @@ class TestSummarizationCompactionInput:
|
||||
)
|
||||
|
||||
def test_compaction_input_redacted(self):
|
||||
mw = self._middleware(PiiRedactionConfig(enabled=True))
|
||||
mw = self._middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
prompt = mw._build_summary_prompt([HumanMessage("reach alice@example.com")], previous_summary=None)
|
||||
assert prompt is not None
|
||||
assert "[EMAIL_1]" in prompt and "alice@example.com" not in prompt
|
||||
assert EMAIL_ALICE in prompt and "alice@example.com" not in prompt
|
||||
|
||||
def test_compaction_input_untouched_when_disabled(self):
|
||||
mw = self._middleware(PiiRedactionConfig(enabled=False))
|
||||
@ -555,7 +577,7 @@ class _RecordingPiiModel(FakeToolCallingModel):
|
||||
self.seen.append(text)
|
||||
if self.echo_summary:
|
||||
# Preserve the exact placeholder received, rather than inventing one.
|
||||
token = re.search(r"\[EMAIL_[0-9]+\]", text).group(0)
|
||||
token = re.search(r"\[EMAIL_[a-z]{27}\]", text).group(0)
|
||||
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=f"Alice's email is {token}"))])
|
||||
return super()._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
||||
|
||||
@ -568,7 +590,7 @@ async def test_async_graph_redacts_configured_title_model_input(monkeypatch, ena
|
||||
from deerflow.config.title_config import TitleConfig
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
|
||||
pii = PiiRedactionConfig(enabled=enabled)
|
||||
pii = PiiRedactionConfig(enabled=enabled, token_secret=_TOKEN_SECRET)
|
||||
config = _wiring_app_config(pii_redaction=pii, title=TitleConfig(enabled=True, model_name="title-model"))
|
||||
primary = _RecordingPiiModel(responses=[AIMessage(content="Reply to charlie@example.net")])
|
||||
title = Mock(ainvoke=AsyncMock(return_value=AIMessage(content="Contact records")))
|
||||
@ -585,7 +607,7 @@ async def test_async_graph_redacts_configured_title_model_input(monkeypatch, ena
|
||||
assert "alice@example.com" not in primary.seen[0]
|
||||
assert "alice@example.com" not in prompt
|
||||
assert "charlie@example.net" not in prompt
|
||||
assert "[EMAIL_1]" in prompt and "[EMAIL_2]" in prompt
|
||||
assert EMAIL_ALICE in prompt and EMAIL_CHARLIE in prompt
|
||||
else:
|
||||
assert "alice@example.com" in primary.seen[0]
|
||||
assert "alice@example.com" in prompt and "charlie@example.net" in prompt
|
||||
@ -599,7 +621,7 @@ def test_compiled_graph_keeps_summary_and_retained_pii_distinct(async_mode):
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware
|
||||
from deerflow.agents.thread_state import ThreadState
|
||||
|
||||
pii = PiiRedactionConfig(enabled=True)
|
||||
pii = PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)
|
||||
config = _wiring_app_config(pii_redaction=pii)
|
||||
summary = _RecordingPiiModel(responses=[AIMessage(content="unused")], echo_summary=True)
|
||||
primary = _RecordingPiiModel(responses=[AIMessage(content="done")])
|
||||
@ -618,49 +640,107 @@ def test_compiled_graph_keeps_summary_and_retained_pii_distinct(async_mode):
|
||||
result = asyncio.run(graph.ainvoke(state)) if async_mode else graph.invoke(state)
|
||||
|
||||
assert summary.seen and "alice@example.com" not in summary.seen[0]
|
||||
assert result["summary_text"] == "Alice's email is [EMAIL_1]"
|
||||
assert "Alice's email is [EMAIL_1]" in primary.seen[0]
|
||||
assert "Bob's email is [EMAIL_2]" in primary.seen[0]
|
||||
assert result["summary_text"] == f"Alice's email is {EMAIL_ALICE}"
|
||||
assert f"Alice's email is {EMAIL_ALICE}" in primary.seen[0]
|
||||
assert f"Bob's email is {EMAIL_BOB_COM}" in primary.seen[0]
|
||||
assert "bob@example.com" not in primary.seen[0]
|
||||
assert any(message.content == messages[2].content for message in result["messages"])
|
||||
|
||||
|
||||
def test_summary_redaction_reserves_existing_placeholders():
|
||||
config = PiiRedactionConfig(enabled=True)
|
||||
assert redact_text("Alice [EMAIL_1], Bob bob@example.com", config) == "Alice [EMAIL_1], Bob [EMAIL_2]"
|
||||
def test_existing_placeholder_text_passes_through_and_new_value_gets_token():
|
||||
config = PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)
|
||||
# Placeholder text is not a detector match, so pre-existing tokens pass
|
||||
# through unchanged while the new raw value gets its own token.
|
||||
assert redact_text("Alice [EMAIL_1], Bob bob@example.com", config) == f"Alice [EMAIL_1], Bob {EMAIL_BOB_COM}"
|
||||
|
||||
|
||||
def test_existing_placeholder_in_later_content_block_is_reserved_first():
|
||||
def test_existing_placeholder_in_later_content_block_keeps_identity():
|
||||
messages, _ = _run_model_call(_make_middleware(), [HumanMessage(content=["bob@example.com", {"type": "text", "text": "Alice [EMAIL_1]"}])])
|
||||
assert messages[0].content == ["[EMAIL_2]", {"type": "text", "text": "Alice [EMAIL_1]"}]
|
||||
assert messages[0].content == [EMAIL_BOB_COM, {"type": "text", "text": "Alice [EMAIL_1]"}]
|
||||
|
||||
|
||||
def test_raw_legacy_summary_and_retained_messages_share_request_allocation():
|
||||
from deerflow.agents.middlewares.durable_context_middleware import DurableContextMiddleware
|
||||
|
||||
pii = PiiRedactionConfig(enabled=True)
|
||||
pii = PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET)
|
||||
state = {"summary_text": "Alice alice@example.com"}
|
||||
request = _StateRequest(state, [HumanMessage(content="Alice alice@example.com; Bob bob@example.com")])
|
||||
redacted = PiiRedactionMiddleware(pii)._process_request(request)
|
||||
final = DurableContextMiddleware(pii_redaction_config=pii)._inject(redacted)
|
||||
text = get_buffer_string(final.messages)
|
||||
assert "Alice [EMAIL_1]" in text and "Bob [EMAIL_2]" in text
|
||||
assert f"Alice {EMAIL_ALICE}" in text and f"Bob {EMAIL_BOB_COM}" in text
|
||||
assert "alice@example.com" not in text and "bob@example.com" not in text
|
||||
assert state == {"summary_text": "Alice alice@example.com"}
|
||||
assert request.messages[0].content == "Alice alice@example.com; Bob bob@example.com"
|
||||
|
||||
|
||||
def test_repeated_compaction_reserves_prior_summary_tokens():
|
||||
middleware = TestSummarizationCompactionInput()._middleware(PiiRedactionConfig(enabled=True))
|
||||
def test_repeated_compaction_keeps_prior_summary_tokens():
|
||||
middleware = TestSummarizationCompactionInput()._middleware(PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
prompt = middleware._build_summary_prompt([HumanMessage("Carol carol@example.com")], previous_summary="Alice [EMAIL_1], Bob [EMAIL_2]")
|
||||
assert "Alice [EMAIL_1], Bob [EMAIL_2]" in prompt
|
||||
assert "Carol [EMAIL_3]" in prompt
|
||||
assert f"Carol {EMAIL_CAROL}" in prompt
|
||||
|
||||
|
||||
def test_title_redacts_identifiers_before_field_truncation():
|
||||
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
|
||||
|
||||
config = _wiring_app_config(pii_redaction=PiiRedactionConfig(enabled=True))
|
||||
config = _wiring_app_config(pii_redaction=PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
prompt, fallback = TitleMiddleware(app_config=config)._build_title_prompt({"messages": [HumanMessage(content="x " * 246 + "alice@example.com"), AIMessage(content="done")]})
|
||||
assert "alice" not in prompt
|
||||
assert fallback.endswith("alice@example.com") # Local display fallback preserves the original user text.
|
||||
|
||||
|
||||
def test_minted_tokens_survive_later_detectors():
|
||||
# Review round 7 on #5577: hex digests carry digit runs that later
|
||||
# digit-anchored detectors re-scan, corrupting minted tokens into nested
|
||||
# placeholders (user280@example.com -> [EMAIL_…[PHONE_…]…]). Base-26
|
||||
# letters contain no digits, so the token survives the full pinned order.
|
||||
result = redact_text("contact user280@example.com today", PiiRedactionConfig(enabled=True, token_secret=_TOKEN_SECRET))
|
||||
assert result.startswith("contact [EMAIL_") and result.endswith("] today")
|
||||
assert result.count("[") == 1 and result.count("]") == 1
|
||||
assert "@" not in result and "user280" not in result
|
||||
|
||||
|
||||
def test_token_secret_scopes_linkability():
|
||||
# Review round 10: unkeyed digests are publicly computable fingerprints,
|
||||
# linkable across deployments. A deployment secret scopes the tokens.
|
||||
cfg_a = PiiRedactionConfig(enabled=True, token_secret="deployment-a-secret-value")
|
||||
cfg_b = PiiRedactionConfig(enabled=True, token_secret="deployment-b-secret-value")
|
||||
assert redact_text("alice@example.com", cfg_a) != redact_text("alice@example.com", cfg_b)
|
||||
assert redact_text("alice@example.com", cfg_a) == redact_text("alice@example.com", cfg_a)
|
||||
|
||||
|
||||
def test_hmac_tokens_still_letters_only():
|
||||
cfg = PiiRedactionConfig(enabled=True, token_secret="s3cret-with-enough-length")
|
||||
result = redact_text("contact user280@example.com today", cfg)
|
||||
assert result.count("[") == 1 and result.count("]") == 1
|
||||
assert "@" not in result and "user280" not in result
|
||||
|
||||
|
||||
def test_enabled_config_rejects_missing_or_weak_token_secret():
|
||||
# Review round 11 on #5577: with token_secret optional, the default
|
||||
# enabled configuration minted empty-key HMAC digests — publicly
|
||||
# computable, globally linkable fingerprints of the raw values. Enabling
|
||||
# redaction now requires a usable deployment key, with the unkeyed mode
|
||||
# gone entirely.
|
||||
with pytest.raises(ValidationError):
|
||||
PiiRedactionConfig(enabled=True)
|
||||
with pytest.raises(ValidationError):
|
||||
PiiRedactionConfig(enabled=True, token_secret=" ")
|
||||
with pytest.raises(ValidationError):
|
||||
PiiRedactionConfig(enabled=True, token_secret="short")
|
||||
|
||||
|
||||
def test_disabled_config_still_defaults_off_without_secret():
|
||||
config = PiiRedactionConfig(enabled=False)
|
||||
assert config.token_secret is None
|
||||
assert redact_text("alice@example.com", config) == "alice@example.com"
|
||||
|
||||
|
||||
def test_user_message_boundary_mints_the_keyed_token():
|
||||
# The request-scoped redactor must carry the configured deployment key
|
||||
# like every other seam: _process_request built its redactor without the
|
||||
# token key, so user-message placeholders silently diverged from
|
||||
# redact_text / tool-result / memory-queue placeholders for the same value.
|
||||
messages, _ = _run_model_call(_make_middleware(), [HumanMessage("Contact alice@example.com")])
|
||||
assert messages[0].content == f"Contact {EMAIL_ALICE}"
|
||||
|
||||
@ -700,7 +700,9 @@ def test_factory_attaches_memory_flush_hook_by_default(monkeypatch):
|
||||
middleware = create_summarization_middleware(app_config=app_config)
|
||||
|
||||
assert middleware is not None
|
||||
assert memory_flush_hook in middleware._before_summarization_hooks
|
||||
# The hook is wrapped in functools.partial to carry the pii_redaction
|
||||
# config; unwrap it for the identity check.
|
||||
assert any(getattr(h, "func", h) is memory_flush_hook for h in middleware._before_summarization_hooks)
|
||||
|
||||
|
||||
def test_factory_skip_memory_flush_omits_hook(monkeypatch):
|
||||
|
||||
@ -1411,11 +1411,13 @@ loop_detection:
|
||||
# phone numbers, API keys, credit-card numbers (Luhn-validated) and national
|
||||
# IDs found in genuine user messages and remote-content tool results
|
||||
# (web_fetch / web_search / image_search / web_capture and MCP tools) are
|
||||
# rewritten to irreversible placeholders like [EMAIL_1] before they reach the
|
||||
# model. No mapping is stored, so redaction cannot be reversed. Raw text stays
|
||||
# in thread state for the UI; the memory-extraction path is not covered.
|
||||
# Compaction and configured LLM title inputs are redacted too. Existing
|
||||
# placeholders reserve indices; compacted raw identities cannot be relinked.
|
||||
# rewritten to keyed, value-derived placeholders like
|
||||
# [EMAIL_bdqkxwcfmqvrqzblcchjkqbcdqw] before they reach the model. Tokens are
|
||||
# HMAC digests keyed by token_secret, so they cannot be re-derived offline.
|
||||
# Raw text stays in thread state for the UI; queued memory-extraction
|
||||
# payloads are redacted too. Compaction and configured LLM title inputs are
|
||||
# redacted as well; placeholders are value-derived, so an identity redacted
|
||||
# before compaction renders the same token afterwards.
|
||||
pii_redaction:
|
||||
enabled: false
|
||||
# redact_email: true
|
||||
@ -1423,6 +1425,13 @@ pii_redaction:
|
||||
# redact_credit_card: true
|
||||
# redact_phone: true
|
||||
# redact_national_id: true
|
||||
# Required when enabled: true (non-empty, at least 16 characters) — the
|
||||
# deployment-scoped HMAC secret for the placeholder digests. Tokens are
|
||||
# linkable only within this deployment and cannot be re-derived offline
|
||||
# without it. Enabling redaction without a valid secret fails config
|
||||
# validation: an unkeyed digest would be a publicly computable, globally
|
||||
# linkable fingerprint of the raw values (guessable for phone numbers).
|
||||
# token_secret: change-me-to-a-long-random-string
|
||||
|
||||
# ============================================================================
|
||||
# Tool Progress State Machine Configuration (RFC #3177)
|
||||
|
||||
@ -45,7 +45,7 @@ fi
|
||||
|
||||
cd "$REPO_ROOT/backend" && CONFIG_WIN_PATH="$CONFIG_WIN" EXAMPLE_WIN_PATH="$EXAMPLE_WIN" uv run python -c "
|
||||
import os
|
||||
import sys, shutil, copy, re
|
||||
import sys, shutil, copy, re, secrets
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
@ -156,6 +156,30 @@ MIGRATIONS = {
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def migrate_pii_token_secret(data):
|
||||
# token_secret became mandatory whenever pii_redaction is enabled (v47).
|
||||
# A v46 deployment that enabled redaction without a secret would fail
|
||||
# startup after the upgrade, so generate a random deployment-scoped
|
||||
# secret and persist it here; the .bak backup taken below covers the
|
||||
# original file.
|
||||
pii = data.get('pii_redaction')
|
||||
changes = []
|
||||
if not isinstance(pii, dict) or not pii.get('enabled'):
|
||||
return changes
|
||||
secret = pii.get('token_secret')
|
||||
if isinstance(secret, str) and secret.strip():
|
||||
return changes
|
||||
pii['token_secret'] = secrets.token_urlsafe(32)
|
||||
changes.append('pii_redaction.token_secret generated (required for enabled redaction; a random value was persisted to config.yaml)')
|
||||
return changes
|
||||
|
||||
|
||||
MIGRATIONS[47] = {
|
||||
'description': 'Generate a token_secret for deployments with pii_redaction enabled (now mandatory)',
|
||||
'data_transform': migrate_pii_token_secret,
|
||||
}
|
||||
|
||||
# Apply migrations in order for versions (user_version, example_version]
|
||||
migrated = []
|
||||
for version in range(user_version + 1, example_version + 1):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user