From 5f31a6b12f839e10133c3463e0b2aca263a5ba5f Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 26 May 2026 13:07:21 +0300
Subject: [PATCH] Fix round-robin docling-serve

---
 .../haiku/rag/providers/docling_serve.py      | 30 ++++++++++++-------
 tests/test_docling_serve_client.py            | 22 +++++++++++++-
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/providers/docling_serve.py b/haiku_rag_slim/haiku/rag/providers/docling_serve.py
index 29b53879..9d7b7930 100644
--- a/haiku_rag_slim/haiku/rag/providers/docling_serve.py
+++ b/haiku_rag_slim/haiku/rag/providers/docling_serve.py
@@ -1,21 +1,28 @@
 """Shared client for docling-serve async API."""
 
 import asyncio
+import itertools
 from typing import Any
 
 import httpx
 
+# Process-global registry of round-robin rotators over docling-serve
+# instances, keyed by the sorted tuple of base URLs. Clients are constructed
+# per-job by `get_converter` / `get_chunker`, so the rotation cursor has to
+# live OUTSIDE the instance to actually advance across jobs. Two clients
+# pointing at the same instance set share one rotator; clients with
+# different sets get independent rotators.
+_instance_rotators: dict[tuple[str, ...], "itertools.cycle[str]"] = {}
+
 
 class DoclingServeClient:
     """Client for docling-serve async workflow.
 
-    Handles the submit → poll → fetch pattern used by both conversion and chunking.
-
-    Accepts a list of base URLs and round-robins jobs across them. Each job's
+    Handles the submit → poll → fetch pattern used by both conversion and
+    chunking. Accepts a list of base URLs and round-robins jobs across
+    them via a process-wide rotator keyed by the URL set. Each job's
     submit/poll/result trio stays on the same instance — task IDs are
-    instance-local, so picking a different URL mid-job would 404. The counter
-    is per-process; concurrent processes pick independently, so over many
-    jobs the distribution is even but not coordinated.
+    instance-local, so picking a different URL mid-job would 404.
     """
 
     def __init__(
@@ -33,7 +40,12 @@ class DoclingServeClient:
         self.timeout = timeout
         # transport is for testing — production callers leave it None.
         self._transport = transport
-        self._counter = 0
+        # setdefault is atomic under the GIL — concurrent constructors with
+        # the same instance set will end up sharing one rotator.
+        key = tuple(sorted(self.base_urls))
+        self._instance_rotator = _instance_rotators.setdefault(
+            key, itertools.cycle(self.base_urls)
+        )
 
     def _httpx_client(self) -> httpx.AsyncClient:
         return httpx.AsyncClient(timeout=self.timeout, transport=self._transport)
@@ -45,9 +57,7 @@ class DoclingServeClient:
         return self.base_urls[0]
 
     def _pick_url(self) -> str:
-        url = self.base_urls[self._counter % len(self.base_urls)]
-        self._counter += 1
-        return url
+        return next(self._instance_rotator)
 
     def _get_headers(self) -> dict[str, str]:
         """Get headers for API requests."""
diff --git a/tests/test_docling_serve_client.py b/tests/test_docling_serve_client.py
index de40d705..b19360d4 100644
--- a/tests/test_docling_serve_client.py
+++ b/tests/test_docling_serve_client.py
@@ -1,4 +1,10 @@
-"""DoclingServeClient round-robin distribution tests."""
+"""DoclingServeClient round-robin distribution tests.
+
+Each test here uses a unique base-URL set so the process-global cycle
+map gives it a fresh itertools.cycle. Don't reuse URL strings across
+tests — cycles persist for the lifetime of the process and would
+resume mid-rotation, breaking specific-order assertions.
+"""
 
 import httpx
 import pytest
@@ -90,6 +96,20 @@ async def test_task_lifecycle_pinned_to_same_url():
     assert len(set(seen)) == 1
 
 
+def test_round_robin_shared_across_fresh_clients():
+    """get_converter / get_chunker build a NEW DoclingServeClient per job.
+    The cycle has to live outside the instance so successive jobs (each
+    with its own freshly-constructed client) actually rotate."""
+    urls = ["http://x:5001", "http://y:5001", "http://z:5001"]
+    c1 = DoclingServeClient(base_urls=urls)
+    c2 = DoclingServeClient(base_urls=urls)
+    c3 = DoclingServeClient(base_urls=urls)
+    c4 = DoclingServeClient(base_urls=urls)
+
+    picks = [c1._pick_url(), c2._pick_url(), c3._pick_url(), c4._pick_url()]
+    assert picks == [urls[0], urls[1], urls[2], urls[0]]
+
+
 @pytest.mark.asyncio
 async def test_zip_endpoint_uses_round_robin_too():
     transport, seen = _scripted_transport(