fix: persist session after HITL resume and emit tool completed events

- After HITL resume, checkpoint restore creates separate session copies
  per executor. Extract the most complete session (highest message count)
  and persist it to KV storage so the next turn has valid history.
- Handle AgentExecutorResponse wrapper in _extract_tool_state_events and
  _extract_contents so function_result from executor_completed data is
  properly found.
- Emit ToolCallEnd in close_message() for pending tool calls interrupted
  by HITL suspension (clears stale _pending_tool_calls state).
- Track pending tool nodes (STARTED without COMPLETED) across stream
  iterations and synthesize COMPLETED events on HITL resume.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: cristipufu

packages/uipath-agent-framework/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-agent-framework"
-version = "0.0.6"
+version = "0.0.7"
 description = "Python SDK that enables developers to build and deploy Microsoft Agent Framework agents to the UiPath Cloud Platform"
 readme = "README.md"
 requires-python = ">=3.11"

packages/uipath-agent-framework/samples/hitl-workflow/main.py

@@ -54,7 +54,7 @@ def issue_refund(order_id: str, amount: float, reason: str) -> str:
     return f"Refund of ${amount:.2f} issued for order {order_id}: {reason}"
 
 
-client = UiPathOpenAIChatClient(model="gpt-5-mini-2025-08-07")
+client = UiPathOpenAIChatClient()
 
 triage = client.as_agent(
     name="triage",

packages/uipath-agent-framework/samples/hitl-workflow/pyproject.toml

@@ -19,3 +19,6 @@ dev = [
 [tool.uv]
 prerelease = "allow"
 
+[tool.uv.sources]
+uipath-dev = { path = "../../../../../uipath-dev-python", editable = true }
+uipath-agent-framework = { path = "../../", editable = true }

packages/uipath-agent-framework/src/uipath_agent_framework/chat/openai.py

@@ -67,15 +67,15 @@ class UiPathOpenAIChatClient(OpenAIChatClient):
 
         from uipath_agent_framework.chat import UiPathOpenAIChatClient
 
-        client = UiPathOpenAIChatClient(model="gpt-4o-mini")
+        client = UiPathOpenAIChatClient(model="gpt-4.1-mini-2025-04-14")
         agent = client.as_agent(
             name="assistant",
             instructions="You are a helpful assistant.",
             tools=[my_tool],
         )
     """
 
-    def __init__(self, model: str = "gpt-4o-mini", **kwargs: Any):
+    def __init__(self, model: str = "gpt-4.1-mini-2025-04-14", **kwargs: Any):
         uipath_url, token = get_uipath_config()
         gateway_url = build_gateway_url("openai", model, uipath_url)
 

packages/uipath-agent-framework/src/uipath_agent_framework/runtime/messages.py

@@ -184,12 +184,20 @@ def map_streaming_content(
 
     def close_message(self) -> list[UiPathConversationMessageEvent]:
         """Close the current message if open. Safety net for end of stream."""
+        events: list[UiPathConversationMessageEvent] = []
+        # Emit ToolCallEnd for any tool calls that were started but never
+        # completed (e.g. HITL suspension interrupted before function_result).
+        if self._pending_tool_calls:
+            for tool_call_id, message_id in self._pending_tool_calls.items():
+                events.append(
+                    self._make_tool_call_end_event(message_id, tool_call_id, {})
+                )
+            self._pending_tool_calls.clear()
         if self._message_started and self._current_message_id:
-            events = [self._make_message_end_event(self._current_message_id)]
+            events.append(self._make_message_end_event(self._current_message_id))
             self._message_started = False
             self._current_message_id = None
-            return events
-        return []
+        return events
 
     @staticmethod
     def _extract_text_from_content(content: Content) -> str:

packages/uipath-agent-framework/src/uipath_agent_framework/runtime/runtime.py

@@ -6,6 +6,7 @@
 
 from agent_framework import (
     AgentExecutor,
+    AgentExecutorResponse,
     AgentResponse,
     AgentResponseUpdate,
     AgentSession,
@@ -67,6 +68,12 @@ def __init__(
         self._last_breakpoint_node: str | None = None
         self._last_checkpoint_id: str | None = None
         self._resumed_from_checkpoint_id: str | None = None
+        # Track tool nodes that emitted STARTED but not yet COMPLETED.
+        # Persists across _stream_workflow() calls (same runtime instance
+        # reused by UiPathChatRuntime's while loop), allowing us to emit
+        # synthetic COMPLETED events on HITL resume when the framework
+        # doesn't surface function_result in output/executor_completed.
+        self._pending_tool_nodes: set[str] = set()
 
     # ------------------------------------------------------------------
     # Checkpoint helpers
@@ -171,6 +178,37 @@ def _apply_session_to_executors(self, session: AgentSession) -> None:
             if isinstance(executor, AgentExecutor):
                 executor._session = session
 
+    def _get_session_from_executors(self) -> AgentSession | None:
+        """Extract the most complete session from AgentExecutors in the workflow.
+
+        After checkpoint restore each executor receives its own independent
+        session copy (unlike fresh runs where all executors share one object).
+        Only the executor that processed the HITL/breakpoint response will
+        have the updated conversation history. We return the session with the
+        most messages to ensure the complete history is persisted.
+        """
+        workflow = self.agent.workflow
+        best_session: AgentSession | None = None
+        best_msg_count = -1
+        for executor in workflow.executors.values():
+            if isinstance(executor, AgentExecutor) and executor._session is not None:
+                msg_count = self._count_session_messages(executor._session)
+                if msg_count > best_msg_count:
+                    best_msg_count = msg_count
+                    best_session = executor._session
+        return best_session
+
+    @staticmethod
+    def _count_session_messages(session: AgentSession) -> int:
+        """Count total messages across all provider keys in a session's state."""
+        count = 0
+        for value in session.state.values():
+            if isinstance(value, dict) and "messages" in value:
+                messages = value["messages"]
+                if isinstance(messages, list):
+                    count += len(messages)
+        return count
+
     # ------------------------------------------------------------------
     # HITL helpers (tool approval flow)
     # ------------------------------------------------------------------
@@ -332,6 +370,12 @@ async def execute(
                     checkpoint_storage=self._checkpoint_storage,
                 )
 
+            # After resume paths the checkpoint restores the session into
+            # executors directly, so the local ``session`` is still None.
+            # Extract it so it can be persisted after completion.
+            if session is None:
+                session = self._get_session_from_executors()
+
             # Check for HITL suspension (framework's request_info mechanism)
             request_info_events = result.get_request_info_events()
             hitl_requests = {
@@ -462,6 +506,19 @@ async def _stream_workflow(
             phase=UiPathRuntimeStatePhase.STARTED,
         )
 
+        # On HITL resume, emit COMPLETED for tool nodes that were left
+        # pending when the previous stream suspended. The framework
+        # doesn't surface function_result in output/executor_completed
+        # for handoff workflows, so we synthesize these events here.
+        if is_resuming and self._pending_tool_nodes:
+            for tool_node in list(self._pending_tool_nodes):
+                yield UiPathRuntimeStateEvent(
+                    payload={},
+                    node_name=tool_node,
+                    phase=UiPathRuntimeStatePhase.COMPLETED,
+                )
+            self._pending_tool_nodes.clear()
+
         # Choose workflow.run() mode based on resume type
         if self._resume_responses:
             # HITL resume: pass responses to workflow with checkpoint
@@ -495,10 +552,13 @@ async def _stream_workflow(
 
         request_info_map: dict[str, Any] = {}
         is_suspended = False
-        # Track executors whose tool events were emitted via output events.
-        # When the workflow filters output events (e.g. GroupChat), tool events
-        # are instead extracted from executor_completed data as a fallback.
-        executors_with_tool_outputs: set[str] = set()
+        # Track which tool event phases were emitted per executor via output
+        # events. When the workflow filters output events (e.g. GroupChat),
+        # tool events are extracted from executor_completed data as a fallback.
+        # Tracking phases (not just executor_ids) lets us handle HITL resume
+        # where function_call (STARTED) is in output but function_result
+        # (COMPLETED) is only in executor_completed.
+        executor_tool_phases: dict[str, set[UiPathRuntimeStatePhase]] = {}
 
         # Emit an early STARTED event for the start executor so the graph
         # visualization shows it immediately rather than after it finishes.
@@ -534,17 +594,36 @@ async def _stream_workflow(
                         phase=UiPathRuntimeStatePhase.STARTED,
                     )
                 elif event.type == "executor_completed":
-                    # When output events were filtered by the workflow (e.g.
-                    # GroupChat where participants are not output executors),
-                    # extract tool state events from the completed data instead.
-                    if (
-                        event.executor_id
-                        and event.executor_id not in executors_with_tool_outputs
-                    ):
+                    # Extract tool state events from executor_completed data,
+                    # skipping phases already emitted via output events.
+                    # This handles three scenarios:
+                    # 1. GroupChat (no output events): emit all from completed
+                    # 2. Normal (both in output): skip all from completed
+                    # 3. HITL resume (only STARTED in output): emit COMPLETED
+                    if event.executor_id:
+                        emitted_phases = executor_tool_phases.get(
+                            event.executor_id, set()
+                        )
                         for tool_event in self._extract_tool_state_events(
                             event.data, event.executor_id
                         ):
-                            yield tool_event
+                            if tool_event.phase not in emitted_phases:
+                                # Track pending tool nodes
+                                if (
+                                    tool_event.phase
+                                    == UiPathRuntimeStatePhase.STARTED
+                                ):
+                                    self._pending_tool_nodes.add(
+                                        tool_event.node_name
+                                    )
+                                elif (
+                                    tool_event.phase
+                                    == UiPathRuntimeStatePhase.COMPLETED
+                                ):
+                                    self._pending_tool_nodes.discard(
+                                        tool_event.node_name
+                                    )
+                                yield tool_event
                     yield UiPathRuntimeStateEvent(
                         payload=self._serialize_event_data(
                             self._filter_completed_data(event.data)
@@ -557,9 +636,15 @@ async def _stream_workflow(
                     tool_events = self._extract_tool_state_events(
                         event.data, executor_id
                     )
-                    if tool_events:
-                        executors_with_tool_outputs.add(executor_id)
                     for tool_event in tool_events:
+                        executor_tool_phases.setdefault(
+                            executor_id, set()
+                        ).add(tool_event.phase)
+                        # Track pending tool nodes across stream iterations
+                        if tool_event.phase == UiPathRuntimeStatePhase.STARTED:
+                            self._pending_tool_nodes.add(tool_event.node_name)
+                        elif tool_event.phase == UiPathRuntimeStatePhase.COMPLETED:
+                            self._pending_tool_nodes.discard(tool_event.node_name)
                         yield tool_event
                     for msg_event in self._extract_workflow_messages(event.data):
                         yield UiPathRuntimeMessageEvent(payload=msg_event)
@@ -581,6 +666,10 @@ async def _stream_workflow(
             for msg_event in self.chat.close_message():
                 yield UiPathRuntimeMessageEvent(payload=msg_event)
 
+            # After resume paths the checkpoint restores the session into
+            # executors directly, so the local ``session`` may still be None.
+            if session is None:
+                session = self._get_session_from_executors()
             if session is not None:
                 await self._save_session(session)
 
@@ -619,6 +708,10 @@ async def _stream_workflow(
         for msg_event in self.chat.close_message():
             yield UiPathRuntimeMessageEvent(payload=msg_event)
 
+        # After resume paths the checkpoint restores the session into
+        # executors directly, so the local ``session`` may still be None.
+        if session is None:
+            session = self._get_session_from_executors()
         if session is not None:
             await self._save_session(session)
 
@@ -681,7 +774,11 @@ def _extract_tool_state_events(
         """
         contents: list[Any] = []
 
-        if isinstance(data, AgentResponseUpdate):
+        if isinstance(data, AgentExecutorResponse):
+            return UiPathAgentFrameworkRuntime._extract_tool_state_events(
+                data.agent_response, executor_id
+            )
+        elif isinstance(data, AgentResponseUpdate):
             contents = list(data.contents or [])
         elif isinstance(data, AgentResponse):
             for message in data.messages or []:
@@ -724,7 +821,9 @@ def _extract_tool_state_events(
     def _extract_contents(data: Any) -> list[Any]:
         """Extract Content objects from any workflow data type."""
         contents: list[Any] = []
-        if isinstance(data, AgentResponseUpdate):
+        if isinstance(data, AgentExecutorResponse):
+            return UiPathAgentFrameworkRuntime._extract_contents(data.agent_response)
+        elif isinstance(data, AgentResponseUpdate):
             contents = list(data.contents or [])
         elif isinstance(data, AgentResponse):
             for message in data.messages or []: