diff --git a/CHANGELOG.md b/CHANGELOG.md index f1a526e7e..98fe0a0bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -56,6 +56,7 @@ to include examples, links to docs, or any other relevant information. ### Fixed +- `contrib.deepagents`: prevent duplicate input messages after continue-as-new. - `temporalio.contrib.deepagents` no longer dedups repeated identical tool, model, and backend-op calls: each dispatch runs its own Activity, and the continue-as-new result cache is retired for new executions (a continued run diff --git a/temporalio/contrib/deepagents/workflow.py b/temporalio/contrib/deepagents/workflow.py index e9bf9f203..d0635db18 100644 --- a/temporalio/contrib/deepagents/workflow.py +++ b/temporalio/contrib/deepagents/workflow.py @@ -25,6 +25,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 @@ -189,7 +190,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: @@ -199,6 +211,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)]} @@ -217,10 +235,9 @@ async def run_deep_agent( ) -> Any: """Drive ``agent.ainvoke(input)`` with continue-as-new state carry. - Once the completed turn leaves pending todos AND history has grown past the - limit, the turn's state (messages + the model/tool result cache) is - snapshotted and carried into a fresh run via ``workflow.continue_as_new``, - so long conversations do not accumulate unbounded history. + Once a completed turn leaves pending todos and history has grown past the + limit, its messages are carried into a fresh run via + ``workflow.continue_as_new``. Only legacy executions carry the result cache. By default (``continue_as_new_after=None``) the limit is the server's own recommendation — ``workflow.info().is_continue_as_new_suggested()`` — which @@ -233,6 +250,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: # Only legacy executions consult the carried cache; on new @@ -268,13 +289,36 @@ async def run_deep_agent( workflow.info().get_current_history_length() >= continue_as_new_after ) if should_continue and _has_pending_work(result): - snapshot: dict[str, Any] = {"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} if _legacy_result_cache(): snapshot[_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_continue_as_new.py b/tests/contrib/deepagents/test_continue_as_new.py index a66f95827..491e77336 100644 --- a/tests/contrib/deepagents/test_continue_as_new.py +++ b/tests/contrib/deepagents/test_continue_as_new.py @@ -29,7 +29,8 @@ ) from temporalio import workflow from temporalio.contrib.deepagents import DeepAgentsPlugin, _serde, run_deep_agent -from temporalio.worker import Worker +from temporalio.contrib.deepagents.workflow import _CACHE_KEY, _merge_snapshot +from temporalio.worker import Replayer, Worker class FakeAgent: @@ -43,7 +44,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": [ @@ -57,7 +58,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, @@ -83,9 +84,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" @@ -115,7 +114,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": [ @@ -162,7 +161,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( @@ -174,9 +173,141 @@ async def test_can_defaults_to_server_suggestion( ) +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: + """Continue-as-new carries one transcript and produces replayable histories.""" + 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 + assert _CACHE_KEY not in snapshot + + replayer = Replayer(workflows=[ContinueAsNewWorkflow], plugins=[DeepAgentsPlugin()]) + await replayer.replay_workflow(hist) + await replayer.replay_workflow(await handle.fetch_history()) + + class DiskCountingBackend: - """Each read appends to a log and reports the total — disk state, so the - count survives sandbox re-imports, replays, and continue-as-new.""" + """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" @@ -187,6 +318,77 @@ def read(self, _file_path: str) -> str: 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 + + class _CrossBoundaryAgent: """ainvoke-shaped driver issuing the SAME read every run; reports pending until the second read has observably executed."""