From a53f3c2586edf6b169627ebc675fae2292fd2b16 Mon Sep 17 00:00:00 2001
From: Liam Bigelow <40188355+bglw@users.noreply.github.com>
Date: Mon, 30 Sep 2024 11:51:20 +1300
Subject: [PATCH] Python API: Improve error handling for messages that failed
 to parse

---
 .../py-pagefind-error-handling.toolproof.yml  | 38 +++++++++++++++++++
 .../python/src/pagefind/service/__init__.py   |  9 ++++-
 2 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 pagefind/integration_tests/python_api/py-pagefind-error-handling.toolproof.yml

diff --git a/pagefind/integration_tests/python_api/py-pagefind-error-handling.toolproof.yml b/pagefind/integration_tests/python_api/py-pagefind-error-handling.toolproof.yml
new file mode 100644
index 00000000..49c11d48
--- /dev/null
+++ b/pagefind/integration_tests/python_api/py-pagefind-error-handling.toolproof.yml
@@ -0,0 +1,38 @@
+name: Python API > Pagefind error handling
+platforms:
+  - linux
+  - mac
+
+steps:
+  - ref: ./background.toolproof.yml
+  - step: I have a "public/run.py" file with the content {python}
+    python: |2-
+      import sys
+      sys.path.append('%repo_wd%/wrappers/python/src')
+
+      import asyncio
+      import json
+      import logging
+      import os
+      from pagefind.index import PagefindIndex, IndexConfig
+
+      async def main():
+        async with PagefindIndex() as index:
+          await index.delete_index()
+          try:
+            files = await index.get_files()
+          except AssertionError:
+            print("Index deleted")
+
+        try:
+          async with PagefindIndex(IndexConfig(root_selector=5)) as index:
+            await index.delete_index()
+        except Exception as e:
+          print(f"Caught error {e}")
+
+        print("Complete")
+      if __name__ == "__main__":
+        asyncio.run(main())
+  - step: I run "cd public && PAGEFIND_BINARY_PATH=%pagefind_exec_path% python3 run.py"
+  - step: 'stdout should contain "invalid type: integer `5`"'
+  - step: stdout should contain "Index deleted"
diff --git a/wrappers/python/src/pagefind/service/__init__.py b/wrappers/python/src/pagefind/service/__init__.py
index 5a0239b5..dba1864b 100644
--- a/wrappers/python/src/pagefind/service/__init__.py
+++ b/wrappers/python/src/pagefind/service/__init__.py
@@ -160,7 +160,14 @@ async def _wait_for_responses(self) -> None:
             if (resp := json.loads(base64.b64decode(output[:-1]))) is None:
                 continue
             resp = cast(InternalServiceResponse, resp)
-            if (message_id := resp.get("message_id")) is not None:
+            message_id = resp.get("message_id")
+            if message_id is None:
+                # If the backend service failed to parse the message, it won't return the ID
+                # However it does return the message itself, so we can retrieve the ID we sent
+                if (orginal := resp["payload"].get("original_message")) is not None:
+                    if (sent := json.loads(orginal)) is not None:
+                        message_id = sent.get("message_id")
+            if message_id is not None:
                 log.debug(f"received response for message {message_id}")
                 assert (
                     self._message_id >= message_id