chore: import upstream snapshot with attribution

tests/test_embedding_server_manager.py

@@ -0,0 +1,648 @@
+import json
+import os
+import threading
+import time
+from typing import Any, cast
+
+import pytest
+from leann.embedding_server_manager import EmbeddingServerManager
+
+
+class DummyProcess:
+    def __init__(self, pid=12345):
+        self.pid = pid
+        self._terminated = False
+
+    def poll(self):
+        return 0 if self._terminated else None
+
+    def terminate(self):
+        self._terminated = True
+
+    def kill(self):
+        self._terminated = True
+
+    def wait(self, timeout=None):
+        self._terminated = True
+        return 0
+
+
+@pytest.fixture
+def embedding_manager(monkeypatch):
+    manager = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+
+    def fake_get_available_port(start_port):
+        return start_port
+
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._get_available_port",
+        fake_get_available_port,
+    )
+
+    start_calls = []
+
+    def fake_start_new_server(self, port, model_name, embedding_mode, **kwargs):
+        config_signature = kwargs.get("config_signature")
+        start_calls.append(config_signature)
+        self.server_process = DummyProcess()
+        self.server_port = port
+        self._server_config = config_signature
+        return True, port
+
+    monkeypatch.setattr(
+        EmbeddingServerManager,
+        "_start_new_server",
+        fake_start_new_server,
+    )
+
+    # Ensure stop_server doesn't try to operate on real subprocesses
+    def fake_stop_server(self):
+        self.server_process = None
+        self.server_port = None
+        self._server_config = None
+
+    monkeypatch.setattr(EmbeddingServerManager, "stop_server", fake_stop_server)
+
+    return manager, start_calls
+
+
+def _write_meta(meta_path, passages_name, index_name, total):
+    meta_path.write_text(
+        json.dumps(
+            {
+                "backend_name": "hnsw",
+                "embedding_model": "test-model",
+                "embedding_mode": "sentence-transformers",
+                "dimensions": 3,
+                "backend_kwargs": {},
+                "passage_sources": [
+                    {
+                        "type": "jsonl",
+                        "path": passages_name,
+                        "index_path": index_name,
+                    }
+                ],
+                "total_passages": total,
+            }
+        ),
+        encoding="utf-8",
+    )
+
+
+def test_server_restarts_when_metadata_changes(tmp_path, embedding_manager):
+    manager, start_calls = embedding_manager
+
+    meta_path = tmp_path / "example.meta.json"
+    passages_path = tmp_path / "example.passages.jsonl"
+    index_path = tmp_path / "example.passages.idx"
+
+    passages_path.write_text("first\n", encoding="utf-8")
+    index_path.write_bytes(b"index")
+    _write_meta(meta_path, passages_path.name, index_path.name, total=1)
+
+    # Initial start populates signature
+    ok, port = manager.start_server(
+        port=6000,
+        model_name="test-model",
+        passages_file=str(meta_path),
+        use_daemon=False,
+    )
+    assert ok
+    assert port == 6000
+    assert len(start_calls) == 1
+
+    initial_signature = start_calls[0]["passages_signature"]
+
+    # No metadata change => reuse existing server
+    ok, port_again = manager.start_server(
+        port=6000,
+        model_name="test-model",
+        passages_file=str(meta_path),
+        use_daemon=False,
+    )
+    assert ok
+    assert port_again == 6000
+    assert len(start_calls) == 1
+
+    # Modify passage data and metadata to force signature change
+    time.sleep(0.01)  # Ensure filesystem timestamps move forward
+    passages_path.write_text("second\n", encoding="utf-8")
+    _write_meta(meta_path, passages_path.name, index_path.name, total=2)
+
+    ok, port_third = manager.start_server(
+        port=6000,
+        model_name="test-model",
+        passages_file=str(meta_path),
+        use_daemon=False,
+    )
+    assert ok
+    assert port_third == 6000
+    assert len(start_calls) == 2
+
+    updated_signature = start_calls[1]["passages_signature"]
+    assert updated_signature != initial_signature
+
+
+def test_list_daemons_ignores_stale_records(tmp_path, monkeypatch):
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+
+    stale = registry_dir / "stale.json"
+    stale.write_text(
+        json.dumps(
+            {
+                "pid": 999999,
+                "port": 65531,
+                "backend_module_name": "leann_backend_hnsw.hnsw_embedding_server",
+                "config_signature": {"passages_file": "/tmp/a.meta.json"},
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    records = EmbeddingServerManager.list_daemons()
+    assert records == []
+    assert not stale.exists()
+
+
+def test_stop_daemons_filters_by_backend_and_passages(tmp_path, monkeypatch):
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+
+    meta_path = (tmp_path / "x.meta.json").resolve()
+    record = registry_dir / "daemon.json"
+    record.write_text(
+        json.dumps(
+            {
+                "pid": 12345,
+                "port": 6001,
+                "backend_module_name": "leann_backend_hnsw.hnsw_embedding_server",
+                "config_signature": {"passages_file": str(meta_path)},
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    monkeypatch.setattr(
+        EmbeddingServerManager,
+        "list_daemons",
+        classmethod(
+            lambda cls: [
+                {
+                    "pid": 12345,
+                    "port": 6001,
+                    "backend_module_name": "leann_backend_hnsw.hnsw_embedding_server",
+                    "config_signature": {"passages_file": str(meta_path)},
+                    "record_path": str(record),
+                }
+            ]
+        ),
+    )
+
+    killed: list[tuple[int, int]] = []
+
+    def fake_kill(pid: int, sig: int):
+        killed.append((pid, sig))
+
+    monkeypatch.setattr(os, "kill", fake_kill)
+
+    stopped = EmbeddingServerManager.stop_daemons(
+        backend_module_name="leann_backend_hnsw.hnsw_embedding_server",
+        passages_file=str(meta_path),
+    )
+    assert stopped == 1
+    assert killed == [(12345, 15)]
+    assert not record.exists()
+
+
+def test_daemon_registry_reuse_across_manager_instances(tmp_path, monkeypatch):
+    """Second manager should adopt the daemon started by first manager."""
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._get_available_port",
+        lambda start_port: start_port,
+    )
+    monkeypatch.setattr("leann.embedding_server_manager._check_port", lambda port: True)
+    monkeypatch.setattr("leann.embedding_server_manager._pid_is_alive", lambda pid: pid == 22222)
+
+    starts = []
+
+    def fake_start_new_server(self, port, model_name, embedding_mode, **kwargs):
+        starts.append((port, model_name, embedding_mode))
+        self.server_process = DummyProcess(pid=22222)
+        self.server_port = port
+        self._server_config = kwargs.get("config_signature")
+        return True, port
+
+    monkeypatch.setattr(EmbeddingServerManager, "_start_new_server", fake_start_new_server)
+
+    manager1 = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    ok1, port1 = manager1.start_server(
+        port=6011,
+        model_name="test-model",
+        use_daemon=True,
+        daemon_ttl_seconds=120,
+    )
+    assert ok1 and port1 == 6011
+    assert len(starts) == 1
+
+    manager2 = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    ok2, port2 = manager2.start_server(
+        port=6011,
+        model_name="test-model",
+        use_daemon=True,
+        daemon_ttl_seconds=120,
+    )
+    assert ok2 and port2 == 6011
+    # No second process spawn: adopted from registry.
+    assert len(starts) == 1
+    assert manager2.server_process is None
+
+
+def test_stale_registry_falls_back_to_fresh_start(tmp_path, monkeypatch):
+    """If registry points to dead daemon, manager should start a new process."""
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._get_available_port",
+        lambda start_port: start_port,
+    )
+
+    starts = []
+
+    def fake_start_new_server(self, port, model_name, embedding_mode, **kwargs):
+        starts.append(port)
+        self.server_process = DummyProcess(pid=33333)
+        self.server_port = port
+        self._server_config = kwargs.get("config_signature")
+        return True, port
+
+    monkeypatch.setattr(EmbeddingServerManager, "_start_new_server", fake_start_new_server)
+
+    manager = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    signature = manager._build_config_signature(
+        model_name="test-model",
+        embedding_mode="sentence-transformers",
+        provider_options=None,
+        passages_file=None,
+        distance_metric=None,
+    )
+    stale_file = registry_dir / f"{manager._registry_key(signature)}.json"
+    stale_file.write_text(
+        json.dumps(
+            {
+                "pid": 999999,
+                "port": 6012,
+                "backend_module_name": "leann_backend_hnsw.hnsw_embedding_server",
+                "config_signature": signature,
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    monkeypatch.setattr("leann.embedding_server_manager._pid_is_alive", lambda pid: False)
+    monkeypatch.setattr("leann.embedding_server_manager._check_port", lambda port: False)
+
+    ok, port = manager.start_server(
+        port=6012,
+        model_name="test-model",
+        use_daemon=True,
+    )
+    assert ok and port == 6012
+    assert starts == [6012]
+    assert stale_file.exists()
+    refreshed = json.loads(stale_file.read_text(encoding="utf-8"))
+    assert refreshed["pid"] == 33333
+
+
+def test_build_server_command_includes_daemon_and_warmup_flags():
+    manager = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    command = manager._build_server_command(
+        port=6020,
+        model_name="m",
+        embedding_mode="sentence-transformers",
+        distance_metric="mips",
+        enable_warmup=True,
+        use_daemon=True,
+        daemon_ttl_seconds=321,
+    )
+    assert "--enable-warmup" in command
+    assert "--daemon-mode" in command
+    assert "--daemon-ttl" in command
+    ttl_idx = command.index("--daemon-ttl")
+    assert command[ttl_idx + 1] == "321"
+
+    command_no_daemon = manager._build_server_command(
+        port=6020,
+        model_name="m",
+        embedding_mode="sentence-transformers",
+        enable_warmup=False,
+        use_daemon=False,
+    )
+    assert "--daemon-mode" not in command_no_daemon
+    assert "--enable-warmup" not in command_no_daemon
+
+
+def test_corrupted_registry_file_is_recovered_on_start(tmp_path, monkeypatch):
+    """Invalid registry json should not block startup; file is replaced."""
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._get_available_port",
+        lambda start_port: start_port,
+    )
+
+    starts = []
+
+    def fake_start_new_server(self, port, model_name, embedding_mode, **kwargs):
+        starts.append(port)
+        self.server_process = DummyProcess(pid=44444)
+        self.server_port = port
+        self._server_config = kwargs.get("config_signature")
+        return True, port
+
+    monkeypatch.setattr(EmbeddingServerManager, "_start_new_server", fake_start_new_server)
+
+    manager = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    signature = manager._build_config_signature(
+        model_name="test-model",
+        embedding_mode="sentence-transformers",
+        provider_options=None,
+        passages_file=None,
+        distance_metric=None,
+    )
+    record_path = registry_dir / f"{manager._registry_key(signature)}.json"
+    record_path.write_text("{invalid-json", encoding="utf-8")
+
+    monkeypatch.setattr("leann.embedding_server_manager._pid_is_alive", lambda pid: False)
+    monkeypatch.setattr("leann.embedding_server_manager._check_port", lambda port: False)
+
+    ok, port = manager.start_server(port=6030, model_name="test-model", use_daemon=True)
+    assert ok and port == 6030
+    assert starts == [6030]
+    data = json.loads(record_path.read_text(encoding="utf-8"))
+    assert data["pid"] == 44444
+
+
+def test_stop_server_detaches_when_daemon_mode(monkeypatch):
+    """Daemon mode should detach manager without terminating shared process."""
+    manager = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    manager.server_process = cast(Any, DummyProcess(pid=55555))
+    manager.server_port = 6031
+    manager._server_config = {"model_name": "m"}
+    manager._daemon_mode = True
+
+    # If terminate is called this test should fail.
+    called = {"terminate": 0}
+
+    def fail_terminate():
+        called["terminate"] += 1
+        raise AssertionError("terminate should not be called in daemon detach path")
+
+    manager.server_process.terminate = fail_terminate  # type: ignore[method-assign]
+
+    manager.stop_server()
+    assert called["terminate"] == 0
+    assert manager.server_process is None
+    assert manager.server_port is None
+    assert manager._server_config is None
+
+
+def test_concurrent_daemon_start_only_spawns_once(tmp_path, monkeypatch):
+    """Concurrent calls should be serialized by registry lock for same signature."""
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._get_available_port",
+        lambda start_port: start_port,
+    )
+    monkeypatch.setattr("leann.embedding_server_manager._check_port", lambda port: True)
+    monkeypatch.setattr("leann.embedding_server_manager._pid_is_alive", lambda pid: pid in (77777,))
+
+    starts = []
+
+    def fake_start_new_server(self, port, model_name, embedding_mode, **kwargs):
+        # Force overlap window between two starters.
+        time.sleep(0.05)
+        starts.append((self, port))
+        self.server_process = DummyProcess(pid=77777)
+        self.server_port = port
+        self._server_config = kwargs.get("config_signature")
+        return True, port
+
+    monkeypatch.setattr(EmbeddingServerManager, "_start_new_server", fake_start_new_server)
+
+    results = []
+
+    def runner():
+        manager = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+        ok, port = manager.start_server(port=6040, model_name="test-model", use_daemon=True)
+        results.append((ok, port))
+
+    t1 = threading.Thread(target=runner)
+    t2 = threading.Thread(target=runner)
+    t1.start()
+    t2.start()
+    t1.join(timeout=5)
+    t2.join(timeout=5)
+
+    assert len(results) == 2
+    assert all(ok and port == 6040 for ok, port in results)
+    # Exactly one actual process start, one adopts registry record.
+    assert len(starts) == 1
+
+
+def test_registry_record_write_is_atomic(tmp_path, monkeypatch):
+    """Registry writes should go through temp file + os.replace."""
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+
+    manager = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    manager.server_process = cast(Any, DummyProcess(pid=88888))
+
+    calls = []
+    real_replace = os.replace
+
+    def tracked_replace(src, dst):
+        calls.append((str(src), str(dst)))
+        return real_replace(src, dst)
+
+    monkeypatch.setattr(os, "replace", tracked_replace)
+
+    path = manager._write_registry_record(
+        port=6060,
+        config_signature={"model_name": "m"},
+        daemon_ttl_seconds=123,
+    )
+    assert path.exists()
+    assert calls, "os.replace should be used for atomic write"
+    src, dst = calls[0]
+    assert src.endswith(".json.tmp")
+    assert dst.endswith(".json")
+
+
+def test_different_passages_files_start_separate_daemons(tmp_path, monkeypatch):
+    """Two managers with different passages files must NOT share a daemon (issue #281).
+
+    Simulates sequential test runs where each test builds its own index in a
+    separate directory but uses the same model and embedding mode.  A stale
+    daemon from run A should never be adopted by run B.
+    """
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+
+    # Each call to _get_available_port returns a strictly increasing port so
+    # that the two daemons land on different ports when they are not sharing.
+    port_seq = [6200]
+
+    def fake_get_available_port(start_port):
+        p = port_seq[0]
+        port_seq[0] += 1
+        return p
+
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._get_available_port",
+        fake_get_available_port,
+    )
+
+    alive_pids: set[int] = set()
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._pid_is_alive", lambda pid: pid in alive_pids
+    )
+    monkeypatch.setattr("leann.embedding_server_manager._check_port", lambda port: True)
+
+    pid_seq = [55001]
+    starts: list[dict] = []
+
+    def fake_start_new_server(self, port, model_name, embedding_mode, **kwargs):
+        pid = pid_seq[0]
+        pid_seq[0] += 1
+        alive_pids.add(pid)
+        starts.append({"port": port, "config": kwargs.get("config_signature")})
+        self.server_process = DummyProcess(pid=pid)
+        self.server_port = port
+        self._server_config = kwargs.get("config_signature")
+        return True, port
+
+    monkeypatch.setattr(EmbeddingServerManager, "_start_new_server", fake_start_new_server)
+
+    # Simulate two separate test runs, each in its own temp subdirectory,
+    # using the *same* filename patterns but *different* parent directories.
+    run_a = tmp_path / "run_a"
+    run_a.mkdir()
+    meta_a = run_a / "myindex.meta.json"
+    passages_a = run_a / "myindex.passages.jsonl"
+    passages_a.write_text("passage: hello world\n", encoding="utf-8")
+    _write_meta(meta_a, passages_a.name, "myindex.passages.idx", total=1)
+
+    run_b = tmp_path / "run_b"
+    run_b.mkdir()
+    meta_b = run_b / "myindex.meta.json"
+    passages_b = run_b / "myindex.passages.jsonl"
+    passages_b.write_text("passage: different content\n", encoding="utf-8")
+    _write_meta(meta_b, passages_b.name, "myindex.passages.idx", total=1)
+
+    # Run A: start daemon with index_a passages
+    manager_a = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    ok_a, port_a = manager_a.start_server(
+        port=6200,
+        model_name="test-model",
+        passages_file=str(meta_a),
+        use_daemon=True,
+    )
+    assert ok_a
+    assert len(starts) == 1, "First manager should start one daemon"
+
+    # Run B: a fresh manager with a DIFFERENT passages file must not adopt run A's daemon
+    manager_b = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    ok_b, port_b = manager_b.start_server(
+        port=6200,
+        model_name="test-model",
+        passages_file=str(meta_b),
+        use_daemon=True,
+    )
+    assert ok_b
+    assert len(starts) == 2, (
+        "A separate daemon must be spawned for each distinct passages file; "
+        "sharing daemons across different indices causes 'Failed to fetch embeddings' errors"
+    )
+    assert port_a != port_b, "Each index must communicate with its own daemon port"
+
+
+def test_same_passages_file_reuses_existing_daemon(tmp_path, monkeypatch):
+    """Two managers with identical passages files SHOULD share the same daemon."""
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._get_available_port", lambda start_port: start_port
+    )
+
+    alive_pids: set[int] = {99001}
+    monkeypatch.setattr(
+        "leann.embedding_server_manager._pid_is_alive", lambda pid: pid in alive_pids
+    )
+    monkeypatch.setattr("leann.embedding_server_manager._check_port", lambda port: True)
+
+    starts: list[int] = []
+
+    def fake_start_new_server(self, port, model_name, embedding_mode, **kwargs):
+        starts.append(port)
+        self.server_process = DummyProcess(pid=99001)
+        self.server_port = port
+        self._server_config = kwargs.get("config_signature")
+        return True, port
+
+    monkeypatch.setattr(EmbeddingServerManager, "_start_new_server", fake_start_new_server)
+
+    meta = tmp_path / "shared.meta.json"
+    passages = tmp_path / "shared.passages.jsonl"
+    passages.write_text("shared passage\n", encoding="utf-8")
+    _write_meta(meta, passages.name, "shared.passages.idx", total=1)
+
+    manager1 = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    ok1, port1 = manager1.start_server(
+        port=6300, model_name="test-model", passages_file=str(meta), use_daemon=True
+    )
+    assert ok1 and len(starts) == 1
+
+    manager2 = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    ok2, port2 = manager2.start_server(
+        port=6300, model_name="test-model", passages_file=str(meta), use_daemon=True
+    )
+    assert ok2
+    assert len(starts) == 1, "Second manager should reuse the existing daemon, not start a new one"
+    assert port1 == port2, "Both managers should connect to the same daemon port"
+
+
+def test_stale_lock_info_removed_when_pid_dead(tmp_path, monkeypatch):
+    registry_dir = tmp_path / "servers"
+    registry_dir.mkdir()
+    monkeypatch.setattr(EmbeddingServerManager, "_registry_dir", staticmethod(lambda: registry_dir))
+    monkeypatch.setattr("leann.embedding_server_manager._pid_is_alive", lambda pid: False)
+
+    manager = EmbeddingServerManager("leann_backend_hnsw.hnsw_embedding_server")
+    signature = manager._build_config_signature(
+        model_name="x",
+        embedding_mode="sentence-transformers",
+        provider_options=None,
+        passages_file=None,
+        distance_metric=None,
+    )
+    key = manager._registry_key(signature)
+    lock_info = registry_dir / f"{key}.lockinfo.json"
+    lock_info.write_text(json.dumps({"pid": 999999, "ts": time.time()}), encoding="utf-8")
+
+    with manager._registry_lock(signature):
+        pass
+
+    assert not lock_info.exists()