diff --git a/CHANGELOG.md b/CHANGELOG.md index 14a225277..12d337258 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/temporalio/contrib/deepagents/_activity.py b/temporalio/contrib/deepagents/_activity.py index 3c1501a42..99c615f02 100644 --- a/temporalio/contrib/deepagents/_activity.py +++ b/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.", diff --git a/temporalio/contrib/deepagents/_tools.py b/temporalio/contrib/deepagents/_tools.py index 1235224c6..b0d8d705f 100644 --- a/temporalio/contrib/deepagents/_tools.py +++ b/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,7 +159,7 @@ 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 @@ -162,10 +167,20 @@ def _unregister_backend(ref: str, inner: Any) -> None: 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]: @@ -173,6 +188,20 @@ def registered_backends() -> dict[str, Any]: 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 # --------------------------------------------------------------------------- diff --git a/temporalio/contrib/deepagents/workflow.py b/temporalio/contrib/deepagents/workflow.py index cb3e45b03..d6d724485 100644 --- a/temporalio/contrib/deepagents/workflow.py +++ b/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 diff --git a/tests/contrib/deepagents/test_backends.py b/tests/contrib/deepagents/test_backends.py index b32a6f12b..b0a9f339d 100644 --- a/tests/contrib/deepagents/test_backends.py +++ b/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 diff --git a/tests/contrib/deepagents/test_continue_as_new.py b/tests/contrib/deepagents/test_continue_as_new.py index c5add7cf7..329e12a12 100644 --- a/tests/contrib/deepagents/test_continue_as_new.py +++ b/tests/contrib/deepagents/test_continue_as_new.py @@ -12,10 +12,14 @@ import sys import uuid +from datetime import timedelta +from pathlib import Path from typing import Any import pytest +from temporalio.client import WorkflowExecutionStatus +from temporalio.common import RetryPolicy from temporalio.testing import WorkflowEnvironment pytestmark = pytest.mark.skipif( @@ -23,6 +27,7 @@ ) from temporalio import workflow from temporalio.contrib.deepagents import DeepAgentsPlugin, _serde, run_deep_agent +from temporalio.contrib.deepagents.workflow import _merge_snapshot from temporalio.worker import Worker @@ -37,7 +42,7 @@ class FakeAgent: async def ainvoke(self, input: Any) -> dict: messages = list(input.get("messages", [])) if isinstance(input, dict) else [] messages = [*messages, "step"] - done = len(messages) >= 3 + done = messages.count("step") >= 3 return { "messages": messages, "todos": [ @@ -51,7 +56,7 @@ class ContinueAsNewWorkflow: @workflow.run async def run(self, input: dict, state_snapshot: dict | None = None) -> dict: # Threshold of 1 means: continue-as-new as soon as there is pending work, - # which the fake agent reports until the conversation reaches 3 messages. + # which the fake agent reports until it has appended 3 steps. return await run_deep_agent( FakeAgent(), input, @@ -77,9 +82,7 @@ async def test_can_threshold_and_cache(env: WorkflowEnvironment) -> None: ) result = await handle.result() - # The only way the conversation reaches >= 3 messages is if the snapshot from - # the pre-continue-as-new run was carried into the continued run and merged. - assert len(result["messages"]) >= 3, result + assert result["messages"] == ["start", "step", "step", "step"], result assert result["todos"][0]["status"] == "completed" @@ -108,7 +111,7 @@ async def ainvoke(self, input: Any) -> dict: await workflow.sleep(0.001) messages = list(input.get("messages", [])) if isinstance(input, dict) else [] messages = [*messages, "step"] - done = len(messages) >= 3 + done = messages.count("step") >= 2 return { "messages": messages, "todos": [ @@ -155,7 +158,7 @@ async def test_can_defaults_to_server_suggestion( # Carry across the suggested continue-as-new: the conversation only reaches # 3 messages if snapshots crossed run boundaries. - assert len(result["messages"]) >= 3, result + assert result["messages"] == ["start", "step", "step"], result assert result["todos"][0]["status"] == "completed" # The first run really did continue-as-new (not complete). first = env.client.get_workflow_handle( @@ -165,3 +168,215 @@ async def test_can_defaults_to_server_suggestion( assert desc.status is not None and desc.status.name == "CONTINUED_AS_NEW", ( desc.status ) + + +def test_merge_snapshot_preserves_new_input_messages() -> None: + # External resume: a saved snapshot plus a NEW user message composes — + # carried history first, the new message after. (Replace semantics here + # would silently drop the user's latest message.) + merged = _merge_snapshot( + {"messages": ["new question"], "config": {"k": "v"}}, + {"messages": ["old q", "old a"]}, + ) + assert merged["messages"] == ["old q", "old a", "new question"] + assert merged["config"] == {"k": "v"} + + # Non-Mapping input: a bare prompt appends after the carried history. + merged = _merge_snapshot("new question", {"messages": ["old q", "old a"]}) + assert merged["messages"] == ["old q", "old a", "new question"] + + +def test_merge_snapshot_internal_carry_has_no_duplicates() -> None: + # The driver strips messages from the carried input, so the internal + # continue-as-new path resumes from the snapshot alone. + merged = _merge_snapshot({"config": {"k": "v"}}, {"messages": ["start", "step"]}) + assert merged["messages"] == ["start", "step"] + assert merged["config"] == {"k": "v"} + + +class BareInputAgent: + """ainvoke-shaped agent for a BARE-STRING input: first turn folds the + prompt into the transcript; finishes after three steps.""" + + async def ainvoke(self, input: Any) -> dict: + if isinstance(input, dict): + messages = list(input.get("messages", [])) + else: + messages = [input] + messages = [*messages, "step"] + done = messages.count("step") >= 3 + return { + "messages": messages, + "todos": [ + {"content": "work", "status": "completed" if done else "pending"} + ], + } + + +@workflow.defn +class BareInputCanWorkflow: + @workflow.run + async def run(self, input: str, state_snapshot: dict | None = None) -> dict: + # A STR-typed run signature: the carried input must decode as str + # after every continue-as-new, or the workflow stalls on task retry. + return await run_deep_agent( + BareInputAgent(), + input, + continue_as_new_after=1, + state_snapshot=state_snapshot, + ) + + +@pytest.mark.asyncio +async def test_bare_string_input_survives_continue_as_new( + env: WorkflowEnvironment, +) -> None: + """A bare-prompt input with a str-typed run signature crosses multiple + continue-as-new boundaries: the type survives (no decode failure) and the + prompt appears exactly once in the final transcript.""" + plugin = DeepAgentsPlugin() + async with Worker( + env.client, + task_queue="da-can-bare", + workflows=[BareInputCanWorkflow], + plugins=[plugin], + ): + handle = await env.client.start_workflow( + BareInputCanWorkflow.run, + "start", + id=f"da-can-bare-{uuid.uuid4()}", + task_queue="da-can-bare", + ) + result = await handle.result() + + assert result["messages"] == ["start", "step", "step", "step"], result + first = env.client.get_workflow_handle( + handle.id, run_id=handle.first_execution_run_id + ) + desc = await first.describe() + assert desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW, desc.status + + +@pytest.mark.asyncio +async def test_can_args_do_not_carry_messages(env: WorkflowEnvironment) -> None: + """The continue-as-new command's carried input omits "messages" — the + transcript rides only in the snapshot (single copy in the payload).""" + import json + + plugin = DeepAgentsPlugin() + async with Worker( + env.client, + task_queue="da-can-args", + workflows=[ContinueAsNewWorkflow], + plugins=[plugin], + ): + handle = await env.client.start_workflow( + ContinueAsNewWorkflow.run, + {"messages": ["start"], "config": {"k": "v"}}, + id=f"da-can-args-{uuid.uuid4()}", + task_queue="da-can-args", + ) + await handle.result() + first = env.client.get_workflow_handle( + handle.id, run_id=handle.first_execution_run_id + ) + hist = await first.fetch_history() + can_events = [ + e + for e in hist.events + if e.HasField("workflow_execution_continued_as_new_event_attributes") + ] + assert can_events, "first run did not continue-as-new" + payloads = can_events[ + 0 + ].workflow_execution_continued_as_new_event_attributes.input.payloads + carried_input = json.loads(payloads[0].data) + snapshot = json.loads(payloads[1].data) + + assert "messages" not in carried_input, carried_input + assert carried_input.get("config") == {"k": "v"} + assert snapshot["messages"], snapshot + + +class DiskCountingBackend: + """Each read appends to a log and reports the total (disk state survives + sandbox re-imports, replays, and continue-as-new).""" + + def __init__(self, root: str) -> None: + self._log = Path(root) / "reads.log" + + def read(self, _file_path: str) -> str: + with self._log.open("a") as f: + f.write("r\n") + return f"read:{len(self._log.read_text().splitlines())}" + + +class NoMessagesAgent: + """Returns an EMPTY transcript with pending todos on the first turn, then + answers. Turn tracking lives on disk via an activity-backed counter — the + agent object is re-created each run/replay, so in-memory state cannot + distinguish turns.""" + + def __init__(self, backend: Any) -> None: + self._backend = backend + + async def ainvoke(self, input: Any) -> dict: + turn = int((await self._backend.read("turn")).split(":")[1]) + if turn == 1: + return {"messages": [], "todos": [{"content": "w", "status": "pending"}]} + messages = list(input.get("messages", [])) if isinstance(input, dict) else [] + return { + "messages": [*messages, "answered"], + "todos": [{"content": "w", "status": "completed"}], + } + + +@workflow.defn +class EmptyTranscriptCanWorkflow: + @workflow.run + async def run(self, input: dict, state_snapshot: dict | None = None) -> dict: + from temporalio.contrib.deepagents import TemporalBackend + + backend = TemporalBackend( + DiskCountingBackend(input["root"]), + activity_options={ + "start_to_close_timeout": timedelta(seconds=30), + "retry_policy": RetryPolicy(maximum_attempts=1), + }, + ) + return await run_deep_agent( + NoMessagesAgent(backend), + input, + continue_as_new_after=1, + state_snapshot=state_snapshot, + ) + + +@pytest.mark.asyncio +async def test_empty_transcript_can_preserves_prompt( + env: WorkflowEnvironment, tmp_path: Any +) -> None: + """A turn ending with pending todos and an EMPTY transcript still carries + the conversation across a REAL continue-as-new boundary.""" + plugin = DeepAgentsPlugin() + async with Worker( + env.client, + task_queue="da-can-empty", + workflows=[EmptyTranscriptCanWorkflow], + plugins=[plugin], + max_cached_workflows=0, + ): + handle = await env.client.start_workflow( + EmptyTranscriptCanWorkflow.run, + {"messages": ["the question"], "root": str(tmp_path)}, + id=f"da-can-empty-{uuid.uuid4()}", + task_queue="da-can-empty", + ) + result = await handle.result() + first = env.client.get_workflow_handle( + handle.id, run_id=handle.first_execution_run_id + ) + desc = await first.describe() + + assert desc.status == WorkflowExecutionStatus.CONTINUED_AS_NEW, desc.status + assert result["messages"] == ["the question", "answered"], result