Prevent duplicate Deep Agents messages after continue-as-new

CHANGELOG.md

@@ -52,6 +52,8 @@ to include examples, links to docs, or any other relevant information.
 
 ### Fixed
 
+- `temporalio.contrib.deepagents.run_deep_agent` no longer duplicates the original
+  input messages when carrying state through continue-as-new.
 - **Experimental**: External storage metrics now report the wall-clock time storage was in flight.
   Previously each batch's duration was summed, over-reporting the time whenever storage operations
   ran concurrently.

temporalio/contrib/deepagents/_activity.py

@@ -338,9 +338,9 @@ async def invoke_tool(self, input: ToolActivityInput) -> ToolActivityOutput:
     @_auto_heartbeater
     async def backend_op(self, input: BackendOpInput) -> BackendOpOutput:
         """Run one operation against a registered (real-I/O) backend."""
-        from temporalio.contrib.deepagents._tools import registered_backends
+        from temporalio.contrib.deepagents._tools import lookup_backend
 
-        backend = registered_backends().get(input.backend_ref)
+        backend = lookup_backend(input.backend_ref)
         if backend is None:
             raise ApplicationError(
                 f"Backend {input.backend_ref!r} is not registered on this worker.",

temporalio/contrib/deepagents/_tools.py

@@ -55,6 +55,11 @@
 
 _TOOL_REGISTRY: dict[str, "BaseTool"] = {}
 _BACKEND_REGISTRY: dict[str, Any] = {}
+# Recently retired backends (insertion-ordered, oldest evicted first). Keeps a
+# GC'd wrapper's backend reachable for the eviction -> activity-start window;
+# see _unregister_backend. Bounded so retired backends cannot accumulate.
+_RETIRED_BACKENDS: dict[str, Any] = {}
+_RETIRED_BACKENDS_MAX = 512
 # Serializes registration against the GC-time unregister in
 # _unregister_backend, which may run on another thread.
 _BACKEND_REGISTRY_LOCK = threading.Lock()
@@ -154,25 +159,49 @@ def register_backend(ref: str, backend: Any) -> None:
 
 
 def _unregister_backend(ref: str, inner: Any) -> None:
-    """Drop ``ref`` from the registry if it still maps to ``inner``.
+    """Retire ``ref`` from the registry if it still maps to ``inner``.
 
     GC hook for :class:`TemporalBackend` (via ``weakref.finalize``): a wrapper
     is typically constructed per workflow run, so without cleanup a long-lived
     worker accumulates one registry entry per run. The identity guard is
     load-bearing: refs are deterministic per run, so after a cache eviction a
     replay re-registers the *same* ref with a fresh inner backend — the evicted
     wrapper's finalizer must not remove that live registration.
+
+    Retired entries move to the bounded :data:`_RETIRED_BACKENDS` store instead
+    of vanishing: a ``backend_op`` activity scheduled just before a cache
+    eviction can be delivered *after* the evicted wrapper is collected, and the
+    replay that would re-register the ref only happens once that very activity
+    completes. :func:`lookup_backend` still resolves the ref in that window.
     """
     with _BACKEND_REGISTRY_LOCK:
         if _BACKEND_REGISTRY.get(ref) is inner:
             del _BACKEND_REGISTRY[ref]
+            _RETIRED_BACKENDS.pop(ref, None)
+            _RETIRED_BACKENDS[ref] = inner
+            while len(_RETIRED_BACKENDS) > _RETIRED_BACKENDS_MAX:
+                _RETIRED_BACKENDS.pop(next(iter(_RETIRED_BACKENDS)))
 
 
 def registered_backends() -> dict[str, Any]:
     """Return the live backend registry (read by the plugin at worker build)."""
     return _BACKEND_REGISTRY
 
 
+def lookup_backend(ref: str) -> Any | None:
+    """Resolve ``ref`` for a ``backend_op`` activity.
+
+    Prefers the live registry, then falls back to recently retired entries so
+    an activity dispatched before a cache eviction still resolves (see
+    :func:`_unregister_backend`).
+    """
+    with _BACKEND_REGISTRY_LOCK:
+        backend = _BACKEND_REGISTRY.get(ref)
+        if backend is not None:
+            return backend
+        return _RETIRED_BACKENDS.get(ref)
+
+
 # ---------------------------------------------------------------------------
 # activity_as_tool
 # ---------------------------------------------------------------------------

temporalio/contrib/deepagents/workflow.py

@@ -24,6 +24,7 @@
 
 # Reserved key under which the CAN result cache rides inside a state snapshot.
 _CACHE_KEY = "__temporal_cache__"
+_INPUT_CARRIED_KEY = "__temporal_input_in_transcript__"
 
 # Checkpointer classes that keep their state in the workflow's own memory and are
 # therefore rehydrated for free by deterministic replay. Anything else does its
@@ -158,7 +159,18 @@ async def call_backend_op(
 
 
 def _merge_snapshot(input: Any, snapshot: Mapping[str, Any]) -> Any:
-    """Prepend a snapshot's carried messages onto the next turn's input."""
+    """Prepend a snapshot's carried messages onto the next turn's input.
+
+    The driver's own continue-as-new re-invocation avoids duplicating the
+    original prompt: a Mapping input travels without its "messages" key, and
+    a bare (non-Mapping) prompt travels as-is with a snapshot marker telling
+    this merge not to re-append it (the type must survive for the user's
+    ``@workflow.run`` signature). An externally supplied ``state_snapshot``
+    plus a fresh input still composes: carried history first, new input
+    after. Agents are expected to return the accumulated transcript in
+    ``result["messages"]`` (as deepagents/LangGraph reducers do) — the carry
+    only strips input messages when the transcript is non-empty.
+    """
     raw_prior: Any = snapshot.get("messages") or []
     prior = list(raw_prior)
     if not prior:
@@ -168,6 +180,12 @@ def _merge_snapshot(input: Any, snapshot: Mapping[str, Any]) -> Any:
         raw_next: Any = input.get("messages") or []
         merged["messages"] = [*prior, *list(raw_next)]
         return merged
+    if _INPUT_CARRIED_KEY in snapshot and snapshot[_INPUT_CARRIED_KEY] == input:
+        # Internal continue-as-new of a bare prompt: this exact input is
+        # already in the transcript; it rode along only to preserve its type.
+        # A DIFFERENT bare input (an externally harvested snapshot plus a
+        # fresh prompt) falls through and composes as usual.
+        return {"messages": prior}
     return {"messages": [*prior, *_as_message_list(input)]}
 
 
@@ -202,6 +220,10 @@ async def run_deep_agent(
     its signature is ``(input, state_snapshot=None)`` — because that is how the
     carried state is threaded into the next run.
     """
+    # The carry across continue-as-new derives from the ORIGINAL input:
+    # re-threading the merged input would hand a dict to str-typed run
+    # signatures and re-prepend carried messages on every later boundary.
+    original_input = input
     # Resume path: rehydrate the result cache and fold carried messages in.
     if state_snapshot is not None:
         _serde.set_result_cache(dict(state_snapshot.get(_CACHE_KEY) or {}))
@@ -230,14 +252,37 @@ async def run_deep_agent(
             workflow.info().get_current_history_length() >= continue_as_new_after
         )
     if should_continue and _has_pending_work(result):
-        snapshot = {
-            "messages": _extract_messages(result),
+        carried = _extract_messages(result)
+        if not carried:
+            # A turn may report pending todos with an empty/pruned transcript;
+            # the conversation the agent SAW must still cross the boundary.
+            if isinstance(input, Mapping):
+                carried = _extract_messages(input)
+            else:
+                carried = _as_message_list(input)
+        snapshot: dict[str, Any] = {
+            "messages": carried,
             _CACHE_KEY: _serde.result_cache_snapshot() or {},
         }
         # ``continue_as_new`` threads positional args into the next run via
         # ``args=``; the enclosing ``@workflow.run`` receives them as
-        # ``(input, state_snapshot)``.
-        workflow.continue_as_new(args=[input, snapshot])
+        # ``(input, state_snapshot)``, so the carried input must keep the
+        # user's declared input TYPE (a dict cannot decode into a run method
+        # typed for a bare-string prompt). When the transcript already
+        # carries the conversation (including the original input messages),
+        # a Mapping input travels without its "messages" key, and a
+        # non-Mapping input travels as-is with a snapshot marker telling
+        # _merge_snapshot not to re-append it. An EMPTY transcript re-sends
+        # the input unchanged so the original prompt is never lost.
+        carry_input: Any = original_input
+        if carried:
+            if isinstance(original_input, Mapping):
+                carry_input = {
+                    k: v for k, v in original_input.items() if k != "messages"
+                }
+            else:
+                snapshot[_INPUT_CARRIED_KEY] = original_input
+        workflow.continue_as_new(args=[carry_input, snapshot])
 
     return result
 

tests/contrib/deepagents/test_backends.py

@@ -31,6 +31,7 @@
 from temporalio import workflow
 from temporalio.contrib.deepagents import DeepAgentsPlugin, TemporalBackend
 from temporalio.contrib.deepagents._tools import (
+    lookup_backend,
     register_backend,
     registered_backends,
 )
@@ -152,6 +153,21 @@ def test_temporal_backend_unregisters_on_gc() -> None:
     assert ref not in registered_backends()
 
 
+def test_temporal_backend_gc_keeps_ref_resolvable_for_inflight_activity() -> None:
+    # A backend_op activity scheduled just before a cache eviction can start
+    # AFTER the evicted wrapper is collected, and the replay that would
+    # re-register the ref only happens once that activity completes. The
+    # activity-side lookup must therefore still resolve a retired ref.
+    inner = RecordingBackend()
+    before = set(registered_backends())
+    wrapper = TemporalBackend(inner)
+    (ref,) = set(registered_backends()) - before
+    del wrapper
+    gc.collect()
+    assert ref not in registered_backends()
+    assert lookup_backend(ref) is inner
+
+
 def test_temporal_backend_gc_keeps_reregistered_ref() -> None:
     # Refs are deterministic per run: after a cache eviction, a replay
     # re-registers the SAME ref with a fresh inner backend. The evicted