fix: preserve async tool result order

private_gpt/components/engines/chat/resumable_runner.py

@@ -34,6 +34,7 @@ if TYPE_CHECKING:
 
     from private_gpt.components.engines.chat.async_chat_engine import AsyncChatEngine
     from private_gpt.components.engines.chat.models.chat_state import ChatState
+    from private_gpt.components.tools.remote_execution import ToolExecutionResponse
     from private_gpt.events.models import Event
 
 
@@ -135,7 +136,8 @@ class ResumableChatRunner:
             return
         channel = BrokerEventChannel(self._events, execution_id)
         try:
-            responses = list((await self._state.get_results(execution_id)).values())
+            results = await self._state.get_results(execution_id)
+            responses = self._ordered_results(saved, results)
             request_data = dict(saved.request_data)
             request_data["messages"] = [
                 *list(request_data.get("messages", [])),
@@ -250,6 +252,17 @@ class ResumableChatRunner:
         await self._state.cleanup(execution_id)
         await self._events.finish(execution_id)
 
+    @staticmethod
+    def _ordered_results(
+        checkpoint: ChatCheckpoint,
+        results: dict[str, ToolExecutionResponse],
+    ) -> list[ToolExecutionResponse]:
+        return [
+            results[tool_id]
+            for tool_id in checkpoint.checkpoint_payload.pending_async_tools
+            if tool_id in results
+        ]
+
     @staticmethod
     def _request(request_data: dict[str, Any]) -> ResolvedChatRequest:
         request = ResolvedChatRequest.model_validate(request_data)

tests/engines/test_resumable_runtime.py

@@ -264,6 +264,48 @@ def test_runner_restores_additional_kwargs_after_serialization() -> None:
     assert isinstance(additional_kwargs["tool_calls"][0], ToolSelection)
 
 
+def test_runner_orders_results_by_pending_tool_order() -> None:
+    from private_gpt.components.engines.chat.resumable_runner import ResumableChatRunner
+
+    results = {
+        "tool-2": ToolExecutionResponse(
+            tool_name="second",
+            tool_id="tool-2",
+            result_content=[TextBlock(text="second")],
+            tool_message={
+                "role": "tool",
+                "content": "second",
+                "additional_kwargs": {"tool_call_id": "tool-2"},
+            },
+        ),
+        "tool-1": ToolExecutionResponse(
+            tool_name="first",
+            tool_id="tool-1",
+            result_content=[TextBlock(text="first")],
+            tool_message={
+                "role": "tool",
+                "content": "first",
+                "additional_kwargs": {"tool_call_id": "tool-1"},
+            },
+        ),
+    }
+    checkpoint = ChatCheckpoint(
+        correlation_id="execution-order",
+        request_data={},
+        stream_type="chat_completion",
+        metadata={},
+        iteration=0,
+        checkpoint="tools",
+        checkpoint_payload=IterationCheckpointPayload(
+            pending_async_tools={"tool-1": "job-1", "tool-2": "job-2"}
+        ),
+    )
+
+    ordered = ResumableChatRunner._ordered_results(checkpoint, results)
+
+    assert [response.tool_id for response in ordered] == ["tool-1", "tool-2"]
+
+
 def test_resolved_chat_request_is_json_roundtrip_serializable() -> None:
     import json