agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/task/teardown.py

@@ -0,0 +1,142 @@
+"""Tear down the sandboxes a finished run deployed.
+
+``teardown_run`` terminates every sandbox a run's context records: its bare sandboxes', its envs' and its
+agents'. It removes compute only; the instance, its outputs and the artifacts it collected stay. A local
+sandbox's work folder goes too, once the sandbox is down. ``agent-env run`` and ``agent-env eval run`` call
+it after each run. ``Task.run`` never does, so a caller that keeps a run's environment up, as the hub and
+the worker do, owns its teardown.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass
+from typing import NamedTuple, Optional
+
+from agent_env.env.env import DeployedEnv
+from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, remove_local_work_dir
+from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider
+from agent_env.task.interrupts import tearing_down
+from agent_env.task_step.context import TaskStepContext
+from agent_env.task_step.task_steps.teardown_sandboxes import _DEFAULT_PROVIDERS, TORN_DOWN_KEY, _env_sandbox_ids
+
+logger = logging.getLogger(__name__)
+
+TERMINATE_TIMEOUT_SECONDS = 120
+
+
+class RecordedSandbox(NamedTuple):
+    sandbox_id: str
+    sandbox_type: Optional[str]  # the backend that made it; None leaves it to the slot's configured provider
+    slot: str  # "sandbox", "env" or "agent"
+
+
+@dataclass(frozen=True)
+class TeardownReport:
+    terminated: tuple[RecordedSandbox, ...] = ()
+    failed: tuple[tuple[RecordedSandbox, str], ...] = ()  # each with why
+    left: tuple[RecordedSandbox, ...] = ()  # not reached, because the teardown was stopped
+
+    @property
+    def still_up(self) -> tuple[RecordedSandbox, ...]:
+        """What may still be running: what failed to terminate and what the teardown didn't reach."""
+        return (*(sandbox for sandbox, _ in self.failed), *self.left)
+
+    @classmethod
+    def skipped(cls, context: TaskStepContext) -> TeardownReport:
+        """A teardown that didn't start: everything ``context`` records that's still up is left."""
+        return cls(left=_pending(context))
+
+
+def recorded_sandboxes(context: TaskStepContext) -> tuple[RecordedSandbox, ...]:
+    """Every sandbox ``context`` records, once each. An agent placed on another record's sandbox shares its
+    id, so the sandbox is attributed to whichever made it."""
+    found: dict[str, RecordedSandbox] = {}
+    for sandbox in context.deployed_sandboxes:
+        found.setdefault(sandbox.sandbox_id, RecordedSandbox(sandbox.sandbox_id, sandbox.sandbox_type, "sandbox"))
+    for env in context.deployed_envs:
+        for sandbox_id, sandbox_type in _env_sandbox_ids(env):
+            found.setdefault(sandbox_id, RecordedSandbox(sandbox_id, sandbox_type, "env"))
+    for agent in context.deployed_agents:
+        if agent.sandbox_id:
+            found.setdefault(agent.sandbox_id, RecordedSandbox(agent.sandbox_id, agent.sandbox_type, "agent"))
+    return tuple(found.values())
+
+
+def sandbox_ids(record) -> list[str]:
+    """The ids of the sandboxes behind one deployed env, agent or bare sandbox record."""
+    if isinstance(record, DeployedEnv):
+        return [sandbox_id for sandbox_id, _ in _env_sandbox_ids(record)]
+    return [record.sandbox_id] if record.sandbox_id else []
+
+
+async def teardown_run(context: TaskStepContext, *, timeout: float = TERMINATE_TIMEOUT_SECONDS) -> TeardownReport:
+    """Terminate every sandbox ``context`` records, all at once, each within ``timeout`` seconds, and remove a
+    local one's work folder once it's down.
+
+    Best-effort: a failure is logged and reported, never raised. A sandbox already torn down, by a
+    ``teardown_sandboxes`` step or an earlier call, isn't terminated again, though a local one's folder
+    still goes. Under ``Interrupts.gather`` the first signal spares it. Cancelling it stops it early: it
+    returns what it reached, with the rest in ``left``, rather than raising."""
+    tearing_down()
+    pending = _pending(context)
+    terminated, failed = [], []
+
+    async def take_down(sandbox: RecordedSandbox) -> None:
+        if sandbox not in pending:  # torn down already: only a local one's folder can be left
+            if why := await _remove_folder(sandbox):
+                failed.append((sandbox, why))
+            return
+        try:
+            async with asyncio.timeout(timeout):
+                handle = await _terminate(sandbox)
+        except Exception as e:
+            reason = f"timed out after {timeout:g}s" if isinstance(e, TimeoutError) else _described(e)
+            logger.warning("teardown: %s (%s) was not terminated: %s", sandbox.sandbox_id, kind(sandbox), reason)
+            failed.append((sandbox, reason))
+            return
+        context.metadata[TORN_DOWN_KEY] = [*context.metadata.get(TORN_DOWN_KEY, []), sandbox.sandbox_id]
+        if isinstance(handle, LocalSandbox) and (why := await _remove_folder(sandbox)):
+            failed.append((sandbox, why))
+        else:
+            terminated.append(sandbox)
+
+    try:
+        await asyncio.gather(*(take_down(sandbox) for sandbox in recorded_sandboxes(context)))
+    except asyncio.CancelledError:
+        pass
+    reached = {sandbox.sandbox_id for sandbox in terminated} | {sandbox.sandbox_id for sandbox, _ in failed}
+    return TeardownReport(tuple(terminated), tuple(failed), tuple(s for s in pending if s.sandbox_id not in reached))
+
+
+def _pending(context: TaskStepContext) -> tuple[RecordedSandbox, ...]:
+    torn_down = set(context.metadata.get(TORN_DOWN_KEY) or [])
+    return tuple(sandbox for sandbox in recorded_sandboxes(context) if sandbox.sandbox_id not in torn_down)
+
+
+async def _remove_folder(sandbox: RecordedSandbox) -> str | None:
+    """Remove a local sandbox's work folder, and say why when it can't be."""
+    try:
+        await asyncio.to_thread(remove_local_work_dir, sandbox.sandbox_id)
+    except OSError as e:
+        logger.warning("teardown: %s's work folder was not removed: %s", sandbox.sandbox_id, e)
+        return f"terminated, but its work folder wasn't removed: {e}"
+    return None
+
+
+async def _terminate(sandbox: RecordedSandbox):
+    provider = (build_sandbox_provider(sandbox.sandbox_type) if sandbox.sandbox_type
+                else _DEFAULT_PROVIDERS[sandbox.slot]())
+    handle = await provider.get_sandbox(sandbox.sandbox_id)
+    await handle.terminate()
+    return handle
+
+
+def _described(error: Exception) -> str:
+    return f"{type(error).__name__}: {error}" if str(error) else type(error).__name__
+
+
+def kind(sandbox: RecordedSandbox) -> str:
+    """The backend a sandbox is on, as the CLI prints it."""
+    return sandbox.sandbox_type or f"the configured {sandbox.slot} provider"

agent_env/task_step/__init__.py

@@ -0,0 +1,47 @@
+"""TaskStep module for AgentEnv."""
+
+from agent_env.env.env import DeployedEnv
+from .context import DeployedAgent, PromptResponse, TaskStepContext
+from .task_step import TaskStep
+from .task_steps import AddSkillsTaskStep, AgentPromptResponseVerifierTaskStep, BuildMcpCliTaskStep, CombineUniverseVerdictsStep, DeployAgentTaskStep, DeployEnvTaskStep, EnvOutcomeVerifierTaskStep, LoadArtifactTaskStep, ModifyEnvToolAccessStep, PromptAgentTaskStep, RubricsVerifierTaskStep, ScoreAggregator, ValidationGateAggregatorStep, VerifyCoreA2AProtocolStep, VerifyCoreEnvironmentProtocolStep, VerifyEnvironmentCardStep, VerifyMCPEnvAssessmentStep, VerifyMCPToolSchemaTaskStep, VerifySpecConformanceTaskStep, VerifyUniverseLoadExportRoundtripStep, VerifyA2AAgentCardStep, VerifyA2AAgentMCPStep, VerifyA2ASkillConfigStep, VerifyA2ATrajectoryStep
+from .task_steps.verifiers.judge_utils.judge_output_format import JudgeOutputFormat, get_judge_output_format_spec
+from .store import TaskStepStore, TaskStepQuery, get_task_step_store, set_task_step_store, reset_task_step_store
+
+__all__ = [
+    "DeployedAgent",
+    "DeployedEnv",
+    "JudgeOutputFormat",
+    "get_judge_output_format_spec",
+    "PromptResponse",
+    "ScoreAggregator",
+    "TaskStepContext",
+    "TaskStep",
+    "AddSkillsTaskStep",
+    "AgentPromptResponseVerifierTaskStep",
+    "BuildMcpCliTaskStep",
+    "DeployAgentTaskStep",
+    "DeployEnvTaskStep",
+    "EnvOutcomeVerifierTaskStep",
+    "LoadArtifactTaskStep",
+    "ModifyEnvToolAccessStep",
+    "PromptAgentTaskStep",
+    "RubricsVerifierTaskStep",
+    "VerifyMCPEnvAssessmentStep",
+    "VerifyMCPToolSchemaTaskStep",
+    "VerifySpecConformanceTaskStep",
+    "ValidationGateAggregatorStep",
+    "VerifyUniverseLoadExportRoundtripStep",
+    "CombineUniverseVerdictsStep",
+    "VerifyEnvironmentCardStep",
+    "VerifyCoreEnvironmentProtocolStep",
+    "VerifyA2AAgentCardStep",
+    "VerifyA2AAgentMCPStep",
+    "VerifyCoreA2AProtocolStep",
+    "VerifyA2ASkillConfigStep",
+    "VerifyA2ATrajectoryStep",
+    "TaskStepStore",
+    "TaskStepQuery",
+    "get_task_step_store",
+    "set_task_step_store",
+    "reset_task_step_store",
+]

agent_env/task_step/context.py

@@ -0,0 +1,195 @@
+"""Context dataclass for task step execution."""
+
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass, field
+from typing import Any
+
+from agent_env.env.env import DeployedEnv
+
+_REDACTED_KEYS = {
+    "litellm_api_key", "judge_litellm_api_key", "usersim_api_key", "remote_tokens", "cf_access_client_secret",
+}
+
+
+def _holds_redacted_key(value: Any) -> bool:
+    """Whether ``value`` carries a ``_REDACTED_KEYS`` entry at any depth."""
+    return isinstance(value, dict) and any(
+        k in _REDACTED_KEYS or _holds_redacted_key(v) for k, v in value.items()
+    )
+
+
+def regraft_redacted_keys(source: dict, target: dict) -> None:
+    """Copy every ``_REDACTED_KEYS`` value in ``source`` into ``target`` in place; the inverse
+    of ``_strip_redacted_keys`` for a context rebuilt from a stored document."""
+    for k, v in source.items():
+        if k in _REDACTED_KEYS:
+            target[k] = v
+        elif isinstance(v, dict):
+            child = target.get(k)
+            if not isinstance(child, dict):
+                # A secrets-only parent leaves nothing to rebuild from, so recreate it —
+                # but never over a value the rebuild did write.
+                if target.get(k) is not None or not _holds_redacted_key(v):
+                    continue
+                child = target[k] = {}
+            regraft_redacted_keys(v, child)
+
+
+def _strip_redacted_keys(value: Any) -> Any:
+    """Recursively strip `_REDACTED_KEYS` from any nested dicts in `value`.
+
+    Single source of truth used by both `to_safe_dict` (full-blob writes via
+    `record_task_failure`) and the path-level diff in `context_ops` so secrets
+    never reach Mongo regardless of which code path persists them.
+    """
+    if isinstance(value, dict):
+        return {k: _strip_redacted_keys(v) for k, v in value.items() if k not in _REDACTED_KEYS}
+    if isinstance(value, list):
+        return [_strip_redacted_keys(v) for v in value]
+    return value
+
+
+@dataclass
+class DeployedAgent:
+    agent_name: str
+    api_url: str
+    sandbox_id: str | None = None
+    sandbox_type: str | None = None
+    a2a_url: str | None = None
+    a2a_card: dict | None = None
+    instance_id: str | None = None
+    role: str | None = None
+    network_policy: dict | None = None
+    # Installed straight onto the sandbox host (install_agent host mode), not into a container.
+    on_host: bool = False
+
+    @classmethod
+    def from_dict(cls, data: dict) -> DeployedAgent:
+        return cls(
+            agent_name=data["agent_name"],
+            api_url=data["api_url"],
+            sandbox_id=data.get("sandbox_id"),
+            sandbox_type=data.get("sandbox_type"),
+            a2a_url=data.get("a2a_url"),
+            a2a_card=data.get("a2a_card"),
+            instance_id=data.get("instance_id"),
+            role=data.get("role"),
+            network_policy=data.get("network_policy"),
+            on_host=bool(data.get("on_host", False)),
+        )
+
+
+@dataclass
+class DeployedSandbox:
+    sandbox_name: str
+    sandbox_id: str
+    sandbox_mode: str
+    sandbox_type: str | None = None
+    tunnel_urls: dict[str, str] | None = None
+    vnc_url: str | None = None
+    instance_id: str | None = None
+    created_at_utc: str | None = None
+    expires_at_utc: str | None = None
+    network_policy: dict | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict) -> DeployedSandbox:
+        return cls(
+            sandbox_name=data["sandbox_name"],
+            sandbox_id=data["sandbox_id"],
+            sandbox_mode=data["sandbox_mode"],
+            sandbox_type=data.get("sandbox_type"),
+            tunnel_urls=data.get("tunnel_urls"),
+            vnc_url=data.get("vnc_url"),
+            instance_id=data.get("instance_id"),
+            created_at_utc=data.get("created_at_utc"),
+            expires_at_utc=data.get("expires_at_utc"),
+            network_policy=data.get("network_policy"),
+        )
+
+
+@dataclass
+class PromptResponse:
+    prompt_id: str
+    response: str
+    prompt_text: str | None = None
+    agent_trajectory_s3_uri: str | None = None
+    agent_trajectory_s3_prefix: str | None = None
+    agent_trajectory_file_path: str | None = None
+    target_agent_per_turn_trajectory_s3_uris: list[str | None] | None = None
+    # A None entry means "identical to prompt_text" — the first turn of a
+    # prompt-mode step is not stored twice.
+    source_agent_per_turn_prompt_parts: list[list[dict] | None] | None = None
+    compact_trajectory_s3_uri: str | None = None
+    tool_call_count: int | None = None
+    model: str | None = None
+    error_type: str | None = None
+    error_code: str | None = None
+    error_message: str | None = None
+    agent_session_id: str = ""
+    a2a_context_id: str | None = None
+    agent_name: str | None = None
+    step_id: str | None = None
+    structured_output: Any = None
+
+    @classmethod
+    def from_dict(cls, data: dict) -> PromptResponse:
+        return cls(
+            prompt_id=data["prompt_id"],
+            response=data["response"],
+            prompt_text=data.get("prompt_text"),
+            agent_trajectory_s3_uri=data.get("agent_trajectory_s3_uri"),
+            agent_trajectory_s3_prefix=data.get("agent_trajectory_s3_prefix"),
+            agent_trajectory_file_path=data.get("agent_trajectory_file_path"),
+            target_agent_per_turn_trajectory_s3_uris=data.get("target_agent_per_turn_trajectory_s3_uris"),
+            source_agent_per_turn_prompt_parts=data.get("source_agent_per_turn_prompt_parts"),
+            compact_trajectory_s3_uri=data.get("compact_trajectory_s3_uri"),
+            tool_call_count=data.get("tool_call_count"),
+            model=data.get("model"),
+            error_type=data.get("error_type"),
+            error_code=data.get("error_code"),
+            error_message=data.get("error_message"),
+            agent_session_id=data.get("agent_session_id", ""),
+            a2a_context_id=data.get("a2a_context_id"),
+            agent_name=data.get("agent_name"),
+            step_id=data.get("step_id"),
+            structured_output=data.get("structured_output"),
+        )
+
+
+@dataclass
+class TaskStepContext:
+    deployed_envs: list[DeployedEnv] = field(default_factory=list)
+    deployed_agents: list[DeployedAgent] = field(default_factory=list)
+    deployed_sandboxes: list[DeployedSandbox] = field(default_factory=list)
+    prompt_responses: list[PromptResponse] = field(default_factory=list)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    agent_model: str | None = None
+    default_agent_model: str | None = None
+    agent_artifact_id: str | None = None
+    agent_harness: str | None = None
+    instance_id: str | None = None
+
+    def to_safe_dict(self) -> dict[str, Any]:
+        """Return a dict representation with sensitive keys recursively removed."""
+        d = dataclasses.asdict(self)
+        if "metadata" in d:
+            d["metadata"] = _strip_redacted_keys(d["metadata"])
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> TaskStepContext:
+        return cls(
+            deployed_envs=[DeployedEnv.from_dict(e) for e in data.get("deployed_envs", [])],
+            deployed_agents=[DeployedAgent.from_dict(a) for a in data.get("deployed_agents", [])],
+            deployed_sandboxes=[DeployedSandbox.from_dict(s) for s in data.get("deployed_sandboxes", [])],
+            prompt_responses=[PromptResponse.from_dict(p) for p in data.get("prompt_responses", [])],
+            metadata=data.get("metadata", {}),
+            agent_model=data.get("agent_model"),
+            default_agent_model=data.get("default_agent_model"),
+            agent_artifact_id=data.get("agent_artifact_id"),
+            agent_harness=data.get("agent_harness"),
+            instance_id=data.get("instance_id"),
+        )

agent_env/task_step/context_ops.py

@@ -0,0 +1,224 @@
+"""Build MongoDB update specs from TaskStepContext diffs.
+
+Replaces the read-merge-write CAS pattern with path-level Mongo ops. A single
+document's disjoint-path writes are atomic in Mongo, so concurrent workers no
+longer need application-layer merging or a retry loop.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import types
+from dataclasses import dataclass, field
+from typing import Any, Union, get_args, get_origin, get_type_hints
+
+from agent_env.task_step.context import TaskStepContext, _REDACTED_KEYS, _strip_redacted_keys
+
+
+@dataclass
+class ContextUpdateOps:
+    """A backend-agnostic update spec describing a single step's changes to TaskInstance.context.
+
+    Dot-paths are rooted at the TaskInstance document (they include the
+    ``context.`` prefix). Compile to a backend-agnostic ``UpdateSpec`` via `to_update_spec`.
+
+    Lists are set-semantic: when a step's "post" snapshot reads the shared in-memory
+    context it sees concurrent siblings' appends too, so the diff's tail may include
+    their items and the write must dedupe (``$addToSet`` via ``to_update_spec``, or
+    the Python ``_union`` the completion CAS in ``task/store.py`` applies). **Constraint:**
+    items in any list must be uniquely identifiable by value, or duplicates are silently
+    dropped on write; ``deployed_envs`` are told apart by ``instance_id`` (``step_journal._item_key``).
+    """
+
+    sets: dict[str, Any] = field(default_factory=dict)
+    unsets: set[str] = field(default_factory=set)
+    add_to_sets: dict[str, list[Any]] = field(default_factory=dict)
+
+    def is_empty(self) -> bool:
+        return not (self.sets or self.unsets or self.add_to_sets)
+
+    def to_update_spec(self) -> "UpdateSpec":
+        """Convert to a backend-agnostic UpdateSpec (1:1 field mapping)."""
+        from agent_env.store.document_store import UpdateSpec
+
+        return UpdateSpec(
+            set=dict(self.sets),
+            unset=set(self.unsets),
+            add_to_set={p: list(v) for p, v in self.add_to_sets.items()},
+        )
+
+    def to_journal_dict(self) -> dict:
+        """Mongo-safe encoding: dotted paths as values, never as keys."""
+        return {
+            "sets": [{"path": p, "value": v} for p, v in self.sets.items()],
+            "unsets": sorted(self.unsets),
+            "add_to_sets": [{"path": p, "items": list(v)} for p, v in self.add_to_sets.items()],
+        }
+
+    @classmethod
+    def from_journal_dict(cls, d: dict) -> "ContextUpdateOps":
+        return cls(
+            sets={e["path"]: e["value"] for e in d.get("sets", [])},
+            unsets=set(d.get("unsets", [])),
+            add_to_sets={e["path"]: list(e["items"]) for e in d.get("add_to_sets", [])},
+        )
+
+
+@dataclass(frozen=True)
+class _FieldClassification:
+    scalars: tuple[str, ...]
+    lists: tuple[str, ...]
+    dicts: tuple[str, ...]
+
+
+def _classify_context_fields() -> _FieldClassification:
+    """Partition TaskStepContext fields into scalar/list/dict via type introspection.
+
+    Lists are assumed to be set-semantic (emitted via `$addToSet`). Dicts are
+    walked recursively for leaf-level path diffs. If a future field needs
+    different semantics, classification must be made explicit here.
+    """
+    scalars: list[str] = []
+    lists: list[str] = []
+    dicts: list[str] = []
+    hints = get_type_hints(TaskStepContext)
+    for f in dataclasses.fields(TaskStepContext):
+        kind = _origin_kind(hints[f.name])
+        if kind is list:
+            lists.append(f.name)
+        elif kind is dict:
+            dicts.append(f.name)
+        else:
+            scalars.append(f.name)
+    return _FieldClassification(tuple(scalars), tuple(lists), tuple(dicts))
+
+
+def _origin_kind(tp: Any) -> Any:
+    """Resolve ``tp`` to list/dict/other, stripping Optional/Union where needed."""
+    origin = get_origin(tp)
+    if origin is Union or origin is types.UnionType:
+        non_none = [a for a in get_args(tp) if a is not type(None)]
+        if len(non_none) == 1:
+            return _origin_kind(non_none[0])
+        return None
+    return origin
+
+
+_FIELDS = _classify_context_fields()
+
+
+def build_context_update_ops(
+    pre: TaskStepContext | None,
+    post: TaskStepContext,
+) -> ContextUpdateOps:
+    """Diff two context snapshots into a path-level Mongo update spec.
+
+    - Scalar fields: `$set` on change, `$unset` on value→None.
+    - List fields (set-semantic): new items (by equality) go to `$addToSet`.
+    - Dict fields (e.g. `metadata`): walked recursively so only leaf-level
+      differences emit. Concurrent writers to sibling leaves of the same
+      parent don't overwrite each other. Lists inside dicts use `$addToSet`
+      for prefix-extending writes (so concurrent appenders don't double-add);
+      non-prefix mutations (reorder, pop, replacement) emit wholesale `$set`.
+    - Redaction (`_REDACTED_KEYS`) applied at every emission point.
+
+    Dict keys containing ``.`` or starting with ``$`` fall back to a
+    wholesale ``$set`` of the nearest enclosing subtree (see ``_diff_dict``).
+    """
+    ops = ContextUpdateOps()
+    pre_ctx = pre if pre is not None else TaskStepContext()
+
+    for name in _FIELDS.scalars:
+        pre_v = getattr(pre_ctx, name)
+        post_v = getattr(post, name)
+        if pre_v == post_v:
+            continue
+        path = f"context.{name}"
+        if post_v is None:
+            ops.unsets.add(path)
+        else:
+            ops.sets[path] = post_v
+
+    for name in _FIELDS.lists:
+        pre_items = getattr(pre_ctx, name)
+        post_items = getattr(post, name)
+        new_items = [x for x in post_items if x not in pre_items]
+        if new_items:
+            ops.add_to_sets[f"context.{name}"] = [
+                dataclasses.asdict(x) if dataclasses.is_dataclass(x) else x
+                for x in new_items
+            ]
+
+    for name in _FIELDS.dicts:
+        _diff_dict(f"context.{name}", getattr(pre_ctx, name), getattr(post, name), ops)
+
+    return ops
+
+
+def _diff_dict(
+    path: str, pre: dict[str, Any], post: dict[str, Any], ops: ContextUpdateOps,
+) -> None:
+    """Recursively diff `pre` vs `post`, emitting path-level ops:
+
+    - Keys removed from pre → `$unset`.
+    - New dict-valued keys → recurse against `{}` so concurrent writers of
+      disjoint sub-keys emit distinct leaf paths (not the whole subtree).
+    - New list-valued keys → `$addToSet` (creates array if missing; dedupes
+      against what concurrent first-creators already wrote).
+    - New scalar keys → `$set`.
+    - Both dicts → recurse.
+    - Both lists, pre is prefix of post → `$addToSet` of the tail.
+    - Lists otherwise (reorder, pop, replacement) → wholesale `$set` (last write wins).
+    - Scalars differ → `$set`.
+
+    Redacted keys (`_REDACTED_KEYS`) are skipped at every level; redacted
+    descendants of emitted subtrees are stripped via `_strip_redacted_keys`.
+
+    If immediate keys on either side are unmappable to Mongo paths
+    (contain ``.`` or start with ``$``), emit a wholesale ``$set`` of the
+    whole dict at ``path`` instead of recursing.
+    """
+    if _has_unmappable_keys(pre) or _has_unmappable_keys(post):
+        ops.sets[path] = _strip_redacted_keys(post)
+        return
+
+    for k in pre.keys() - post.keys():
+        if k in _REDACTED_KEYS:
+            continue
+        ops.unsets.add(f"{path}.{k}")
+
+    for k, post_v in post.items():
+        if k in _REDACTED_KEYS:
+            continue
+        sub_path = f"{path}.{k}"
+        if k not in pre:
+            if isinstance(post_v, dict):
+                _diff_dict(sub_path, {}, post_v, ops)
+            elif isinstance(post_v, list):
+                if post_v:
+                    ops.add_to_sets[sub_path] = [_strip_redacted_keys(x) for x in post_v]
+            else:
+                ops.sets[sub_path] = _strip_redacted_keys(post_v)
+            continue
+        pre_v = pre[k]
+        if isinstance(pre_v, dict) and isinstance(post_v, dict):
+            _diff_dict(sub_path, pre_v, post_v, ops)
+        elif isinstance(pre_v, list) and isinstance(post_v, list):
+            if _is_prefix(pre_v, post_v):
+                # A re-append is a no-op on the write, so the journaled diff must be one too.
+                tail = [x for x in post_v[len(pre_v):] if x not in pre_v]
+                if tail:
+                    ops.add_to_sets[sub_path] = [_strip_redacted_keys(x) for x in tail]
+            elif pre_v != post_v:
+                ops.sets[sub_path] = _strip_redacted_keys(post_v)
+        elif pre_v != post_v:
+            ops.sets[sub_path] = _strip_redacted_keys(post_v)
+
+
+def _is_prefix(pre: list[Any], post: list[Any]) -> bool:
+    return len(post) >= len(pre) and post[: len(pre)] == pre
+
+
+def _has_unmappable_keys(d: dict[str, Any]) -> bool:
+    """True if any immediate key in `d` can't be a Mongo dot-notation path component."""
+    return any(not isinstance(k, str) or "." in k or k.startswith("$") for k in d)