pygpt-net 2.6.62__py3-none-any.whl → 2.6.63__py3-none-any.whl

pygpt_net/CHANGELOG.txt

@@ -1,3 +1,8 @@
+2.6.63 (2025-09-27)
+
+- Improved agents' workflows.
+- Enhanced the display of agents' steps in the UI.
+
 2.6.62 (2025-09-26)
 
 - Enhanced agent workflow execution.

pygpt_net/__init__.py

@@ -6,15 +6,15 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.09.26 00:00:00                  #
+# Updated Date: 2025.09.27 00:00:00                  #
 # ================================================== #
 
 __author__ = "Marcin Szczygliński"
 __copyright__ = "Copyright 2025, Marcin Szczygliński"
 __credits__ = ["Marcin Szczygliński"]
 __license__ = "MIT"
-__version__ = "2.6.62"
-__build__ = "2025-09-26"
+__version__ = "2.6.63"
+__build__ = "2025-09-27"
 __maintainer__ = "Marcin Szczygliński"
 __github__ = "https://github.com/szczyglis-dev/py-gpt"
 __report__ = "https://github.com/szczyglis-dev/py-gpt/issues"

pygpt_net/controller/presets/editor.py

@@ -6,7 +6,7 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.09.24 00:00:00                  #
+# Updated Date: 2025.09.27 00:00:00                  #
 # ================================================== #
 
 import datetime
@@ -393,6 +393,9 @@ class Editor:
                         value=extra_options[key].get('default', None),
                     )
 
+            # ensure combo defaults are effectively applied for this tab (only empty values are updated)
+            self._apply_combo_defaults_for_group(option_key, extra_options)
+
     def load_extra_defaults(self):
         """Load extra options defaults for preset editor"""
         if not self.tab_options_idx:
@@ -423,6 +426,8 @@ class Editor:
                             option=extra_options[key],
                             value=value,
                         )
+                # ensure combo defaults are effectively applied for this tab (only empty values are updated)
+                self._apply_combo_defaults_for_group(option_key, extra_options)
 
     def load_extra_defaults_current(self):
         """Load extra options defaults on mode change"""
@@ -479,6 +484,8 @@ class Editor:
                             option=extra_options[key],
                             value=value,
                         )
+                # ensure combo defaults are effectively applied for this tab (only empty values are updated)
+                self._apply_combo_defaults_for_group(option_key, extra_options)
 
     def append_extra_options(self, preset: PresetItem):
         """
@@ -785,6 +792,9 @@ class Editor:
                             value=opt_schema.get('default'),
                         )
 
+            # ensure combo defaults are effectively applied for this tab (only empty values are updated)
+            self._apply_combo_defaults_for_group(config_id, schema_options)
+
         # 4) Recompute mapping fully based on actual tabs and their 'agent_id' properties.
         self._rebuild_tab_index_mapping()
 
@@ -1521,6 +1531,9 @@ class Editor:
                                 value=opt_schema.get('default'),
                             )
 
+                    # ensure combo defaults are effectively applied for this tab (only empty values are updated)
+                    self._apply_combo_defaults_for_group(config_id, schema_options)
+
             # 7) Recompute the index mapping strictly from the QTabWidget
             self._rebuild_tab_index_mapping()
 
@@ -1529,3 +1542,54 @@ class Editor:
 
         finally:
             tabs.setUpdatesEnabled(True)
+
+    # ---------- Helpers for reliable combo defaults in agent extra options ----------
+
+    def _apply_combo_defaults_for_group(self, parent_id: str, schema_options: Dict[str, Any]) -> None:
+        """
+        Ensure that combo-type inputs inside a given UI config group have their default values applied
+        when the current value is empty ("", None or "_"). This avoids the situation where combo boxes
+        remain uninitialized while other field types receive defaults correctly.
+
+        This function never overrides a non-empty value set by the user or loaded from a preset.
+        """
+        if not schema_options:
+            return
+
+        get_value = self.window.controller.config.get_value
+        apply_value = self.window.controller.config.apply_value
+
+        for key, opt_schema in schema_options.items():
+            if not isinstance(opt_schema, dict):
+                continue
+            if opt_schema.get('type') != 'combo':
+                continue
+
+            default_val = opt_schema.get('default', None)
+            if default_val is None:
+                continue
+
+            current_val = get_value(
+                parent_id=parent_id,
+                key=key,
+                option=opt_schema,
+            )
+
+            # Treat "_", "", None as empty and safe to replace with default
+            if current_val in (None, "", "_"):
+                # First try apply_value (standard path)
+                apply_value(
+                    parent_id=parent_id,
+                    key=key,
+                    option=opt_schema,
+                    value=default_val,
+                )
+                # Additionally set directly on widget if accessible to guard against timing of key population
+                try:
+                    widget_group = self.window.ui.config.get(parent_id, {})
+                    widget = widget_group.get(key)
+                    if widget and hasattr(widget, "set_value"):
+                        widget.set_value(default_val)
+                except Exception:
+                    # Silent fallback; apply_value above should already handle most cases
+                    pass

pygpt_net/core/agents/custom/llama_index/runner.py

@@ -1,3 +1,5 @@
+# core/agents/runners/llama_workflow.py
+
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 # ================================================== #
@@ -6,7 +8,7 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.09.26 17:00:00                  #
+# Updated Date: 2025.09.27 06:00:00                  #
 # ================================================== #
 
 from __future__ import annotations
@@ -207,37 +209,6 @@ class DynamicFlowWorkflowLI(Workflow):
                 return False
         return False
 
-    def _friendly_map(self) -> Dict[str, str]:
-        return {aid: a.name or aid for aid, a in self.fs.agents.items()}
-
-    def _friendly_map_for_routes(self, route_ids: List[str]) -> Dict[str, Any]:
-        """
-        Build a friendly map for the given route ids:
-        - Always include a human-friendly name.
-        - Include role only if provided in preset options or schema and non-empty.
-        """
-        out: Dict[str, Any] = {}
-        for rid in route_ids or []:
-            a = self.fs.agents.get(rid)
-            name = (a.name if a and a.name else rid)
-            # Prefer preset option, then schema role
-            role_opt = None
-            try:
-                role_opt = self.option_get(rid, "role", None)
-            except Exception:
-                role_opt = None
-            role_schema = getattr(a, "role", None) if a is not None else None
-            role_val = None
-            if isinstance(role_opt, str) and role_opt.strip():
-                role_val = role_opt.strip()
-            elif isinstance(role_schema, str) and role_schema.strip():
-                role_val = role_schema.strip()
-            item = {"name": name}
-            if role_val:
-                item["role"] = role_val
-            out[rid] = item
-        return out
-
     async def _emit(self, ctx: Context, ev: Any):
         if self.dbg.event_echo:
             self.logger.debug(f"[event] emit {ev.__class__.__name__}")
@@ -245,8 +216,8 @@ class DynamicFlowWorkflowLI(Workflow):
 
     async def _emit_agent_text(self, ctx: Context, text: str, agent_name: str = "Agent"):
         """
-        Emit AgentStream(delta=text) robustly. If your env requires extra fields,
-        fall back to extended AgentStream like in your SupervisorWorkflow.
+        Emit AgentStream(delta=text) robustly. If env requires extra fields,
+        fall back to extended AgentStream.
         """
         try:
             if self.dbg.event_echo:
@@ -266,16 +237,11 @@ class DynamicFlowWorkflowLI(Workflow):
             )
 
     async def _emit_header(self, ctx: Context, name: str):
-        if self.dbg.event_echo:
-            self.logger.debug(f"[event] header emit begin name='{name}'")
+        # Lightweight header to ensure agent name is known before tokens.
         await self._emit_agent_text(ctx, "", agent_name=name)
-        # await self._emit_agent_text(ctx, f"\n\n**{name}**\n\n", agent_name=name)
-        if self.dbg.event_echo:
-            self.logger.debug("[event] header emit done")
 
     async def _emit_step_sep(self, ctx: Context, node_id: str):
         try:
-            # Include human-friendly agent name in StepEvent meta for downstream ctx propagation.
             a = self.fs.agents.get(node_id)
             friendly_name = (a.name if a and a.name else node_id)
             await self._emit(
@@ -350,6 +316,9 @@ class DynamicFlowWorkflowLI(Workflow):
             return user_msg, [], "no-mem:last_output"
 
     async def _update_memory_after_step(self, node_id: str, user_msg_text: str, display_text: str):
+        """
+        Update per-node memory after a step, storing baton user message and assistant output.
+        """
         mem_id = self.g.agent_to_memory.get(node_id)
         mem_state = self.mem.get(mem_id) if mem_id else None
         if not mem_state:
@@ -374,7 +343,7 @@ class DynamicFlowWorkflowLI(Workflow):
     # ============== Workflow steps ==============
 
     def run(self, query: str, ctx: Optional[Context] = None, memory: Any = None, verbose: bool = False, on_stop=None):
-        """Entry point used by your LlamaWorkflow runner."""
+        """Entry point used by LlamaWorkflow runner."""
         self._on_stop = on_stop
 
         # Build initial chat once
@@ -444,8 +413,9 @@ class DynamicFlowWorkflowLI(Workflow):
             return FlowTickEvent() if self._current_ids else FlowStopEvent(final_answer=self._last_plain_output or "")
 
         node: AgentNode = self.fs.agents[current_id]
-        if self._steps > 1:
-            await self._emit_step_sep(ctx, current_id)
+
+        # IMPORTANT: emit StepEvent also for the very first agent step.
+        await self._emit_step_sep(ctx, current_id)
         await self._emit_header(ctx, node.name or current_id)
 
         # Resolve runtime + per-node LLM/tools
@@ -474,11 +444,10 @@ class DynamicFlowWorkflowLI(Workflow):
                 f"user='{ellipsize(user_msg_text, self.dbg.preview_chars)}'"
             )
 
-        # Prepare friendly map with optional roles for this node's allowed routes
+        # Build agent
         allowed_routes_now = list(node.outputs or [])
-        friendly_map = self._friendly_map_for_routes(allowed_routes_now)
+        friendly_map = {rid: self.fs.agents.get(rid).name or rid for rid in allowed_routes_now if rid in self.fs.agents}
 
-        # Build agent (chat_history/max_iterations in ctor – best practice)
         built = self.factory.build(
             node=node,
             node_runtime=node_rt,
@@ -528,9 +497,6 @@ class DynamicFlowWorkflowLI(Workflow):
             display_text = decision.content or ""
             if display_text:
                 await self._emit_agent_text(ctx, display_text, agent_name=(node.name or current_id))
-            if self.dbg.log_memory_dump:
-                self.logger.debug(f"[mem.prep] node={current_id} save user='{ellipsize(user_msg_text, self.dbg.preview_chars)}' "
-                                  f"assist='{ellipsize(display_text, self.dbg.preview_chars)}'")
             await self._update_memory_after_step(current_id, user_msg_text, display_text)
             next_id = decision.route if decision.valid else (allowed_routes[0] if allowed_routes else None)
             if self.dbg.log_routes:
@@ -541,9 +507,6 @@ class DynamicFlowWorkflowLI(Workflow):
             display_text = raw_text_clean or ""
             if display_text:
                 await self._emit_agent_text(ctx, display_text, agent_name=(node.name or current_id))
-            if self.dbg.log_memory_dump:
-                self.logger.debug(f"[mem.prep] node={current_id} save user='{ellipsize(user_msg_text, self.dbg.preview_chars)}' "
-                                  f"assist='{ellipsize(display_text, self.dbg.preview_chars)}'")
             await self._update_memory_after_step(current_id, user_msg_text, display_text)
             outs = self.g.get_next(current_id)
             next_id = outs[0] if outs else self.g.first_connected_end(current_id)

pygpt_net/core/agents/custom/runner.py

@@ -6,7 +6,7 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2025.09.26 17:00:00                  #
+# Updated Date: 2025.09.27 06:00:00                  #
 # ================================================== #
 
 from __future__ import annotations
@@ -62,13 +62,133 @@ class FlowOrchestrator:
       - First agent (in the whole flow) gets full initial messages from the app.
       - Next agent WITHOUT memory gets only last step's displayed content as a single 'user' message.
       - Agent WITH memory:
-          * if memory has items -> use them;
-          * if memory empty -> seed input from last displayed content (or initial messages as fallback).
+          * if memory has items -> use base history (items[:-1]) and pass last displayed content as a single 'user' baton;
+          * if memory empty -> seed baton from last displayed content (or initial messages as fallback).
     """
     def __init__(self, window, logger: Optional[Logger] = None) -> None:
         self.window = window
         self.logger = logger or NullLogger()
 
+    # ---------- Helpers (production-ready) ----------
+
+    def _extract_text_from_item(self, item: TResponseInputItem) -> str:
+        """Best-effort extract plain text from TResponseInputItem."""
+        if isinstance(item, dict):
+            content = item.get("content", "")
+            if isinstance(content, str):
+                return content
+            if isinstance(content, list):
+                parts = []
+                for p in content:
+                    if isinstance(p, dict):
+                        t = p.get("text")
+                        if isinstance(t, str):
+                            parts.append(t)
+                return "\n".join(parts)
+            return ""
+        if isinstance(item, str):
+            return item
+        return ""
+
+    def _build_baton_input(
+        self,
+        *,
+        node_id: str,
+        g: FlowGraph,
+        mem: MemoryManager,
+        initial_messages: List[TResponseInputItem],
+        first_dispatch_done: bool,
+        last_plain_output: str,
+        dbg: DebugConfig,
+    ) -> tuple[List[TResponseInputItem], str, Optional[str], Any, str]:
+        """
+        Returns: (prepared_items, baton_user_text, mem_id, mem_state, source_tag)
+        Mirrors LI baton/memory policy.
+        """
+        mem_id = g.agent_to_memory.get(node_id)
+        mem_state = mem.get(mem_id) if mem_id else None
+
+        baton_user_text = ""
+        source = ""
+
+        if mem_state and mem_state.items:
+            # memory with history -> base history + baton from last output (preferred)
+            base_items = list(mem_state.items[:-1]) if len(mem_state.items) >= 1 else []
+            if last_plain_output and last_plain_output.strip():
+                baton_user_text = last_plain_output
+                prepared = base_items + [{"role": "user", "content": baton_user_text}]
+                source = "memory:existing_to_user_baton"
+            else:
+                # fallback: use last assistant content as baton
+                last_ass = mem_state.items[-1] if isinstance(mem_state.items[-1], dict) else {}
+                if isinstance(last_ass.get("content"), str):
+                    baton_user_text = last_ass.get("content", "")
+                elif isinstance(last_ass.get("content"), list) and last_ass["content"]:
+                    baton_user_text = last_ass["content"][0].get("text", "") or ""
+                else:
+                    baton_user_text = ""
+                prepared = base_items + [{"role": "user", "content": baton_user_text}]
+                source = "memory:existing_to_last_assistant"
+            return sanitize_input_items(prepared), baton_user_text, mem_id, mem_state, source
+
+        if mem_state:
+            # memory attached but empty -> seed from last output else from initial (use last user msg as baton)
+            if last_plain_output and last_plain_output.strip():
+                baton_user_text = last_plain_output
+                prepared = [{"role": "user", "content": baton_user_text}]
+                source = "memory:seed_from_last_output"
+            else:
+                base_items = list(initial_messages[:-1]) if initial_messages else []
+                last_item = initial_messages[-1] if initial_messages else {"role": "user", "content": ""}
+                baton_user_text = self._extract_text_from_item(last_item)
+                prepared = base_items + [{"role": "user", "content": baton_user_text}]
+                source = "memory:seed_from_initial"
+            return sanitize_input_items(prepared), baton_user_text, mem_id, mem_state, source
+
+        # no memory attached
+        if not first_dispatch_done:
+            # first agent: pass initial messages as-is; baton is last user text (for potential external memory)
+            last_item = initial_messages[-1] if initial_messages else {"role": "user", "content": ""}
+            baton_user_text = self._extract_text_from_item(last_item)
+            return sanitize_input_items(list(initial_messages)), baton_user_text, None, None, "no-mem:first_initial"
+        else:
+            baton_user_text = last_plain_output if last_plain_output and last_plain_output.strip() else (
+                self._extract_text_from_item(initial_messages[-1]) if initial_messages else ""
+            )
+            prepared = [{"role": "user", "content": baton_user_text}]
+            return sanitize_input_items(prepared), baton_user_text, None, None, "no-mem:last_output"
+
+    def _update_memory_after_step(
+        self,
+        *,
+        node_id: str,
+        mem_state: Any,
+        baton_user_text: str,
+        display_text: str,
+        last_response_id: Optional[str],
+        dbg: DebugConfig,
+    ) -> None:
+        """Update memory strictly with [user baton, assistant display_text], mirroring LI semantics."""
+        if not mem_state:
+            return
+        base_items = list(mem_state.items[:-1]) if getattr(mem_state, "items", None) else []
+        new_mem = (base_items or []) + [
+            {"role": "user", "content": baton_user_text or ""},
+            {"role": "assistant", "content": [{"type": "output_text", "text": display_text or ""}]},
+        ]
+        try:
+            mem_state.set_from(new_mem, last_response_id)
+            if dbg.log_inputs:
+                self.logger.debug(
+                    f"[memory] {node_id} updated len {len(base_items)} -> {len(new_mem)} "
+                    f"user='{ellipsize(baton_user_text or '', dbg.preview_chars)}' "
+                    f"assist='{ellipsize(display_text or '', dbg.preview_chars)}'"
+                )
+        except Exception as e:
+            self.logger.error(f"[memory] update failed for {node_id}: {e}")
+
+    # ---------- Main flow ----------
+
     async def run_flow(
         self,
         schema: List[Dict[str, Any]],
@@ -174,55 +294,34 @@ class FlowOrchestrator:
                     f" role='{node_rt.role}'"
                 )
 
-            # Memory selection and INPUT BUILD (memory-first policy)
-            mem_id = g.agent_to_memory.get(current_id)
-            mem_state = mem.get(mem_id) if mem_id else None
-            if dbg.log_runtime:
-                mem_info = f"{mem_id} (len={len(mem_state.items)})" if mem_state and mem_state.items else (mem_id or "-")
-                self.logger.debug(f"[memory] attached={bool(mem_id)} mem_id={mem_info}")
-
-            input_items: List[TResponseInputItem]
-            input_source = ""
-
-            if mem_state and not mem_state.is_empty():
-                # memory present and already has history
-                input_items = list(mem_state.items)
-                input_source = "memory:existing"
-            elif mem_state:
-                # memory present but empty -> seed from last output or initial messages
-                if last_plain_output and last_plain_output.strip():
-                    input_items = [{"role": "user", "content": last_plain_output}]
-                    input_source = "memory:seeded_from_last"
-                else:
-                    input_items = list(initial_messages)
-                    input_source = "memory:seeded_from_initial"
-            else:
-                # no memory -> first agent gets initial messages; others get only last output
-                if not first_dispatch_done:
-                    input_items = list(initial_messages)
-                    input_source = "no-mem:initial"
-                else:
-                    if last_plain_output and last_plain_output.strip():
-                        input_items = [{"role": "user", "content": last_plain_output}]
-                        input_source = "no-mem:last_output"
-                    else:
-                        input_items = list(initial_messages)
-                        input_source = "no-mem:fallback_initial"
-
-            prepared_items = sanitize_input_items(input_items)
+            # Input build using baton policy (LI parity)
+            prepared_items, baton_user_text, mem_id, mem_state, input_source = self._build_baton_input(
+                node_id=current_id,
+                g=g,
+                mem=mem,
+                initial_messages=initial_messages,
+                first_dispatch_done=first_dispatch_done,
+                last_plain_output=last_plain_output,
+                dbg=dbg,
+            )
 
             if dbg.log_inputs:
                 self.logger.debug(f"[input] source={input_source} items={len(prepared_items)} "
                                   f"preview={items_preview(prepared_items, dbg.preview_chars)}")
+                if mem_id:
+                    mem_info = f"{mem_id} (len={len(mem_state.items) if mem_state else 0})"
+                    self.logger.debug(f"[memory] attached={bool(mem_id)} mem_id={mem_info}")
 
             # Build agent with per-node runtime
+            # Restrict friendly_map only to allowed outgoing routes of current node
+            allowed_map = {rid: fs.agents[rid].name or rid for rid in (node.outputs or []) if rid in fs.agents}
             built = factory.build(
                 node=node,
                 node_runtime=node_rt,
                 preset=preset,
                 function_tools=function_tools,
                 force_router=False,  # auto on multi-output
-                friendly_map={aid: a.name or aid for aid, a in fs.agents.items()},
+                friendly_map=allowed_map,
                 handoffs_enabled=True,
                 context=agent_kwargs.get("context"),
             )
@@ -243,37 +342,42 @@ class FlowOrchestrator:
 
             # Header for UI
             ctx.set_agent_name(agent.name)
-            # title = f"\n\n**{built.name}**\n\n"
-            # ctx.stream = title
             bridge.on_step(ctx, begin)
             begin = False
             handler.begin = begin
-            # if not use_partial_ctx:
-                # handler.to_buffer(title)
 
             display_text = ""  # what we show to UI for this step
             next_id: Optional[str] = None
 
             # --- EXECUTION ---
             if stream and not multi_output:
-                # Full token streaming (single-output agent)
+                # Full token streaming (single-output agent) – collect full buffer for baton
                 result = Runner.run_streamed(agent, **run_kwargs)
                 handler.reset()
+                # Optional local accumulator; prefer handler.buffer after loop
+                last_chunk = ""
 
                 async for event in result.stream_events():
                     if bridge.stopped():
                         result.cancel()
                         bridge.on_stop(ctx)
                         break
-                    display_text, last_response_id = handler.handle(event, ctx)
-
-                # Prepare next inputs from result for memory update (if any)
-                input_items_next = result.to_input_list()
-                input_items_next = sanitize_input_items(input_items_next)
-
-                # Update memory only if attached
-                if mem_state:
-                    mem_state.update_from_result(input_items_next, last_response_id)
+                    chunk, last_response_id = handler.handle(event, ctx)
+                    if chunk:
+                        last_chunk = chunk
+
+                # Use full buffer if available (ensures baton sees complete output)
+                display_text = getattr(handler, "buffer", "") or last_chunk or ""
+
+                # Update memory strictly with baton + displayed text
+                self._update_memory_after_step(
+                    node_id=current_id,
+                    mem_state=mem_state,
+                    baton_user_text=baton_user_text,
+                    display_text=display_text,
+                    last_response_id=last_response_id,
+                    dbg=dbg,
+                )
 
                 # Route: first edge or END
                 outs = g.get_next(current_id)
@@ -308,13 +412,15 @@ class FlowOrchestrator:
                         decision = parse_route_output(raw_text, allowed_routes)
                         display_text = decision.content or ""
 
-                        # Prepare next inputs from streamed result, patch assistant content -> content
-                        input_items_next = result.to_input_list()
-                        input_items_next = patch_last_assistant_output(input_items_next, decision.content or "")
-
-                        # Update memory if attached
-                        if mem_state:
-                            mem_state.update_from_result(input_items_next, last_response_id)
+                        # Update memory with baton + displayed content
+                        self._update_memory_after_step(
+                            node_id=current_id,
+                            mem_state=mem_state,
+                            baton_user_text=baton_user_text,
+                            display_text=display_text,
+                            last_response_id=last_response_id,
+                            dbg=dbg,
+                        )
 
                         # Route decision
                         if decision.valid:
@@ -348,11 +454,15 @@ class FlowOrchestrator:
                             if not use_partial_ctx:
                                 handler.to_buffer(display_text)
 
-                        input_items_next = result.to_input_list()
-                        input_items_next = patch_last_assistant_output(input_items_next, display_text)
-
-                        if mem_state:
-                            mem_state.update_from_result(input_items_next, last_response_id)
+                        # Update memory with baton + displayed content
+                        self._update_memory_after_step(
+                            node_id=current_id,
+                            mem_state=mem_state,
+                            baton_user_text=baton_user_text,
+                            display_text=display_text,
+                            last_response_id=last_response_id,
+                            dbg=dbg,
+                        )
 
                         if decision.valid:
                             next_id = decision.route
@@ -373,11 +483,15 @@ class FlowOrchestrator:
                             if not use_partial_ctx:
                                 handler.to_buffer(display_text)
 
-                        input_items_next = result.to_input_list()
-                        input_items_next = patch_last_assistant_output(input_items_next, display_text)
-
-                        if mem_state:
-                            mem_state.update_from_result(input_items_next, last_response_id)
+                        # Update memory with baton + displayed content
+                        self._update_memory_after_step(
+                            node_id=current_id,
+                            mem_state=mem_state,
+                            baton_user_text=baton_user_text,
+                            display_text=display_text,
+                            last_response_id=last_response_id,
+                            dbg=dbg,
+                        )
 
                         if decision.valid:
                             next_id = decision.route
@@ -397,11 +511,15 @@ class FlowOrchestrator:
                         if not use_partial_ctx:
                             handler.to_buffer(display_text)
 
-                    input_items_next = result.to_input_list()
-                    input_items_next = sanitize_input_items(input_items_next)
-
-                    if mem_state:
-                        mem_state.update_from_result(input_items_next, last_response_id)
+                    # Update memory with baton + displayed text
+                    self._update_memory_after_step(
+                        node_id=current_id,
+                        mem_state=mem_state,
+                        baton_user_text=baton_user_text,
+                        display_text=display_text,
+                        last_response_id=last_response_id,
+                        dbg=dbg,
+                    )
 
                     outs = g.get_next(current_id)
                     next_id = outs[0] if outs else g.first_connected_end(current_id)