diff --git a/src/praisonai-agents/praisonaiagents/trace/context_events.py b/src/praisonai-agents/praisonaiagents/trace/context_events.py index 7e4d7433ad..926338944c 100644 --- a/src/praisonai-agents/praisonaiagents/trace/context_events.py +++ b/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(), diff --git a/src/praisonai-agents/praisonaiagents/workflows/workflows.py b/src/praisonai-agents/praisonaiagents/workflows/workflows.py index 5aec581a0d..83df9a55f6 100644 --- a/src/praisonai-agents/praisonaiagents/workflows/workflows.py +++ b/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,6 +1738,7 @@ 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: @@ -1727,7 +1746,7 @@ def _notify_step_complete( 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, diff --git a/src/praisonai-agents/tests/unit/trace/test_run_events.py b/src/praisonai-agents/tests/unit/trace/test_run_events.py index 722e6f1774..fdcbdd43e7 100644 --- a/src/praisonai-agents/tests/unit/trace/test_run_events.py +++ b/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}] diff --git a/src/praisonai/praisonai/replay/rerun.py b/src/praisonai/praisonai/replay/rerun.py index 9add90bbe8..27df48cc79 100644 --- a/src/praisonai/praisonai/replay/rerun.py +++ b/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,6 +67,7 @@ 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 @@ -73,6 +75,9 @@ def __init__( 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 diff --git a/src/praisonai/tests/unit/replay/test_replay_rerun.py b/src/praisonai/tests/unit/replay/test_replay_rerun.py index 91b19f8ae5..e01c2ae127 100644 --- a/src/praisonai/tests/unit/replay/test_replay_rerun.py +++ b/src/praisonai/tests/unit/replay/test_replay_rerun.py @@ -411,3 +411,189 @@ def polish(ctx): with pytest.raises(RerunError, match="different order"): rerun(AgentFlow(steps=[draft, research, polish]), str(src), "step:draft", input="topic") assert ran == [] + + +def test_flow_rerun_rejects_step_inserted_before_target(traces_dir, fx): + """research -> extra -> draft: 'extra' never ran, so draft would get its output, not research's.""" + from praisonaiagents import AgentFlow + from praisonaiagents.workflows import StepResult + + from praisonai.replay.rerun import RerunError, rerun + + ran = [] + + def research(ctx): + ran.append("research") + return StepResult(output="r") + + def extra(ctx): + ran.append("extra") + return StepResult(output="e") + + def draft(ctx): + ran.append("draft") + return StepResult(output="d") + + def polish(ctx): + ran.append("polish") + return StepResult(output="p") + + src = fx.write_trace(traces_dir, "run-flow01", fx.flow_events()) + with pytest.raises(RerunError, match="did not run in the recording"): + rerun(AgentFlow(steps=[research, extra, draft, polish]), str(src), "step:draft", input="topic") + assert ran == [] + + +def test_recorded_handler_variables_survive_a_rerun(traces_dir): + """End to end: core records StepResult.variables on step_end; rerun restores them.""" + import json + + from praisonaiagents import AgentFlow + from praisonaiagents.workflows import StepResult + + from praisonai.replay import record + from praisonai.replay.rerun import rerun + + def make(customer_id, seen, calls): + def research(ctx): + calls.append("research") + return StepResult(output="notes", variables={"customer_id": customer_id}) + + def draft(ctx): + seen.append(ctx.variables.get("customer_id")) + return StepResult(output="draft") + + return AgentFlow(steps=[research, draft]) + + with record(full_content=True) as rec: + make(42, [], []).start("topic") + ends = [json.loads(line) for line in open(rec.path) if '"step_end"' in line] + assert any((e["data"].get("variables") or {}).get("customer_id") == 42 for e in ends) + + seen, calls = [], [] + rerun(make(0, seen, calls), rec.path, "step:draft", input="topic") + assert calls == [] # research reused from the recording + assert seen == [42] # with the variable it set, not the fresh run's 0 + + +def test_flow_rerun_rejects_unnamed_pattern_inserted_before_target(traces_dir, fx): + """An If (no name) inserted before the target is caught by its graph node id.""" + from praisonaiagents import AgentFlow + from praisonaiagents.workflows import StepResult + from praisonaiagents.workflows.workflows import If + + from praisonai.replay.rerun import RerunError, rerun + + ran = [] + + def step(name): + def handler(ctx): + ran.append(name) + return StepResult(output=name) + handler.__name__ = name + return handler + + src = fx.write_trace(traces_dir, "run-flow01", fx.flow_events()) + flow = AgentFlow(steps=[step("research"), If(condition="1 > 0", then_steps=[step("extra")]), + step("draft"), step("polish")]) + with pytest.raises(RerunError, match="did not run in the recording"): + rerun(flow, str(src), "step:draft", input="topic") + assert ran == [] + + +def test_flow_rerun_rejects_redacted_variables(traces_dir, fx): + """A redacted variable would be restored as the literal marker: refuse instead.""" + from praisonaiagents import AgentFlow + from praisonaiagents.workflows import StepResult + + from praisonai.replay.rerun import RerunError, rerun + + evs = fx.flow_events() + for ev in evs: + if ev["event_type"] == "step_end" and ev.get("node_id") == "step:research": + ev["data"]["variables"] = {"api_key": "[REDACTED]", "customer_id": 42} + src = fx.write_trace(traces_dir, "run-flow01", evs) + + def research(ctx): + return StepResult(output="r") + + def draft(ctx): + return StepResult(output="d") + + def polish(ctx): + return StepResult(output="p") + + with pytest.raises(RerunError, match="'api_key'.*redacted"): + rerun(AgentFlow(steps=[research, draft, polish]), str(src), "step:draft", input="topic") + + +def _nested_flow(ran, extra=False): + from praisonaiagents import AgentFlow + from praisonaiagents.workflows import StepResult + from praisonaiagents.workflows.workflows import If + + def step(name): + def handler(ctx): + ran.append(name) + return StepResult(output=name) + handler.__name__ = name + return handler + + steps = [step("research")] + ([step("extra")] if extra else []) + return AgentFlow(steps=steps + [If(condition="1 > 0", then_steps=[step("draft")]), step("polish")]) + + +def test_flow_rerun_of_a_nested_target_rejects_a_step_inserted_before_its_container(traces_dir): + """research -> If[draft] recorded; research -> extra -> If[draft] would feed draft extra's output.""" + from praisonai.replay import record + from praisonai.replay.rerun import RerunError, rerun + + ran = [] + with record(full_content=True) as rec: + _nested_flow(ran).start("x") + ran.clear() + with pytest.raises(RerunError, match="did not run in the recording"): + rerun(_nested_flow(ran, extra=True), rec.path, "step:draft", input="x") + assert ran == [] + rerun(_nested_flow(ran), rec.path, "step:draft", input="x") # the same flow still re-runs + assert ran == ["draft", "polish"] + + +def test_flow_rerun_keeps_a_recorded_skip_skipped(traces_dir): + """A step skipped by its should_run is re-checked on re-run from the restored + state: the same condition skips it again and the target gets the same input.""" + from praisonaiagents import AgentFlow, Task + from praisonaiagents.workflows import StepResult + + from praisonai.replay import record + from praisonai.replay.rerun import rerun + + ran, inputs = [], [] + + def make(): + def research(ctx): + ran.append("research") + return StepResult(output="notes", variables={"mode": "brief"}) + + def extra(ctx): + ran.append("extra") + return StepResult(output="EXTRA") + + def draft(ctx): + ran.append("draft") + inputs.append(ctx.previous_result) + return StepResult(output="draft") + + return AgentFlow(steps=[ + research, + Task(name="extra", handler=extra, should_run=lambda ctx: ctx.variables.get("mode") == "long"), + draft, + ]) + + with record(full_content=True) as rec: + make().start("x") + assert ran == ["research", "draft"] + ran.clear() + rerun(make(), rec.path, "step:draft", input="x") + assert ran == ["draft"] # research reused, extra still skipped + assert inputs[0] == inputs[1] == "notes"