fix(replay): restore handler step variables on re-run; reject steps inserted before the target

src/praisonai-agents/praisonaiagents/trace/context_events.py

@@ -1294,8 +1294,14 @@ def step_end(
         route: Optional[str] = None,
         node_id: Optional[str] = None,
         agent: Optional[str] = None,
+        variables: Optional[Dict[str, Any]] = None,
     ) -> None:
-        """A flow step ends: ``completed``, ``failed``, ``skipped`` or ``cached``."""
+        """A flow step ends: ``completed``, ``failed``, ``skipped`` or ``cached``.
+
+        ``variables`` are the workflow variables the step itself set (a handler's
+        ``StepResult.variables``); recorded redacted so a re-run that reuses the
+        step can restore them.
+        """
         if not self._enabled:
             return
         data: Dict[str, Any] = {
@@ -1306,6 +1312,8 @@ def step_end(
         }
         if route is not None:
             data["route"] = route
+        if variables:
+            data["variables"] = self._redact_dict(dict(variables))
         self._emit(ContextEvent(
             event_type=ContextEventType.STEP_END,
             timestamp=time.time(),

src/praisonai-agents/praisonaiagents/workflows/workflows.py

@@ -760,7 +760,24 @@ def _replay_step_begin(flow, step, name=None, pattern=None):
             "started": time.perf_counter()}
 
 
-def _replay_step_end(span, status, output=None):
+def _differs(a: Any, b: Any) -> bool:
+    """``a != b`` that treats a failing or ambiguous comparison as "changed"."""
+    try:
+        return bool(a != b)
+    except Exception:
+        return True
+
+
+def _snapshot_scope(span, variables):
+    """Before a recorded leaf step really runs (not skipped or cached), keep a
+    per-value deep copy of the scope on its span, so a list or dict the step
+    changes in place still counts as set by it; only an uncopyable value is
+    shared by reference."""
+    if span is not None:
+        span["before"] = _isolate_scope(variables)
+
+
+def _replay_step_end(span, status, output=None, variables=None):
     """Emit ``step_end`` for a span from :func:`_replay_step_begin`."""
     if span is None:
         return
@@ -772,6 +789,7 @@ def _replay_step_end(span, status, output=None):
         route=span.get("route"),
         node_id=span["node_id"],
         agent=span["agent"],
+        variables=variables,
     )
 
 
@@ -1720,14 +1738,15 @@ def _notify_step_start(self, raw_step: Any, step: Any, context: Any) -> Optional
 
     def _notify_step_complete(
         self, span: Optional[dict], step_name: str, step_result: Any, status: str, output: Any = None,
+        variables: Optional[Dict[str, Any]] = None,
     ) -> None:
         """Call ``on_step_complete`` then emit replay ``step_end``."""
         if self.on_step_complete:
             try:
                 self.on_step_complete(step_name, step_result)
             except Exception as e:
                 logger.error(f"on_step_complete callback failed: {e}")
-        _replay_step_end(span, status, output)
+        _replay_step_end(span, status, output, variables)
 
     def _replay_graph(self) -> Dict[str, Any]:
         """Graph for ``workflow_start``; also maps step objects to node ids."""
@@ -2034,11 +2053,15 @@ def _run_impl(
                         step.status = cached_record["status"]
                     if _cached.get("variables"):
                         all_variables.update(_cached["variables"])
+                    # The cache also holds the step's output variable (its full
+                    # output): trace only what the handler set, like a cold run.
+                    _out_var = step.output_variable or f"{step.name}_output"
                     self._notify_step_complete(
                         _replay, step.name, StepResult(
                             output=previous_output or "",
                             stop_workflow=bool(_cached.get("stop")),
                         ), "cached", previous_output,
+                        variables={k: v for k, v in (_cached.get("variables") or {}).items() if k != _out_var},
                     )
                     if _cached.get("stop"):
                         if verbose:
@@ -2059,6 +2082,8 @@ def _run_impl(
             validation_feedback = None
             guardrail_failed = False
             cached_variable_updates = {} if _cache_key is not None else None
+            # Variables this step set (StepResult.variables), for replay's step_end.
+            step_variable_updates: Dict[str, Any] = {}
 
             while retry_count <= max_retries:
                 step_error = None
@@ -2081,6 +2106,7 @@ def _run_impl(
                                 # accepted attempt that wrote {a} leaves both
                                 # keys in the live run, so both must be cached.
                                 all_variables.update(result.variables)
+                                step_variable_updates.update(result.variables)
                                 if cached_variable_updates is not None:
                                     cached_variable_updates.update(result.variables)
                         else:
@@ -2337,6 +2363,7 @@ def _run_impl(
             self._notify_step_complete(
                 _replay, step.name, step_result,
                 "failed" if step_failed else "completed", output,
+                variables=step_variable_updates,
             )
 
             # Handle output_file - save output to file
@@ -3129,7 +3156,17 @@ def _execute_single_step_internal(
             "skipped" if is_dict and result.get("skipped")
             else "failed" if is_dict and result.get("error") else "completed"
         )
-        _replay_step_end(span, status, result.get("output") if is_dict else None)
+        # Variables this step set, minus its output variable, for replay's step_end.
+        # A cache hit leaves all_variables untouched, so it is its own baseline.
+        before = span.pop("before", all_variables)
+        step_vars = {}
+        if is_dict:
+            step_vars = {
+                k: v for k, v in (result.get("variables") or {}).items()
+                if k not in before or _differs(before[k], v)
+            }
+            step_vars.pop(getattr(step, "output_variable", None) or f"{result.get('step')}_output", None)
+        _replay_step_end(span, status, result.get("output") if is_dict else None, step_vars)
         return result
 
     def _execute_leaf_step(
@@ -3172,6 +3209,7 @@ def _execute_leaf_step(
                 }
         cache = getattr(self, "_step_cache", None)
         if cache is None:
+            _snapshot_scope(_replay, all_variables)
             return self._execute_single_step_uncached(
                 step, previous_output, input, all_variables, model, verbose, index,
                 stream=stream, depth=depth,
@@ -3198,6 +3236,7 @@ def _execute_leaf_step(
                 _replay["status"] = "cached"
             # Copy: a caller mutating a returned result must not edit the cache.
             return dict(hit)
+        _snapshot_scope(_replay, all_variables)
         result = self._execute_single_step_uncached(
             step, previous_output, input, all_variables, model, verbose, index,
             stream=stream, depth=depth,

src/praisonai-agents/tests/unit/trace/test_run_events.py

@@ -821,3 +821,119 @@ def test_named_team_keeps_its_name(self):
 
         team = AgentTeam(agents=[Agent(name="Researcher", output="silent")], name="EV brief")
         assert team._replay_name() == "EV brief"
+
+
+class TestStepVariables:
+    """step_end records the variables a step set, so a re-run can restore them."""
+
+    def test_step_end_records_handler_variables_redacted(self, recorder):
+        from praisonaiagents.workflows import StepResult
+        from praisonaiagents.workflows.workflows import AgentFlow
+
+        def research(ctx):
+            return StepResult(output="notes", variables={"customer_id": 42, "api_key": "sk-abcdefghijklmnop1234"})
+
+        def draft(ctx):
+            return StepResult(output="draft")
+
+        AgentFlow(name="v", steps=[research, draft]).run("go")
+        ends = {e.node_id: e.data for e in _of(recorder, "step_end")}
+        assert ends["step:research"]["variables"]["customer_id"] == 42
+        assert "sk-abcdefghijklmnop1234" not in str(ends["step:research"]["variables"])
+        assert "variables" not in ends["step:draft"]
+
+    def test_nested_step_end_records_handler_variables(self, recorder):
+        """Steps run through _execute_single_step_internal (patterns, hierarchical) record them too."""
+        from praisonaiagents.workflows import StepResult
+        from praisonaiagents.workflows.workflows import AgentFlow, If
+
+        def research(ctx):
+            return StepResult(output="notes", variables={"customer_id": 42})
+
+        AgentFlow(name="v", steps=[If(condition="1 > 0", then_steps=[research])]).run("go")
+        ends = {e.node_id: e.data for e in _of(recorder, "step_end")}
+        variables = next(d["variables"] for n, d in ends.items() if n.endswith("research"))
+        assert variables == {"customer_id": 42}
+
+    def test_cache_hit_step_end_leaves_out_the_output_variable(self, recorder):
+        """The step cache stores the output variable (the full output); a cached
+        step_end must trace only what the handler set, like the cold run."""
+        from praisonaiagents.workflows import StepResult
+        from praisonaiagents.workflows.workflows import AgentFlow
+
+        def research(ctx):
+            return StepResult(output="the full step output " * 50, variables={"customer_id": 42})
+
+        flow = AgentFlow(name="v", steps=[research], cache=True)
+        flow.run("go")
+        flow.run("go")  # cache hit
+        ends = [e.data for e in _of(recorder, "step_end")]
+        assert [d["status"] for d in ends] == ["completed", "cached"]
+        for d in ends:
+            assert d["variables"] == {"customer_id": 42}
+            assert "the full step output" not in str(d["variables"])
+
+    def test_nested_step_records_a_variable_changed_in_place(self, recorder):
+        """A list the step appends to (same object) still counts as set by it."""
+        from praisonaiagents.workflows import StepResult
+        from praisonaiagents.workflows.workflows import AgentFlow, If
+
+        def seed(ctx):
+            return StepResult(output="s", variables={"items": ["a"]})
+
+        def add(ctx):
+            ctx.variables["items"].append("b")
+            return StepResult(output="m")
+
+        AgentFlow(name="v", steps=[seed, If(condition="1 > 0", then_steps=[add])]).run("go")
+        ends = {e.node_id: e.data for e in _of(recorder, "step_end")}
+        variables = next(d.get("variables") for n, d in ends.items() if n.endswith("add"))
+        assert variables == {"items": ["a", "b"]}
+
+    def test_uncopyable_variable_does_not_hide_an_in_place_change(self, recorder):
+        """One variable that cannot be deep-copied must not stop the others being compared by value."""
+        import threading
+        from praisonaiagents.workflows import StepResult
+        from praisonaiagents.workflows.workflows import AgentFlow, If
+
+        def seed(ctx):
+            return StepResult(output="s", variables={"items": ["a"]})
+
+        def add(ctx):
+            ctx.variables["items"].append("b")
+            return StepResult(output="m")
+
+        flow = AgentFlow(name="v", steps=[seed, If(condition="1 > 0", then_steps=[add])],
+                         variables={"lock": threading.Lock()})
+        flow.run("go")
+        ends = {e.node_id: e.data for e in _of(recorder, "step_end")}
+        variables = next(d.get("variables") for n, d in ends.items() if n.endswith("add"))
+        assert variables == {"items": ["a", "b"]}
+
+    def test_skipped_or_cached_nested_step_does_not_copy_the_scope(self, recorder, monkeypatch):
+        """Only a nested step that really runs snapshots the variables; a cache
+        hit still records what the handler set."""
+        from praisonaiagents.workflows import StepResult, Task
+        from praisonaiagents.workflows import workflows as wf
+
+        copies = []
+        real = wf._isolate_scope
+        monkeypatch.setattr(wf, "_isolate_scope", lambda v: copies.append(1) or real(v))
+
+        def research(ctx):
+            return StepResult(output="notes", variables={"customer_id": 42})
+
+        def never(ctx):
+            return StepResult(output="x")
+
+        steps = [wf.If(condition="1 > 0", then_steps=[
+            research, Task(name="never", handler=never, should_run=lambda ctx: False),
+        ])]
+        flow = wf.AgentFlow(name="v", steps=steps, variables={"items": list(range(1000))}, cache=True)
+        flow.run("go")
+        assert len(copies) == 1  # research ran; the skipped step copied nothing
+        copies.clear()
+        flow.run("go")  # research is a cache hit
+        assert copies == []
+        cached = [e.data for e in _of(recorder, "step_end") if e.data["status"] == "cached"]
+        assert [d.get("variables") for d in cached] == [{"customer_id": 42}]

src/praisonai/praisonai/replay/rerun.py

@@ -30,6 +30,7 @@
 from collections import OrderedDict
 from typing import Any, Callable, Dict, Iterator, List, Optional, Union
 
+from praisonaiagents.trace.redact import REDACTED_VALUE
 from praisonaiagents.workflows.step_cache import StepCacheProtocol
 
 from .events import as_dicts, etype, load_events, pair_spans, run_id_of
@@ -66,13 +67,17 @@ def __init__(
         kind: str,
         outputs: "OrderedDict[str, Any]",
         variables: Optional[Dict[str, Dict[str, Any]]] = None,
+        started_before: Optional[List[str]] = None,
     ):
         self.source_run_id = source_run_id
         self.node_id = node_id
         self.kind = kind  # "flow" | "team"
         self.outputs = outputs
         # Variable updates a step's handler made, when its step_end recorded them.
         self.variables = variables or {}
+        # Every node of the target's kind that STARTED before it, whatever its
+        # status (skipped and failed steps too): what the recorded run executed.
+        self.started_before = list(started_before or [])
 
     @property
     def forked_from(self) -> Dict[str, str]:
@@ -150,6 +155,10 @@ def plan_rerun(events: Union[str, List[Any]], node_id: str) -> RerunPlan:
 
     outputs: "OrderedDict[str, Any]" = OrderedDict()
     variables: Dict[str, Dict[str, Any]] = {}
+    started_before = [
+        s["start"]["node_id"] for s in nodes
+        if s["type"] == target["type"] and s["seq"] < target["seq"]
+    ]
     for span in nodes:
         if span["type"] != target["type"] or span["end"] is None:
             continue
@@ -164,8 +173,16 @@ def plan_rerun(events: Union[str, List[Any]], node_id: str) -> RerunPlan:
         outputs[span["start"]["node_id"]] = _recover_output(span, evs)
         recorded_vars = (span["end"].get("data") or {}).get("variables")
         if isinstance(recorded_vars, dict):
+            # A redacted value cannot be restored: the marker would replace the
+            # real value in the re-run's variables.
+            hidden = [k for k, v in recorded_vars.items() if REDACTED_VALUE in repr(v)]
+            if hidden:
+                raise RerunError(
+                    f"{span['start']['node_id']} set {', '.join(map(repr, hidden))}, which the "
+                    "recording redacted, so it cannot be restored; re-run from that node or earlier."
+                )
             variables[span["start"]["node_id"]] = recorded_vars
-    return RerunPlan(run_id_of(evs), node_id, kind, outputs, variables)
+    return RerunPlan(run_id_of(evs), node_id, kind, outputs, variables, started_before)
 
 
 # ---------------------------------------------------------------------------
@@ -268,11 +285,22 @@ def _flow_step_names(steps: Any) -> set:
     return names
 
 
-def _check_flow_order(steps: List[Any], recorded: List[str]) -> None:
+def _check_flow_order(
+    steps: List[Any],
+    recorded: List[str],
+    started_before: Optional[List[str]] = None,
+    node_ids: Optional[Dict[int, str]] = None,
+) -> None:
     """Reject a sequential flow whose top-level steps are not in the recorded order.
 
     Only checked when the flow runs strictly in sequence: ``next_steps`` /
     ``branch_condition`` jumps or a repeated node (``#2``) can revisit steps.
+    Also rejects a step inserted before the target that never ran in the
+    recording: it would run fresh and replace the reused output the target
+    expects as its input. Steps are matched by graph node id (``node_ids``,
+    id(step) -> node id), so unnamed patterns such as ``If`` are covered too.
+    A target nested in a pattern is checked up to that pattern: the last
+    top-level step the recording started before the target.
     """
     if any(getattr(s, "next_steps", None) or getattr(s, "branch_condition", None) for s in steps):
         return
@@ -284,6 +312,20 @@ def _check_flow_order(steps: List[Any], recorded: List[str]) -> None:
             value = getattr(step, attr, None)
             if isinstance(value, str):
                 position.setdefault(value, index)
+    if started_before is not None and node_ids is not None and recorded:
+        ran = set(started_before)
+        top = [node_ids.get(id(step)) for step in steps]
+        if recorded[-1] in top:
+            anchor = top.index(recorded[-1])
+        else:  # nested target: its container is the last top-level step that started
+            anchor = max((i for i, node_id in enumerate(top) if node_id in ran), default=-1)
+        if anchor > 0:
+            for step, node_id in zip(steps, top[:anchor]):
+                if node_id not in ran:
+                    raise RerunError(
+                        f"Step {node_id or type(step).__name__!r} comes before {recorded[-1]!r} in "
+                        "this flow but did not run in the recording; re-run against the flow that recorded it."
+                    )
     last = -1
     for node_id in recorded:  # nested steps have no top-level position: skipped
         index = position.get(_step_name(node_id))
@@ -299,6 +341,7 @@ def _check_flow_order(steps: List[Any], recorded: List[str]) -> None:
 
 def seed_flow(flow: Any, plan: RerunPlan) -> TraceStepCache:
     """Install a :class:`TraceStepCache` for ``plan`` on an AgentFlow."""
+    from praisonaiagents.workflows.diagram import _flow_graph
     from praisonaiagents.workflows.step_cache import resolve_step_cache
 
     # A different or edited flow would silently run unmatched steps fresh.
@@ -307,7 +350,9 @@ def seed_flow(flow: Any, plan: RerunPlan) -> TraceStepCache:
     for node_id in recorded:
         if _step_name(node_id) not in names:
             raise RerunError(f"Recorded node {node_id!r} does not match any step in this flow.")
-    _check_flow_order(getattr(flow, "steps", None) or [], recorded)
+    _check_flow_order(
+        getattr(flow, "steps", None) or [], recorded, plan.started_before, _flow_graph(flow).ids,
+    )
 
     cache = TraceStepCache(plan.outputs, plan.variables).bind(flow)
     flow.cache = cache