chore: import upstream snapshot with attribution

tests/__init__.py

@@ -0,0 +1 @@
+"""Test package markers for shared helper imports."""

tests/api/auth/test_auth.py

@@ -0,0 +1,272 @@
+import importlib
+import sys
+from types import SimpleNamespace
+
+import bcrypt
+import pytest
+from fastapi import HTTPException
+from starlette.status import HTTP_403_FORBIDDEN
+
+from lightrag.api.passwords import BCRYPT_PASSWORD_PREFIX, hash_password
+from lightrag.tools.hash_password import main as hash_password_main
+from lightrag.utils import logger as lightrag_logger
+
+
+def import_real_api_module(module_name: str):
+    sys.modules.pop(module_name, None)
+
+    package_name, _, child_name = module_name.rpartition(".")
+    package = sys.modules.get(package_name)
+    if package is not None and hasattr(package, child_name):
+        delattr(package, child_name)
+
+    return importlib.import_module(module_name)
+
+
+@pytest.fixture
+def auth_module(monkeypatch):
+    config = import_real_api_module("lightrag.api.config")
+
+    mock_global_args = SimpleNamespace(
+        token_secret="test-jwt-secret",
+        jwt_algorithm="HS256",
+        token_expire_hours=48,
+        guest_token_expire_hours=24,
+        auth_accounts="admin:admin_pass",
+    )
+
+    monkeypatch.setattr(config, "global_args", mock_global_args)
+
+    module = import_real_api_module("lightrag.api.auth")
+    module = importlib.reload(module)
+    yield module
+    sys.modules.pop("lightrag.api.auth", None)
+
+
+def build_bcrypt_value(password: str) -> str:
+    hashed = bcrypt.hashpw(password.encode("utf-8"), bcrypt.gensalt()).decode("utf-8")
+    return f"{BCRYPT_PASSWORD_PREFIX}{hashed}"
+
+
+def test_verify_plaintext_password(auth_module):
+    handler = auth_module.AuthHandler()
+    handler.accounts = {"admin": "admin_pass"}
+
+    assert handler.verify_password("admin", "admin_pass")
+    assert not handler.verify_password("admin", "wrong_pass")
+
+
+def test_verify_prefixed_bcrypt_password(auth_module):
+    handler = auth_module.AuthHandler()
+    handler.accounts = {"user": build_bcrypt_value("user_pass")}
+
+    assert handler.verify_password("user", "user_pass")
+    assert not handler.verify_password("user", "wrong_pass")
+
+
+def test_plaintext_password_with_bcrypt_prefix_stays_plaintext(auth_module):
+    handler = auth_module.AuthHandler()
+    handler.accounts = {"user": "$2b$not-a-real-hash"}
+
+    assert handler.verify_password("user", "$2b$not-a-real-hash")
+    assert not handler.verify_password("user", "anything-else")
+
+
+def test_invalid_auth_accounts_raises(monkeypatch):
+    config = import_real_api_module("lightrag.api.config")
+
+    mock_global_args = SimpleNamespace(
+        token_secret="test-jwt-secret",
+        jwt_algorithm="HS256",
+        token_expire_hours=48,
+        guest_token_expire_hours=24,
+        auth_accounts="admin",
+    )
+
+    monkeypatch.setattr(config, "global_args", mock_global_args)
+
+    with pytest.raises(ValueError, match="AUTH_ACCOUNTS must use"):
+        import_real_api_module("lightrag.api.auth")
+
+    sys.modules.pop("lightrag.api.auth", None)
+
+
+def test_initialize_config_rejects_default_token_secret_with_auth_accounts():
+    config = import_real_api_module("lightrag.api.config")
+
+    insecure_args = SimpleNamespace(
+        auth_accounts="admin:admin_pass",
+        token_secret=config.DEFAULT_TOKEN_SECRET,
+    )
+
+    with pytest.raises(ValueError, match="TOKEN_SECRET must be explicitly set"):
+        config.initialize_config(insecure_args, force=True)
+
+
+def test_initialize_config_allows_custom_token_secret_with_auth_accounts():
+    config = import_real_api_module("lightrag.api.config")
+
+    secure_args = SimpleNamespace(
+        auth_accounts="admin:admin_pass",
+        token_secret="custom-jwt-secret",
+    )
+
+    initialized = config.initialize_config(secure_args, force=True)
+
+    assert initialized is secure_args
+
+
+def test_guest_tokens_fall_back_to_default_secret_when_token_secret_missing(
+    monkeypatch,
+):
+    config = import_real_api_module("lightrag.api.config")
+
+    mock_global_args = SimpleNamespace(
+        token_secret=None,
+        jwt_algorithm="HS256",
+        token_expire_hours=48,
+        guest_token_expire_hours=24,
+        auth_accounts="",
+    )
+
+    monkeypatch.setattr(config, "global_args", mock_global_args)
+    warning_messages = []
+
+    def capture_warning(message):
+        warning_messages.append(message)
+
+    monkeypatch.setattr(lightrag_logger, "warning", capture_warning)
+
+    module = import_real_api_module("lightrag.api.auth")
+    module = importlib.reload(module)
+    handler = module.AuthHandler()
+
+    token = handler.create_token("guest", role="guest")
+    token_info = handler.validate_token(token)
+
+    assert handler.secret == config.DEFAULT_TOKEN_SECRET
+    assert token_info["username"] == "guest"
+    assert token_info["role"] == "guest"
+    assert any(
+        "Falling back to the default guest-mode JWT secret" in msg
+        for msg in warning_messages
+    )
+
+    sys.modules.pop("lightrag.api.auth", None)
+
+
+def _load_api_key_only_modules(monkeypatch):
+    """Import auth + utils_api in the API-key-only profile.
+
+    API-key-only = an API key is set but AUTH_ACCOUNTS is empty, so
+    ``auth_configured`` is False and AuthHandler falls back to the public
+    DEFAULT_TOKEN_SECRET. Returns (config, auth, utils_api).
+    """
+    config = import_real_api_module("lightrag.api.config")
+
+    mock_global_args = SimpleNamespace(
+        token_secret=None,  # -> AuthHandler falls back to DEFAULT_TOKEN_SECRET
+        jwt_algorithm="HS256",
+        token_expire_hours=48,
+        guest_token_expire_hours=24,
+        auth_accounts="",  # no password accounts -> auth_configured is False
+        whitelist_paths="/health",  # consumed by utils_api at import time
+        token_auto_renew=False,
+    )
+    monkeypatch.setattr(config, "global_args", mock_global_args)
+
+    auth = import_real_api_module("lightrag.api.auth")
+    auth = importlib.reload(auth)
+    utils_api = import_real_api_module("lightrag.api.utils_api")
+    return config, auth, utils_api
+
+
+async def test_combined_auth_rejects_guest_token_in_api_key_mode(monkeypatch):
+    """A forged/obtained guest token must not bypass the X-API-Key check.
+
+    Regression for GHSA-f4vv-55c2-5789 / GHSA-xr5c-v5r6-c9f9: in the
+    API-key-only profile an anonymous caller can mint a guest JWT with the
+    public default secret (or fetch one from /auth-status) and previously
+    short-circuited authorization before the API key was ever checked.
+    """
+    config, auth, utils_api = _load_api_key_only_modules(monkeypatch)
+    try:
+        assert utils_api.auth_configured is False
+        assert auth.auth_handler.secret == config.DEFAULT_TOKEN_SECRET
+
+        forged_guest = auth.auth_handler.create_token(username="guest", role="guest")
+
+        api_key = "operator-secret-key"
+        dependency = utils_api.get_combined_auth_dependency(api_key=api_key)
+
+        request = SimpleNamespace(url=SimpleNamespace(path="/documents"), scope={})
+        response = SimpleNamespace(headers={})
+
+        # Forged guest token, no API key -> rejected (was HTTP 200 before the fix).
+        with pytest.raises(HTTPException) as rejected:
+            await dependency(
+                request=request,
+                response=response,
+                token=forged_guest,
+                api_key_header_value=None,
+            )
+        assert rejected.value.status_code == HTTP_403_FORBIDDEN
+
+        # Guest token together with the real API key -> allowed (the WebUI sends
+        # both Authorization and X-API-Key headers).
+        result = await dependency(
+            request=request,
+            response=response,
+            token=forged_guest,
+            api_key_header_value=api_key,
+        )
+        assert result is None
+    finally:
+        sys.modules.pop("lightrag.api.utils_api", None)
+        sys.modules.pop("lightrag.api.auth", None)
+
+
+async def test_combined_auth_allows_guest_token_when_fully_open(monkeypatch):
+    """Fully-open mode (no API key, no AUTH_ACCOUNTS) keeps accepting guest tokens.
+
+    Guards against over-fixing: the bypass fix must not break zero-config guest
+    access, which is the intended convenience mode.
+    """
+    _config, auth, utils_api = _load_api_key_only_modules(monkeypatch)
+    try:
+        guest_token = auth.auth_handler.create_token(username="guest", role="guest")
+
+        dependency = utils_api.get_combined_auth_dependency(api_key=None)
+        request = SimpleNamespace(url=SimpleNamespace(path="/documents"), scope={})
+        response = SimpleNamespace(headers={})
+
+        result = await dependency(
+            request=request,
+            response=response,
+            token=guest_token,
+            api_key_header_value=None,
+        )
+        assert result is None
+    finally:
+        sys.modules.pop("lightrag.api.utils_api", None)
+        sys.modules.pop("lightrag.api.auth", None)
+
+
+def test_hash_password_returns_prefixed_value(auth_module):
+    hashed = hash_password("new_password")
+
+    assert hashed.startswith(BCRYPT_PASSWORD_PREFIX)
+    raw_hash = hashed[len(BCRYPT_PASSWORD_PREFIX) :]
+    assert bcrypt.checkpw("new_password".encode("utf-8"), raw_hash.encode("utf-8"))
+
+
+def test_hash_password_cli_outputs_auth_accounts_entry(capsys):
+    exit_code = hash_password_main(["--username", "admin", "secret"])
+
+    assert exit_code == 0
+    output = capsys.readouterr().out.strip()
+    username, hashed = output.split(":", 1)
+    assert username == "admin"
+    assert hashed.startswith(BCRYPT_PASSWORD_PREFIX)
+    raw_hash = hashed[len(BCRYPT_PASSWORD_PREFIX) :]
+    assert bcrypt.checkpw("secret".encode("utf-8"), raw_hash.encode("utf-8"))

tests/api/auth/test_token_auto_renewal.py

@@ -0,0 +1,403 @@
+"""
+Pytest unit tests for token auto-renewal functionality
+
+Tests:
+1. Backend token renewal logic
+2. Rate limiting for token renewals
+3. Token renewal state tracking
+"""
+
+import pytest
+from datetime import datetime, timedelta, timezone
+from unittest.mock import Mock
+from fastapi import Response
+import time
+
+# Create a simple token renewal cache for testing
+_token_renewal_cache = {}
+_RENEWAL_MIN_INTERVAL = 60
+
+
+@pytest.mark.offline
+class TestTokenRenewal:
+    """Tests for token auto-renewal logic"""
+
+    @pytest.fixture
+    def mock_auth_handler(self):
+        """Mock authentication handler"""
+        handler = Mock()
+        handler.guest_expire_hours = 24
+        handler.expire_hours = 24
+        handler.create_token = Mock(return_value="new-token-12345")
+        return handler
+
+    @pytest.fixture
+    def mock_global_args(self):
+        """Mock global configuration"""
+        args = Mock()
+        args.token_auto_renew = True
+        args.token_renew_threshold = 0.5
+        return args
+
+    @pytest.fixture
+    def mock_token_info_guest(self):
+        """Mock token info for guest user"""
+        # Token with 10 hours remaining (below 50% of 24 hours)
+        exp_time = datetime.now(timezone.utc) + timedelta(hours=10)
+        return {
+            "username": "guest",
+            "role": "guest",
+            "exp": exp_time,
+            "metadata": {"auth_mode": "disabled"},
+        }
+
+    @pytest.fixture
+    def mock_token_info_user(self):
+        """Mock token info for regular user"""
+        # Token with 10 hours remaining (below 50% of 24 hours)
+        exp_time = datetime.now(timezone.utc) + timedelta(hours=10)
+        return {
+            "username": "testuser",
+            "role": "user",
+            "exp": exp_time,
+            "metadata": {"auth_mode": "enabled"},
+        }
+
+    @pytest.fixture
+    def mock_token_info_above_threshold(self):
+        """Mock token info with time above renewal threshold"""
+        # Token with 20 hours remaining (above 50% of 24 hours)
+        exp_time = datetime.now(timezone.utc) + timedelta(hours=20)
+        return {
+            "username": "testuser",
+            "role": "user",
+            "exp": exp_time,
+            "metadata": {"auth_mode": "enabled"},
+        }
+
+    def test_token_renewal_when_below_threshold(
+        self, mock_auth_handler, mock_global_args, mock_token_info_user
+    ):
+        """Test that token is renewed when remaining time < threshold"""
+        # Use global cache
+        global _token_renewal_cache
+
+        # Clear cache
+        _token_renewal_cache.clear()
+
+        response = Mock(spec=Response)
+        response.headers = {}
+
+        # Simulate the renewal logic
+        expire_time = mock_token_info_user["exp"]
+        now = datetime.now(timezone.utc)
+        remaining_seconds = (expire_time - now).total_seconds()
+
+        role = mock_token_info_user["role"]
+        total_hours = (
+            mock_auth_handler.expire_hours
+            if role == "user"
+            else mock_auth_handler.guest_expire_hours
+        )
+        total_seconds = total_hours * 3600
+
+        # Should renew because remaining_seconds < total_seconds * 0.5
+        should_renew = (
+            remaining_seconds < total_seconds * mock_global_args.token_renew_threshold
+        )
+        assert should_renew is True
+
+        # Simulate renewal
+        username = mock_token_info_user["username"]
+        current_time = time.time()
+        last_renewal = _token_renewal_cache.get(username, 0)
+        time_since_last_renewal = current_time - last_renewal
+
+        # Should pass rate limit (first renewal)
+        assert time_since_last_renewal >= 60 or last_renewal == 0
+
+        # Perform renewal
+        new_token = mock_auth_handler.create_token(
+            username=username, role=role, metadata=mock_token_info_user["metadata"]
+        )
+        response.headers["X-New-Token"] = new_token
+        _token_renewal_cache[username] = current_time
+
+        # Verify
+        assert "X-New-Token" in response.headers
+        assert response.headers["X-New-Token"] == "new-token-12345"
+        assert username in _token_renewal_cache
+
+    def test_token_no_renewal_when_above_threshold(
+        self, mock_auth_handler, mock_global_args, mock_token_info_above_threshold
+    ):
+        """Test that token is NOT renewed when remaining time > threshold"""
+        response = Mock(spec=Response)
+        response.headers = {}
+
+        expire_time = mock_token_info_above_threshold["exp"]
+        now = datetime.now(timezone.utc)
+        remaining_seconds = (expire_time - now).total_seconds()
+
+        mock_token_info_above_threshold["role"]
+        total_hours = mock_auth_handler.expire_hours
+        total_seconds = total_hours * 3600
+
+        # Should NOT renew because remaining_seconds > total_seconds * 0.5
+        should_renew = (
+            remaining_seconds < total_seconds * mock_global_args.token_renew_threshold
+        )
+        assert should_renew is False
+
+        # No renewal should happen
+        assert "X-New-Token" not in response.headers
+
+    def test_token_renewal_disabled(
+        self, mock_auth_handler, mock_global_args, mock_token_info_user
+    ):
+        """Test that no renewal happens when TOKEN_AUTO_RENEW=false"""
+        mock_global_args.token_auto_renew = False
+        response = Mock(spec=Response)
+        response.headers = {}
+
+        # Auto-renewal is disabled, so even if below threshold, no renewal
+        if not mock_global_args.token_auto_renew:
+            # Skip renewal logic
+            pass
+
+        assert "X-New-Token" not in response.headers
+
+    def test_token_renewal_for_guest_mode(
+        self, mock_auth_handler, mock_global_args, mock_token_info_guest
+    ):
+        """Test that guest tokens are renewed correctly"""
+        # Use global cache
+        global _token_renewal_cache
+
+        _token_renewal_cache.clear()
+
+        response = Mock(spec=Response)
+        response.headers = {}
+
+        expire_time = mock_token_info_guest["exp"]
+        now = datetime.now(timezone.utc)
+        remaining_seconds = (expire_time - now).total_seconds()
+
+        role = mock_token_info_guest["role"]
+        total_hours = mock_auth_handler.guest_expire_hours
+        total_seconds = total_hours * 3600
+
+        should_renew = (
+            remaining_seconds < total_seconds * mock_global_args.token_renew_threshold
+        )
+        assert should_renew is True
+
+        # Renewal for guest
+        username = mock_token_info_guest["username"]
+        new_token = mock_auth_handler.create_token(
+            username=username, role=role, metadata=mock_token_info_guest["metadata"]
+        )
+        response.headers["X-New-Token"] = new_token
+        _token_renewal_cache[username] = time.time()
+
+        assert "X-New-Token" in response.headers
+        assert username in _token_renewal_cache
+
+
+@pytest.mark.offline
+class TestRateLimiting:
+    """Tests for token renewal rate limiting"""
+
+    @pytest.fixture
+    def mock_auth_handler(self):
+        """Mock authentication handler"""
+        handler = Mock()
+        handler.expire_hours = 24
+        handler.create_token = Mock(return_value="new-token-12345")
+        return handler
+
+    def test_rate_limit_prevents_rapid_renewals(self, mock_auth_handler):
+        """Test that second renewal within 60s is blocked"""
+        # Use global cache and constant
+        global _token_renewal_cache, _RENEWAL_MIN_INTERVAL
+
+        username = "testuser"
+        _token_renewal_cache.clear()
+
+        # First renewal
+        current_time_1 = time.time()
+        _token_renewal_cache[username] = current_time_1
+
+        response_1 = Mock(spec=Response)
+        response_1.headers = {}
+        response_1.headers["X-New-Token"] = "new-token-12345"
+
+        # Immediate second renewal attempt (within 60s)
+        current_time_2 = time.time()  # Almost same time
+        last_renewal = _token_renewal_cache.get(username, 0)
+        time_since_last_renewal = current_time_2 - last_renewal
+
+        # Should be blocked by rate limit
+        assert time_since_last_renewal < _RENEWAL_MIN_INTERVAL
+
+        response_2 = Mock(spec=Response)
+        response_2.headers = {}
+
+        # No new token should be issued
+        if time_since_last_renewal < _RENEWAL_MIN_INTERVAL:
+            # Rate limited, skip renewal
+            pass
+
+        assert "X-New-Token" not in response_2.headers
+
+    def test_rate_limit_allows_renewal_after_interval(self, mock_auth_handler):
+        """Test that renewal succeeds after 60s interval"""
+        # Use global cache and constant
+        global _token_renewal_cache, _RENEWAL_MIN_INTERVAL
+
+        username = "testuser"
+        _token_renewal_cache.clear()
+
+        # First renewal at time T
+        first_renewal_time = time.time() - 61  # 61 seconds ago
+        _token_renewal_cache[username] = first_renewal_time
+
+        # Second renewal attempt now
+        current_time = time.time()
+        last_renewal = _token_renewal_cache.get(username, 0)
+        time_since_last_renewal = current_time - last_renewal
+
+        # Should pass rate limit (>60s elapsed)
+        assert time_since_last_renewal >= _RENEWAL_MIN_INTERVAL
+
+        response = Mock(spec=Response)
+        response.headers = {}
+
+        if time_since_last_renewal >= _RENEWAL_MIN_INTERVAL:
+            new_token = mock_auth_handler.create_token(
+                username=username, role="user", metadata={}
+            )
+            response.headers["X-New-Token"] = new_token
+            _token_renewal_cache[username] = current_time
+
+        assert "X-New-Token" in response.headers
+        assert response.headers["X-New-Token"] == "new-token-12345"
+
+    def test_rate_limit_per_user(self, mock_auth_handler):
+        """Test that different users have independent rate limits"""
+        # Use global cache
+        global _token_renewal_cache
+
+        _token_renewal_cache.clear()
+
+        user1 = "user1"
+        user2 = "user2"
+
+        current_time = time.time()
+
+        # User1 gets renewal
+        _token_renewal_cache[user1] = current_time
+
+        # User2 should still be able to get renewal (independent cache)
+        last_renewal_user2 = _token_renewal_cache.get(user2, 0)
+        assert last_renewal_user2 == 0  # No previous renewal
+
+        # User2 can renew
+        _token_renewal_cache[user2] = current_time
+
+        # Both users should have entries
+        assert user1 in _token_renewal_cache
+        assert user2 in _token_renewal_cache
+        assert _token_renewal_cache[user1] == _token_renewal_cache[user2]
+
+
+@pytest.mark.offline
+class TestTokenExpirationCalculation:
+    """Tests for token expiration time calculation"""
+
+    def test_expiration_extraction_from_jwt(self):
+        """Test extracting expiration time from JWT token"""
+        import base64
+        import json
+
+        # Create a mock JWT payload
+        exp_timestamp = int(
+            (datetime.now(timezone.utc) + timedelta(hours=24)).timestamp()
+        )
+        payload = {"sub": "testuser", "role": "user", "exp": exp_timestamp}
+
+        # Encode as base64 (simulating JWT structure: header.payload.signature)
+        payload_b64 = base64.b64encode(json.dumps(payload).encode()).decode()
+        mock_token = f"header.{payload_b64}.signature"
+
+        # Simulate extraction
+        parts = mock_token.split(".")
+        assert len(parts) == 3
+
+        decoded_payload = json.loads(base64.b64decode(parts[1]))
+        assert decoded_payload["exp"] == exp_timestamp
+        assert decoded_payload["sub"] == "testuser"
+
+    def test_remaining_time_calculation(self):
+        """Test calculation of remaining token time"""
+        # Token expires in 10 hours
+        exp_time = datetime.now(timezone.utc) + timedelta(hours=10)
+        now = datetime.now(timezone.utc)
+
+        remaining_seconds = (exp_time - now).total_seconds()
+
+        # Should be approximately 10 hours (36000 seconds)
+        assert 35990 < remaining_seconds < 36010
+
+        # Calculate percentage remaining (for 24-hour token)
+        total_seconds = 24 * 3600
+        percentage_remaining = remaining_seconds / total_seconds
+
+        # Should be approximately 41.67% remaining
+        assert 0.41 < percentage_remaining < 0.42
+
+    def test_threshold_comparison(self):
+        """Test threshold-based renewal decision"""
+        threshold = 0.5
+        total_hours = 24
+        total_seconds = total_hours * 3600
+
+        # Scenario 1: 10 hours remaining -> should renew
+        remaining_seconds_1 = 10 * 3600
+        should_renew_1 = remaining_seconds_1 < total_seconds * threshold
+        assert should_renew_1 is True
+
+        # Scenario 2: 20 hours remaining -> should NOT renew
+        remaining_seconds_2 = 20 * 3600
+        should_renew_2 = remaining_seconds_2 < total_seconds * threshold
+        assert should_renew_2 is False
+
+        # Scenario 3: Exactly 12 hours remaining (at threshold) -> should NOT renew
+        remaining_seconds_3 = 12 * 3600
+        should_renew_3 = remaining_seconds_3 < total_seconds * threshold
+        assert should_renew_3 is False
+
+
+@pytest.mark.offline
+def test_renewal_cache_cleanup():
+    """Test that renewal cache can be cleared"""
+    # Use global cache
+    global _token_renewal_cache
+
+    # Clear first
+    _token_renewal_cache.clear()
+
+    # Add some entries
+    _token_renewal_cache["user1"] = time.time()
+    _token_renewal_cache["user2"] = time.time()
+
+    assert len(_token_renewal_cache) == 2
+
+    # Clear cache
+    _token_renewal_cache.clear()
+
+    assert len(_token_renewal_cache) == 0
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])

tests/api/config/test_api_config_bedrock.py

@@ -0,0 +1,141 @@
+import sys
+
+import pytest
+
+from lightrag.api.config import get_default_host, parse_args
+
+
+pytestmark = pytest.mark.offline
+
+
+def _clear_bedrock_auth_env(monkeypatch):
+    for key in (
+        "LLM_BINDING",
+        "EMBEDDING_BINDING",
+        "QUERY_LLM_BINDING",
+        "QUERY_LLM_BINDING_API_KEY",
+        "AWS_ACCESS_KEY_ID",
+        "AWS_SECRET_ACCESS_KEY",
+        "AWS_SESSION_TOKEN",
+        "AWS_BEARER_TOKEN_BEDROCK",
+        "QUERY_AWS_ACCESS_KEY_ID",
+        "QUERY_AWS_SECRET_ACCESS_KEY",
+    ):
+        monkeypatch.delenv(key, raising=False)
+
+
+def test_bedrock_default_host_uses_sdk_default_endpoint(monkeypatch):
+    monkeypatch.delenv("LLM_BINDING_HOST", raising=False)
+
+    assert get_default_host("bedrock") == "DEFAULT_BEDROCK_ENDPOINT"
+
+
+def test_bedrock_custom_host_is_returned(monkeypatch):
+    monkeypatch.setenv("LLM_BINDING_HOST", "https://proxy.example.com")
+
+    assert get_default_host("bedrock") == "https://proxy.example.com"
+
+
+def test_bedrock_env_binding_alias_is_normalized(monkeypatch):
+    _clear_bedrock_auth_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "aws_bedrock")
+    monkeypatch.setenv("EMBEDDING_BINDING", "aws_bedrock")
+    monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "absk-test")
+
+    args = parse_args()
+
+    assert args.llm_binding == "bedrock"
+    assert args.embedding_binding == "bedrock"
+
+
+def test_bedrock_cli_binding_alias_is_not_supported(monkeypatch):
+    monkeypatch.setattr(
+        sys, "argv", ["lightrag-server", "--llm-binding", "aws_bedrock"]
+    )
+
+    with pytest.raises(SystemExit):
+        parse_args()
+
+
+def test_bedrock_role_env_binding_alias_is_normalized(monkeypatch):
+    _clear_bedrock_auth_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("QUERY_LLM_BINDING", "aws_bedrock")
+    monkeypatch.setenv("QUERY_LLM_MODEL", "us.amazon.nova-lite-v1:0")
+    monkeypatch.setenv("QUERY_AWS_REGION", "us-west-2")
+    monkeypatch.setenv("QUERY_AWS_ACCESS_KEY_ID", "akid")
+    monkeypatch.setenv("QUERY_AWS_SECRET_ACCESS_KEY", "secret")
+
+    args = parse_args()
+
+    assert args.query_llm_binding == "bedrock"
+    assert args.query_aws_region == "us-west-2"
+
+
+def test_bedrock_role_api_key_is_rejected(monkeypatch):
+    _clear_bedrock_auth_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("QUERY_LLM_BINDING", "bedrock")
+    monkeypatch.setenv("QUERY_LLM_MODEL", "us.amazon.nova-lite-v1:0")
+    monkeypatch.setenv("QUERY_LLM_BINDING_API_KEY", "absk-role")
+
+    with pytest.raises(SystemExit, match="does not support QUERY_LLM_BINDING_API_KEY"):
+        parse_args()
+
+
+def test_bedrock_binding_requires_sigv4_pair_or_bearer_token(monkeypatch):
+    _clear_bedrock_auth_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "bedrock")
+    monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
+
+    with pytest.raises(ValueError, match="Bedrock LLM binding requires"):
+        parse_args()
+
+
+def test_bedrock_binding_rejects_partial_sigv4_pair(monkeypatch):
+    _clear_bedrock_auth_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "bedrock")
+    monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
+    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "akid")
+
+    with pytest.raises(ValueError, match="Bedrock LLM binding requires"):
+        parse_args()
+
+
+def test_bedrock_binding_accepts_sigv4_pair(monkeypatch):
+    _clear_bedrock_auth_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "bedrock")
+    monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
+    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "akid")
+    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret")
+
+    args = parse_args()
+
+    assert args.llm_binding == "bedrock"
+
+
+def test_bedrock_binding_accepts_bearer_token(monkeypatch):
+    _clear_bedrock_auth_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "bedrock")
+    monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
+    monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "absk-test")
+
+    args = parse_args()
+
+    assert args.llm_binding == "bedrock"
+
+
+def test_bedrock_role_requires_complete_role_sigv4_pair(monkeypatch):
+    _clear_bedrock_auth_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("QUERY_LLM_BINDING", "bedrock")
+    monkeypatch.setenv("QUERY_LLM_MODEL", "us.amazon.nova-lite-v1:0")
+    monkeypatch.setenv("QUERY_AWS_ACCESS_KEY_ID", "akid")
+
+    with pytest.raises(ValueError, match="Bedrock role 'query' requires"):
+        parse_args()

tests/api/config/test_api_config_gemini.py

@@ -0,0 +1,12 @@
+import pytest
+
+from lightrag.api.config import get_default_host
+
+
+pytestmark = pytest.mark.offline
+
+
+def test_gemini_default_host_uses_sdk_default_endpoint(monkeypatch):
+    monkeypatch.delenv("LLM_BINDING_HOST", raising=False)
+
+    assert get_default_host("gemini") == "DEFAULT_GEMINI_ENDPOINT"

tests/api/config/test_api_config_role_max_async.py

@@ -0,0 +1,95 @@
+import sys
+
+import pytest
+
+from lightrag.api.config import parse_args
+from lightrag.constants import DEFAULT_MAX_ASYNC
+
+
+pytestmark = pytest.mark.offline
+
+
+ROLE_MAX_ASYNC_ENV_KEYS = (
+    "MAX_ASYNC_LLM",
+    "MAX_ASYNC",
+    "EXTRACT_MAX_ASYNC_LLM",
+    "KEYWORD_MAX_ASYNC_LLM",
+    "QUERY_MAX_ASYNC_LLM",
+    "VLM_MAX_ASYNC_LLM",
+)
+
+
+def _clear_max_async_env(monkeypatch):
+    for key in ROLE_MAX_ASYNC_ENV_KEYS:
+        monkeypatch.delenv(key, raising=False)
+
+
+def test_role_max_async_defaults_none_when_env_unset(monkeypatch):
+    _clear_max_async_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("MAX_ASYNC", "10")
+
+    args = parse_args()
+
+    assert args.max_async == 10
+    assert args.extract_llm_max_async is None
+    assert args.keyword_llm_max_async is None
+    assert args.query_llm_max_async is None
+    assert args.vlm_llm_max_async is None
+
+
+def test_role_max_async_env_override_keeps_other_roles_none(monkeypatch):
+    _clear_max_async_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("MAX_ASYNC", "10")
+    monkeypatch.setenv("EXTRACT_MAX_ASYNC_LLM", "7")
+
+    args = parse_args()
+
+    assert args.max_async == 10
+    assert args.extract_llm_max_async == 7
+    assert args.keyword_llm_max_async is None
+    assert args.query_llm_max_async is None
+    assert args.vlm_llm_max_async is None
+
+
+def test_role_max_async_literal_none_string_is_preserved(monkeypatch):
+    _clear_max_async_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("MAX_ASYNC", "10")
+    monkeypatch.setenv("QUERY_MAX_ASYNC_LLM", "None")
+
+    args = parse_args()
+
+    assert args.max_async == 10
+    assert args.query_llm_max_async is None
+
+
+def test_max_async_llm_new_name(monkeypatch):
+    _clear_max_async_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("MAX_ASYNC_LLM", "8")
+
+    args = parse_args()
+
+    assert args.max_async == 8
+
+
+def test_max_async_llm_takes_precedence_over_legacy(monkeypatch):
+    _clear_max_async_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("MAX_ASYNC_LLM", "8")
+    monkeypatch.setenv("MAX_ASYNC", "10")
+
+    args = parse_args()
+
+    assert args.max_async == 8
+
+
+def test_max_async_defaults_when_unset(monkeypatch):
+    _clear_max_async_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+
+    args = parse_args()
+
+    assert args.max_async == DEFAULT_MAX_ASYNC

tests/api/config/test_api_config_vlm_enable.py

@@ -0,0 +1,111 @@
+"""Offline tests for VLM_PROCESS_ENABLE and renamed role timeout vars."""
+
+from __future__ import annotations
+
+import sys
+
+import pytest
+
+from lightrag.api.config import parse_args
+
+
+pytestmark = pytest.mark.offline
+
+
+def _reset_vlm_env(monkeypatch):
+    for key in (
+        "VLM_PROCESS_ENABLE",
+        "VLM_LLM_BINDING",
+        "VLM_LLM_MODEL",
+        "VLM_LLM_BINDING_HOST",
+        "VLM_LLM_BINDING_API_KEY",
+        "VLM_LLM_TIMEOUT",
+        "EXTRACT_LLM_TIMEOUT",
+        "KEYWORD_LLM_TIMEOUT",
+        "QUERY_LLM_TIMEOUT",
+        "LLM_TIMEOUT_VLM_LLM",
+        "LLM_TIMEOUT_EXTRACT_LLM",
+        "LLM_TIMEOUT_KEYWORD_LLM",
+        "LLM_TIMEOUT_QUERY_LLM",
+    ):
+        monkeypatch.delenv(key, raising=False)
+
+
+def test_vlm_process_enable_defaults_to_false(monkeypatch):
+    _reset_vlm_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "openai")
+
+    args = parse_args()
+
+    assert args.vlm_process_enable is False
+
+
+def test_vlm_process_enable_true_with_openai_passes(monkeypatch):
+    _reset_vlm_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "openai")
+    monkeypatch.setenv("VLM_PROCESS_ENABLE", "true")
+
+    args = parse_args()
+
+    assert args.vlm_process_enable is True
+
+
+def test_vlm_process_enable_rejects_lollms_base_binding(monkeypatch):
+    _reset_vlm_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "lollms")
+    monkeypatch.setenv("VLM_PROCESS_ENABLE", "true")
+
+    with pytest.raises(SystemExit) as exc:
+        parse_args()
+    assert "lollms" in str(exc.value).lower()
+
+
+def test_vlm_process_enable_rejects_lollms_role_binding(monkeypatch):
+    _reset_vlm_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "openai")
+    monkeypatch.setenv("VLM_PROCESS_ENABLE", "true")
+    monkeypatch.setenv("VLM_LLM_BINDING", "lollms")
+    # Cross-provider validation needs model + key; fill them so the lollms
+    # branch is the only failure path.
+    monkeypatch.setenv("VLM_LLM_MODEL", "anything")
+    monkeypatch.setenv("VLM_LLM_BINDING_HOST", "http://localhost:9600")
+    monkeypatch.setenv("VLM_LLM_BINDING_API_KEY", "placeholder")
+
+    with pytest.raises(SystemExit) as exc:
+        parse_args()
+    assert "lollms" in str(exc.value).lower()
+
+
+def test_role_timeout_uses_new_variable_names(monkeypatch):
+    _reset_vlm_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "openai")
+    monkeypatch.setenv("EXTRACT_LLM_TIMEOUT", "240")
+    monkeypatch.setenv("KEYWORD_LLM_TIMEOUT", "120")
+    monkeypatch.setenv("QUERY_LLM_TIMEOUT", "60")
+    monkeypatch.setenv("VLM_LLM_TIMEOUT", "300")
+
+    args = parse_args()
+
+    assert args.extract_llm_timeout == 240
+    assert args.keyword_llm_timeout == 120
+    assert args.query_llm_timeout == 60
+    assert args.vlm_llm_timeout == 300
+
+
+def test_role_timeout_legacy_variables_no_longer_have_effect(monkeypatch):
+    """The breaking-change migration: legacy LLM_TIMEOUT_{ROLE}_LLM is silently ignored."""
+    _reset_vlm_env(monkeypatch)
+    monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+    monkeypatch.setenv("LLM_BINDING", "openai")
+    monkeypatch.setenv("LLM_TIMEOUT_EXTRACT_LLM", "999")
+    monkeypatch.setenv("LLM_TIMEOUT_VLM_LLM", "888")
+
+    args = parse_args()
+
+    assert args.extract_llm_timeout is None
+    assert args.vlm_llm_timeout is None

tests/api/config/test_whitelist_exposes_api_routes.py

@@ -0,0 +1,80 @@
+"""Offline tests for whitelist_exposes_api_routes (GHSA-mmg5-8x8q-v934 banner).
+
+The startup banner uses this helper to warn when WHITELIST_PATHS leaves the
+Ollama-compatible /api/* routes unauthenticated on a network-exposed bind. It
+must mirror the prefix/exact matching in get_combined_auth_dependency so a
+catch-all entry like "/*" is detected, not just literal "/api..." entries.
+"""
+
+from __future__ import annotations
+
+import importlib
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+
+pytestmark = pytest.mark.offline
+
+
+@pytest.fixture
+def whitelist_exposes_api_routes(monkeypatch):
+    """Import utils_api with a stub global_args (consumed at import time)."""
+    config = importlib.import_module("lightrag.api.config")
+    monkeypatch.setattr(
+        config,
+        "global_args",
+        SimpleNamespace(
+            token_secret=None,  # -> AuthHandler falls back to DEFAULT_TOKEN_SECRET
+            jwt_algorithm="HS256",
+            token_expire_hours=48,
+            guest_token_expire_hours=24,
+            auth_accounts="",
+            whitelist_paths="/health",  # consumed by utils_api at import time
+            token_auto_renew=False,
+        ),
+    )
+    sys.modules.pop("lightrag.api.auth", None)
+    importlib.reload(importlib.import_module("lightrag.api.auth"))
+    sys.modules.pop("lightrag.api.utils_api", None)
+    utils_api = importlib.import_module("lightrag.api.utils_api")
+    try:
+        yield utils_api.whitelist_exposes_api_routes
+    finally:
+        sys.modules.pop("lightrag.api.utils_api", None)
+
+
+@pytest.mark.parametrize(
+    "whitelist",
+    [
+        "/health,/api/*",  # default whitelist
+        "/api/*",
+        "/api",
+        "/api/chat",  # exact Ollama route
+        "/*",  # catch-all: empty prefix matches every path
+        " /* ",  # catch-all with surrounding whitespace
+        "/health,/*",
+        "/a/*",  # prefix that /api also starts with
+    ],
+)
+def test_exposes_api_routes(whitelist_exposes_api_routes, whitelist: str) -> None:
+    assert whitelist_exposes_api_routes(whitelist) is True
+
+
+@pytest.mark.parametrize(
+    "whitelist",
+    [
+        "/health",
+        "/health,/docs",
+        "/health/*",  # unrelated prefix
+        "/apiary/*",  # prefix under /api... but not /api/ — exempts only /apiary
+        "/apidocs",  # exact path that merely shares the /api substring
+        "",
+        "   ",
+    ],
+)
+def test_does_not_expose_api_routes(
+    whitelist_exposes_api_routes, whitelist: str
+) -> None:
+    assert whitelist_exposes_api_routes(whitelist) is False

tests/api/routes/test_aquery_data_endpoint.py

@@ -0,0 +1,781 @@
+#!/usr/bin/env python3
+"""
+Test script: Demonstrates usage of aquery_data FastAPI endpoint
+Query content: Who is the author of LightRAG
+
+Updated to handle the new data format where:
+- Response includes status, message, data, and metadata fields at top level
+- Actual query results (entities, relationships, chunks, references) are nested under 'data' field
+- Includes backward compatibility with legacy format
+"""
+
+import pytest
+import requests
+import time
+import json
+from typing import Dict, Any, List, Optional
+
+# API configuration
+API_KEY = "your-secure-api-key-here-123"
+BASE_URL = "http://localhost:9621"
+
+# Unified authentication headers
+AUTH_HEADERS = {"Content-Type": "application/json", "X-API-Key": API_KEY}
+
+
+def validate_references_format(references: List[Dict[str, Any]]) -> bool:
+    """Validate the format of references list"""
+    if not isinstance(references, list):
+        print(f"❌ References should be a list, got {type(references)}")
+        return False
+
+    for i, ref in enumerate(references):
+        if not isinstance(ref, dict):
+            print(f"❌ Reference {i} should be a dict, got {type(ref)}")
+            return False
+
+        required_fields = ["reference_id", "file_path"]
+        for field in required_fields:
+            if field not in ref:
+                print(f"❌ Reference {i} missing required field: {field}")
+                return False
+
+            if not isinstance(ref[field], str):
+                print(
+                    f"❌ Reference {i} field '{field}' should be string, got {type(ref[field])}"
+                )
+                return False
+
+    return True
+
+
+def parse_streaming_response(
+    response_text: str,
+) -> tuple[Optional[List[Dict]], List[str], List[str]]:
+    """Parse streaming response and extract references, response chunks, and errors"""
+    references = None
+    response_chunks = []
+    errors = []
+
+    lines = response_text.strip().split("\n")
+
+    for line in lines:
+        line = line.strip()
+        if not line or line.startswith("data: "):
+            if line.startswith("data: "):
+                line = line[6:]  # Remove 'data: ' prefix
+
+        if not line:
+            continue
+
+        try:
+            data = json.loads(line)
+
+            if "references" in data:
+                references = data["references"]
+            if "response" in data:
+                response_chunks.append(data["response"])
+            if "error" in data:
+                errors.append(data["error"])
+
+        except json.JSONDecodeError:
+            # Skip non-JSON lines (like SSE comments)
+            continue
+
+    return references, response_chunks, errors
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_query_endpoint_references():
+    """Test /query endpoint references functionality"""
+
+    print("\n" + "=" * 60)
+    print("Testing /query endpoint references functionality")
+    print("=" * 60)
+
+    query_text = "who authored LightRAG"
+    endpoint = f"{BASE_URL}/query"
+
+    # Test 1: References enabled (default)
+    print("\n🧪 Test 1: References enabled (default)")
+    print("-" * 40)
+
+    try:
+        response = requests.post(
+            endpoint,
+            json={"query": query_text, "mode": "mix", "include_references": True},
+            headers=AUTH_HEADERS,
+            timeout=30,
+        )
+
+        if response.status_code == 200:
+            data = response.json()
+
+            # Check response structure
+            if "response" not in data:
+                print("❌ Missing 'response' field")
+                return False
+
+            if "references" not in data:
+                print("❌ Missing 'references' field when include_references=True")
+                return False
+
+            references = data["references"]
+            if references is None:
+                print("❌ References should not be None when include_references=True")
+                return False
+
+            if not validate_references_format(references):
+                return False
+
+            print(f"✅ References enabled: Found {len(references)} references")
+            print(f"   Response length: {len(data['response'])} characters")
+
+            # Display reference list
+            if references:
+                print("   📚 Reference List:")
+                for i, ref in enumerate(references, 1):
+                    ref_id = ref.get("reference_id", "Unknown")
+                    file_path = ref.get("file_path", "Unknown")
+                    print(f"      {i}. ID: {ref_id} | File: {file_path}")
+
+        else:
+            print(f"❌ Request failed: {response.status_code}")
+            print(f"   Error: {response.text}")
+            return False
+
+    except Exception as e:
+        print(f"❌ Test 1 failed: {str(e)}")
+        return False
+
+    # Test 2: References disabled
+    print("\n🧪 Test 2: References disabled")
+    print("-" * 40)
+
+    try:
+        response = requests.post(
+            endpoint,
+            json={"query": query_text, "mode": "mix", "include_references": False},
+            headers=AUTH_HEADERS,
+            timeout=30,
+        )
+
+        if response.status_code == 200:
+            data = response.json()
+
+            # Check response structure
+            if "response" not in data:
+                print("❌ Missing 'response' field")
+                return False
+
+            references = data.get("references")
+            if references is not None:
+                print("❌ References should be None when include_references=False")
+                return False
+
+            print("✅ References disabled: No references field present")
+            print(f"   Response length: {len(data['response'])} characters")
+
+        else:
+            print(f"❌ Request failed: {response.status_code}")
+            print(f"   Error: {response.text}")
+            return False
+
+    except Exception as e:
+        print(f"❌ Test 2 failed: {str(e)}")
+        return False
+
+    print("\n✅ /query endpoint references tests passed!")
+    return True
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_query_stream_endpoint_references():
+    """Test /query/stream endpoint references functionality"""
+
+    print("\n" + "=" * 60)
+    print("Testing /query/stream endpoint references functionality")
+    print("=" * 60)
+
+    query_text = "who authored LightRAG"
+    endpoint = f"{BASE_URL}/query/stream"
+
+    # Test 1: Streaming with references enabled
+    print("\n🧪 Test 1: Streaming with references enabled")
+    print("-" * 40)
+
+    try:
+        response = requests.post(
+            endpoint,
+            json={"query": query_text, "mode": "mix", "include_references": True},
+            headers=AUTH_HEADERS,
+            timeout=30,
+            stream=True,
+        )
+
+        if response.status_code == 200:
+            # Collect streaming response
+            full_response = ""
+            for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
+                if chunk:
+                    # Ensure chunk is string type
+                    if isinstance(chunk, bytes):
+                        chunk = chunk.decode("utf-8")
+                    full_response += chunk
+
+            # Parse streaming response
+            references, response_chunks, errors = parse_streaming_response(
+                full_response
+            )
+
+            if errors:
+                print(f"❌ Errors in streaming response: {errors}")
+                return False
+
+            if references is None:
+                print("❌ No references found in streaming response")
+                return False
+
+            if not validate_references_format(references):
+                return False
+
+            if not response_chunks:
+                print("❌ No response chunks found in streaming response")
+                return False
+
+            print(f"✅ Streaming with references: Found {len(references)} references")
+            print(f"   Response chunks: {len(response_chunks)}")
+            print(
+                f"   Total response length: {sum(len(chunk) for chunk in response_chunks)} characters"
+            )
+
+            # Display reference list
+            if references:
+                print("   📚 Reference List:")
+                for i, ref in enumerate(references, 1):
+                    ref_id = ref.get("reference_id", "Unknown")
+                    file_path = ref.get("file_path", "Unknown")
+                    print(f"      {i}. ID: {ref_id} | File: {file_path}")
+
+        else:
+            print(f"❌ Request failed: {response.status_code}")
+            print(f"   Error: {response.text}")
+            return False
+
+    except Exception as e:
+        print(f"❌ Test 1 failed: {str(e)}")
+        return False
+
+    # Test 2: Streaming with references disabled
+    print("\n🧪 Test 2: Streaming with references disabled")
+    print("-" * 40)
+
+    try:
+        response = requests.post(
+            endpoint,
+            json={"query": query_text, "mode": "mix", "include_references": False},
+            headers=AUTH_HEADERS,
+            timeout=30,
+            stream=True,
+        )
+
+        if response.status_code == 200:
+            # Collect streaming response
+            full_response = ""
+            for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
+                if chunk:
+                    # Ensure chunk is string type
+                    if isinstance(chunk, bytes):
+                        chunk = chunk.decode("utf-8")
+                    full_response += chunk
+
+            # Parse streaming response
+            references, response_chunks, errors = parse_streaming_response(
+                full_response
+            )
+
+            if errors:
+                print(f"❌ Errors in streaming response: {errors}")
+                return False
+
+            if references is not None:
+                print("❌ References should be None when include_references=False")
+                return False
+
+            if not response_chunks:
+                print("❌ No response chunks found in streaming response")
+                return False
+
+            print("✅ Streaming without references: No references present")
+            print(f"   Response chunks: {len(response_chunks)}")
+            print(
+                f"   Total response length: {sum(len(chunk) for chunk in response_chunks)} characters"
+            )
+
+        else:
+            print(f"❌ Request failed: {response.status_code}")
+            print(f"   Error: {response.text}")
+            return False
+
+    except Exception as e:
+        print(f"❌ Test 2 failed: {str(e)}")
+        return False
+
+    print("\n✅ /query/stream endpoint references tests passed!")
+    return True
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_references_consistency():
+    """Test references consistency across all endpoints"""
+
+    print("\n" + "=" * 60)
+    print("Testing references consistency across endpoints")
+    print("=" * 60)
+
+    query_text = "who authored LightRAG"
+    query_params = {
+        "query": query_text,
+        "mode": "mix",
+        "top_k": 10,
+        "chunk_top_k": 8,
+        "include_references": True,
+    }
+
+    references_data = {}
+
+    # Test /query endpoint
+    print("\n🧪 Testing /query endpoint")
+    print("-" * 40)
+
+    try:
+        response = requests.post(
+            f"{BASE_URL}/query", json=query_params, headers=AUTH_HEADERS, timeout=30
+        )
+
+        if response.status_code == 200:
+            data = response.json()
+            references_data["query"] = data.get("references", [])
+            print(f"✅ /query: {len(references_data['query'])} references")
+        else:
+            print(f"❌ /query failed: {response.status_code}")
+            return False
+
+    except Exception as e:
+        print(f"❌ /query test failed: {str(e)}")
+        return False
+
+    # Test /query/stream endpoint
+    print("\n🧪 Testing /query/stream endpoint")
+    print("-" * 40)
+
+    try:
+        response = requests.post(
+            f"{BASE_URL}/query/stream",
+            json=query_params,
+            headers=AUTH_HEADERS,
+            timeout=30,
+            stream=True,
+        )
+
+        if response.status_code == 200:
+            full_response = ""
+            for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
+                if chunk:
+                    # Ensure chunk is string type
+                    if isinstance(chunk, bytes):
+                        chunk = chunk.decode("utf-8")
+                    full_response += chunk
+
+            references, _, errors = parse_streaming_response(full_response)
+
+            if errors:
+                print(f"❌ Errors: {errors}")
+                return False
+
+            references_data["stream"] = references or []
+            print(f"✅ /query/stream: {len(references_data['stream'])} references")
+        else:
+            print(f"❌ /query/stream failed: {response.status_code}")
+            return False
+
+    except Exception as e:
+        print(f"❌ /query/stream test failed: {str(e)}")
+        return False
+
+    # Test /query/data endpoint
+    print("\n🧪 Testing /query/data endpoint")
+    print("-" * 40)
+
+    try:
+        response = requests.post(
+            f"{BASE_URL}/query/data",
+            json=query_params,
+            headers=AUTH_HEADERS,
+            timeout=30,
+        )
+
+        if response.status_code == 200:
+            data = response.json()
+            query_data = data.get("data", {})
+            references_data["data"] = query_data.get("references", [])
+            print(f"✅ /query/data: {len(references_data['data'])} references")
+        else:
+            print(f"❌ /query/data failed: {response.status_code}")
+            return False
+
+    except Exception as e:
+        print(f"❌ /query/data test failed: {str(e)}")
+        return False
+
+    # Compare references consistency
+    print("\n🔍 Comparing references consistency")
+    print("-" * 40)
+
+    # Convert to sets of (reference_id, file_path) tuples for comparison
+    def refs_to_set(refs):
+        return set(
+            (ref.get("reference_id", ""), ref.get("file_path", "")) for ref in refs
+        )
+
+    query_refs = refs_to_set(references_data["query"])
+    stream_refs = refs_to_set(references_data["stream"])
+    data_refs = refs_to_set(references_data["data"])
+
+    # Check consistency
+    consistency_passed = True
+
+    if query_refs != stream_refs:
+        print("❌ References mismatch between /query and /query/stream")
+        print(f"   /query only: {query_refs - stream_refs}")
+        print(f"   /query/stream only: {stream_refs - query_refs}")
+        consistency_passed = False
+
+    if query_refs != data_refs:
+        print("❌ References mismatch between /query and /query/data")
+        print(f"   /query only: {query_refs - data_refs}")
+        print(f"   /query/data only: {data_refs - query_refs}")
+        consistency_passed = False
+
+    if stream_refs != data_refs:
+        print("❌ References mismatch between /query/stream and /query/data")
+        print(f"   /query/stream only: {stream_refs - data_refs}")
+        print(f"   /query/data only: {data_refs - stream_refs}")
+        consistency_passed = False
+
+    if consistency_passed:
+        print("✅ All endpoints return consistent references")
+        print(f"   Common references count: {len(query_refs)}")
+
+        # Display common reference list
+        if query_refs:
+            print("   📚 Common Reference List:")
+            for i, (ref_id, file_path) in enumerate(sorted(query_refs), 1):
+                print(f"      {i}. ID: {ref_id} | File: {file_path}")
+
+    return consistency_passed
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_aquery_data_endpoint():
+    """Test the /query/data endpoint"""
+
+    # Use unified configuration
+    endpoint = f"{BASE_URL}/query/data"
+
+    # Query request
+    query_request = {
+        "query": "who authored LighRAG",
+        "mode": "mix",  # Use mixed mode to get the most comprehensive results
+        "top_k": 20,
+        "chunk_top_k": 15,
+        "max_entity_tokens": 4000,
+        "max_relation_tokens": 4000,
+        "max_total_tokens": 16000,
+        "enable_rerank": True,
+    }
+
+    print("=" * 60)
+    print("LightRAG aquery_data endpoint test")
+    print(
+        "   Returns structured data including entities, relationships and text chunks"
+    )
+    print("   Can be used for custom processing and analysis")
+    print("=" * 60)
+    print(f"Query content: {query_request['query']}")
+    print(f"Query mode: {query_request['mode']}")
+    print(f"API endpoint: {endpoint}")
+    print("-" * 60)
+
+    try:
+        # Send request
+        print("Sending request...")
+        start_time = time.time()
+
+        response = requests.post(
+            endpoint, json=query_request, headers=AUTH_HEADERS, timeout=30
+        )
+
+        end_time = time.time()
+        response_time = end_time - start_time
+
+        print(f"Response time: {response_time:.2f} seconds")
+        print(f"HTTP status code: {response.status_code}")
+
+        if response.status_code == 200:
+            data = response.json()
+            print_query_results(data)
+        else:
+            print(f"Request failed: {response.status_code}")
+            print(f"Error message: {response.text}")
+
+    except requests.exceptions.ConnectionError:
+        print("❌ Connection failed: Please ensure LightRAG API service is running")
+        print("   Start command: python -m lightrag.api.lightrag_server")
+    except requests.exceptions.Timeout:
+        print("❌ Request timeout: Query processing took too long")
+    except Exception as e:
+        print(f"❌ Error occurred: {str(e)}")
+
+
+def print_query_results(data: Dict[str, Any]):
+    """Format and print query results"""
+
+    # Check for new data format with status and message
+    status = data.get("status", "unknown")
+    message = data.get("message", "")
+
+    print(f"\n📋 Query Status: {status}")
+    if message:
+        print(f"📋 Message: {message}")
+
+    # Handle new nested data format
+    query_data = data.get("data", {})
+
+    # Fallback to old format if new format is not present
+    if not query_data and any(
+        key in data for key in ["entities", "relationships", "chunks"]
+    ):
+        print("   (Using legacy data format)")
+        query_data = data
+
+    entities = query_data.get("entities", [])
+    relationships = query_data.get("relationships", [])
+    chunks = query_data.get("chunks", [])
+    references = query_data.get("references", [])
+
+    print("\n📊 Query result statistics:")
+    print(f"   Entity count: {len(entities)}")
+    print(f"   Relationship count: {len(relationships)}")
+    print(f"   Text chunk count: {len(chunks)}")
+    print(f"   Reference count: {len(references)}")
+
+    # Print metadata (now at top level in new format)
+    metadata = data.get("metadata", {})
+    if metadata:
+        print("\n🔍 Query metadata:")
+        print(f"   Query mode: {metadata.get('query_mode', 'unknown')}")
+
+        keywords = metadata.get("keywords", {})
+        if keywords:
+            high_level = keywords.get("high_level", [])
+            low_level = keywords.get("low_level", [])
+            if high_level:
+                print(f"   High-level keywords: {', '.join(high_level)}")
+            if low_level:
+                print(f"   Low-level keywords: {', '.join(low_level)}")
+
+        processing_info = metadata.get("processing_info", {})
+        if processing_info:
+            print("   Processing info:")
+            for key, value in processing_info.items():
+                print(f"     {key}: {value}")
+
+    # Print entity information
+    if entities:
+        print("\n👥 Retrieved entities (first 5):")
+        for i, entity in enumerate(entities[:5]):
+            entity_name = entity.get("entity_name", "Unknown")
+            entity_type = entity.get("entity_type", "Unknown")
+            description = entity.get("description", "No description")
+            file_path = entity.get("file_path", "Unknown source")
+            reference_id = entity.get("reference_id", "No reference")
+
+            print(f"   {i + 1}. {entity_name} ({entity_type})")
+            print(
+                f"      Description: {description[:100]}{'...' if len(description) > 100 else ''}"
+            )
+            print(f"      Source: {file_path}")
+            print(f"      Reference ID: {reference_id}")
+            print()
+
+    # Print relationship information
+    if relationships:
+        print("🔗 Retrieved relationships (first 5):")
+        for i, rel in enumerate(relationships[:5]):
+            src = rel.get("src_id", "Unknown")
+            tgt = rel.get("tgt_id", "Unknown")
+            description = rel.get("description", "No description")
+            keywords = rel.get("keywords", "No keywords")
+            file_path = rel.get("file_path", "Unknown source")
+            reference_id = rel.get("reference_id", "No reference")
+
+            print(f"   {i + 1}. {src} → {tgt}")
+            print(f"      Keywords: {keywords}")
+            print(
+                f"      Description: {description[:100]}{'...' if len(description) > 100 else ''}"
+            )
+            print(f"      Source: {file_path}")
+            print(f"      Reference ID: {reference_id}")
+            print()
+
+    # Print text chunk information
+    if chunks:
+        print("📄 Retrieved text chunks (first 3):")
+        for i, chunk in enumerate(chunks[:3]):
+            content = chunk.get("content", "No content")
+            file_path = chunk.get("file_path", "Unknown source")
+            chunk_id = chunk.get("chunk_id", "Unknown ID")
+            reference_id = chunk.get("reference_id", "No reference")
+
+            print(f"   {i + 1}. Text chunk ID: {chunk_id}")
+            print(f"      Source: {file_path}")
+            print(f"      Reference ID: {reference_id}")
+            print(
+                f"      Content: {content[:200]}{'...' if len(content) > 200 else ''}"
+            )
+            print()
+
+    # Print references information (new in updated format)
+    if references:
+        print("📚 References:")
+        for i, ref in enumerate(references):
+            reference_id = ref.get("reference_id", "Unknown ID")
+            file_path = ref.get("file_path", "Unknown source")
+            print(f"   {i + 1}. Reference ID: {reference_id}")
+            print(f"      File Path: {file_path}")
+            print()
+
+    print("=" * 60)
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def compare_with_regular_query():
+    """Compare results between regular query and data query"""
+
+    query_text = "LightRAG的作者是谁"
+
+    print("\n🔄 Comparison test: Regular query vs Data query")
+    print("-" * 60)
+
+    # Regular query
+    try:
+        print("1. Regular query (/query):")
+        regular_response = requests.post(
+            f"{BASE_URL}/query",
+            json={"query": query_text, "mode": "mix"},
+            headers=AUTH_HEADERS,
+            timeout=30,
+        )
+
+        if regular_response.status_code == 200:
+            regular_data = regular_response.json()
+            response_text = regular_data.get("response", "No response")
+            print(
+                f"   Generated answer: {response_text[:300]}{'...' if len(response_text) > 300 else ''}"
+            )
+        else:
+            print(f"   Regular query failed: {regular_response.status_code}")
+            if regular_response.status_code == 403:
+                print("   Authentication failed - Please check API Key configuration")
+            elif regular_response.status_code == 401:
+                print("   Unauthorized - Please check authentication information")
+            print(f"   Error details: {regular_response.text}")
+
+    except Exception as e:
+        print(f"   Regular query error: {str(e)}")
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def run_all_reference_tests():
+    """Run all reference-related tests"""
+
+    print("\n" + "🚀" * 20)
+    print("LightRAG References Test Suite")
+    print("🚀" * 20)
+
+    all_tests_passed = True
+
+    # Test 1: /query endpoint references
+    try:
+        if not test_query_endpoint_references():
+            all_tests_passed = False
+    except Exception as e:
+        print(f"❌ /query endpoint test failed with exception: {str(e)}")
+        all_tests_passed = False
+
+    # Test 2: /query/stream endpoint references
+    try:
+        if not test_query_stream_endpoint_references():
+            all_tests_passed = False
+    except Exception as e:
+        print(f"❌ /query/stream endpoint test failed with exception: {str(e)}")
+        all_tests_passed = False
+
+    # Test 3: References consistency across endpoints
+    try:
+        if not test_references_consistency():
+            all_tests_passed = False
+    except Exception as e:
+        print(f"❌ References consistency test failed with exception: {str(e)}")
+        all_tests_passed = False
+
+    # Final summary
+    print("\n" + "=" * 60)
+    print("TEST SUITE SUMMARY")
+    print("=" * 60)
+
+    if all_tests_passed:
+        print("🎉 ALL TESTS PASSED!")
+        print("✅ /query endpoint references functionality works correctly")
+        print("✅ /query/stream endpoint references functionality works correctly")
+        print("✅ References are consistent across all endpoints")
+        print("\n🔧 System is ready for production use with reference support!")
+    else:
+        print("❌ SOME TESTS FAILED!")
+        print("Please check the error messages above and fix the issues.")
+        print("\n🔧 System needs attention before production deployment.")
+
+    return all_tests_passed
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) > 1 and sys.argv[1] == "--references-only":
+        # Run only the new reference tests
+        success = run_all_reference_tests()
+        sys.exit(0 if success else 1)
+    else:
+        # Run original tests plus new reference tests
+        print("Running original aquery_data endpoint test...")
+        test_aquery_data_endpoint()
+
+        print("\nRunning comparison test...")
+        compare_with_regular_query()
+
+        print("\nRunning new reference tests...")
+        run_all_reference_tests()
+
+        print("\n💡 Usage tips:")
+        print("1. Ensure LightRAG API service is running")
+        print("2. Adjust base_url and authentication information as needed")
+        print("3. Modify query parameters to test different retrieval strategies")
+        print("4. Data query results can be used for further analysis and processing")
+        print("5. Run with --references-only flag to test only reference functionality")

tests/api/routes/test_curl_aquery_data.sh

@@ -0,0 +1,271 @@
+#!/bin/bash
+
+# LightRAG aquery_data endpoint test script
+# Use curl command to test the new /query/data endpoint and validate the new data format
+
+echo "🚀 LightRAG aquery_data Endpoint Test (New Data Format Validation)"
+echo "=================================================="
+
+# Base URL (adjust according to actual deployment)
+BASE_URL="http://localhost:9621"
+
+# Color definitions
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+# Test result statistics
+TOTAL_TESTS=0
+PASSED_TESTS=0
+FAILED_TESTS=0
+
+# Function to validate success response format
+validate_success_response() {
+    local response="$1"
+    local test_name="$2"
+    local expected_mode="$3"
+
+    echo -e "${BLUE}Validating $test_name response format...${NC}"
+
+    # Check if valid JSON
+    if ! echo "$response" | jq . >/dev/null 2>&1; then
+        echo -e "${RED}❌ Response is not valid JSON format${NC}"
+        return 1
+    fi
+
+    # Validate required fields
+    local status=$(echo "$response" | jq -r '.status // "missing"')
+    local message=$(echo "$response" | jq -r '.message // "missing"')
+    local data_exists=$(echo "$response" | jq 'has("data")')
+    local metadata_exists=$(echo "$response" | jq 'has("metadata")')
+
+    echo "  Status: $status"
+    echo "  Message: $message"
+
+    # Validate data structure
+    if [[ "$data_exists" == "true" ]]; then
+        local entities_count=$(echo "$response" | jq '.data.entities | length // 0')
+        local relationships_count=$(echo "$response" | jq '.data.relationships | length // 0')
+        local chunks_count=$(echo "$response" | jq '.data.chunks | length // 0')
+        local references_count=$(echo "$response" | jq '.data.references | length // 0')
+
+        echo "  Data.entities: $entities_count"
+        echo "  Data.relationships: $relationships_count"
+        echo "  Data.chunks: $chunks_count"
+        echo "  Data.references: $references_count"
+    else
+        echo -e "${RED}  ❌ Missing 'data' field${NC}"
+        return 1
+    fi
+
+    # Validate metadata
+    if [[ "$metadata_exists" == "true" ]]; then
+        local query_mode=$(echo "$response" | jq -r '.metadata.query_mode // "missing"')
+        local keywords_exists=$(echo "$response" | jq 'has("metadata") and (.metadata | has("keywords"))')
+        local processing_info_exists=$(echo "$response" | jq 'has("metadata") and (.metadata | has("processing_info"))')
+
+        echo "  Metadata.query_mode: $query_mode"
+        echo "  Metadata.keywords: $keywords_exists"
+        echo "  Metadata.processing_info: $processing_info_exists"
+
+        # Validate if query mode matches
+        if [[ "$expected_mode" != "" && "$query_mode" != "$expected_mode" ]]; then
+            echo -e "${YELLOW}  ⚠️  Query mode mismatch: expected '$expected_mode', actual '$query_mode'${NC}"
+        fi
+    else
+        echo -e "${RED}  ❌ Missing 'metadata' field${NC}"
+        return 1
+    fi
+
+    # Validate status
+    if [[ "$status" == "success" ]]; then
+        echo -e "${GREEN}  ✅ Response format validation passed${NC}"
+        return 0
+    else
+        echo -e "${RED}  ❌ Status is not 'success': $status${NC}"
+        return 1
+    fi
+}
+
+# Function to validate error response format
+validate_error_response() {
+    local response="$1"
+    local test_name="$2"
+
+    echo -e "${BLUE}Validating $test_name response format...${NC}"
+
+    # Check if valid JSON
+    if ! echo "$response" | jq . >/dev/null 2>&1; then
+        echo -e "${RED}❌ Response is not valid JSON format${NC}"
+        return 1
+    fi
+
+    # Validate required fields
+    local status=$(echo "$response" | jq -r '.status // "missing"')
+    local message=$(echo "$response" | jq -r '.message // "missing"')
+    local data_exists=$(echo "$response" | jq 'has("data")')
+    local metadata_exists=$(echo "$response" | jq 'has("metadata")')
+
+    echo "  Status: $status"
+    echo "  Message: $message"
+
+    # Validate basic structure exists
+    if [[ "$data_exists" != "true" ]]; then
+        echo -e "${RED}  ❌ Missing 'data' field${NC}"
+        return 1
+    fi
+
+    if [[ "$metadata_exists" != "true" ]]; then
+        echo -e "${RED}  ❌ Missing 'metadata' field${NC}"
+        return 1
+    fi
+
+    echo "  Data: {}"
+    echo "  Metadata: {}"
+
+    # Validate status should be failure
+    if [[ "$status" == "failure" ]]; then
+        echo -e "${GREEN}  ✅ Error response format validation passed${NC}"
+        return 0
+    else
+        echo -e "${RED}  ❌ Status is not 'failure': $status${NC}"
+        return 1
+    fi
+}
+
+# Function to run success test
+run_success_test() {
+    local test_name="$1"
+    local query_data="$2"
+    local expected_mode="$3"
+    local print_json="${4:-false}"  # Optional parameter: whether to print JSON response (default: false)
+
+    echo ""
+    echo "=================================="
+    echo -e "${BLUE}$test_name${NC}"
+    echo "=================================="
+
+    TOTAL_TESTS=$((TOTAL_TESTS + 1))
+
+    # Send request
+    echo "Sending request..."
+    local response=$(curl -s -X POST "${BASE_URL}/query/data" \
+      -H "Content-Type: application/json" \
+      -H "X-API-Key: your-secure-api-key-here-123" \
+      -d "$query_data")
+
+    # Check if curl succeeded
+    if [[ $? -ne 0 ]]; then
+        echo -e "${RED}❌ Request failed - cannot connect to server${NC}"
+        FAILED_TESTS=$((FAILED_TESTS + 1))
+        return 1
+    fi
+
+    # Print JSON response if requested
+    if [[ "$print_json" == "true" ]]; then
+        echo ""
+        echo "Response JSON:"
+        echo "$response" | jq '.' 2>/dev/null || echo "$response"
+        echo ""
+    fi
+
+    # Validate response
+    if validate_success_response "$response" "$test_name" "$expected_mode"; then
+        PASSED_TESTS=$((PASSED_TESTS + 1))
+        echo -e "${GREEN}✅ $test_name test passed${NC}"
+    else
+        FAILED_TESTS=$((FAILED_TESTS + 1))
+        echo -e "${RED}❌ $test_name test failed${NC}"
+        echo "Raw response:"
+        echo "$response" | jq '.' 2>/dev/null || echo "$response"
+    fi
+}
+
+# Function to run error test
+run_error_test() {
+    local test_name="$1"
+    local query_data="$2"
+
+    echo ""
+    echo "=================================="
+    echo -e "${BLUE}$test_name${NC}"
+    echo "=================================="
+
+    TOTAL_TESTS=$((TOTAL_TESTS + 1))
+
+    # Send request
+    echo "Sending request..."
+    local response=$(curl -s -X POST "${BASE_URL}/query/data" \
+      -H "Content-Type: application/json" \
+      -H "X-API-Key: your-secure-api-key-here-123" \
+      -d "$query_data")
+
+    # Check if curl succeeded
+    if [[ $? -ne 0 ]]; then
+        echo -e "${RED}❌ Request failed - cannot connect to server${NC}"
+        FAILED_TESTS=$((FAILED_TESTS + 1))
+        return 1
+    fi
+
+    # Validate response
+    if validate_error_response "$response" "$test_name"; then
+        PASSED_TESTS=$((PASSED_TESTS + 1))
+        echo -e "${GREEN}✅ $test_name test passed${NC}"
+    else
+        FAILED_TESTS=$((FAILED_TESTS + 1))
+        echo -e "${RED}❌ $test_name test failed${NC}"
+        echo "Raw response:"
+        echo "$response" | jq '.' 2>/dev/null || echo "$response"
+    fi
+}
+
+# Start tests
+echo "Starting tests for new /query/data endpoint data format..."
+echo ""
+
+# Test 1: Basic query test (mix mode)
+run_success_test "1. Basic Query Test (mix mode)" '{
+    "query": "What is GraphRAG",
+    "mode": "mix",
+    "top_k": 5
+}' "mix" "true" # Output full JSON
+
+# Test 2: Detailed parameter query test (hybrid mode)
+run_success_test "2. Detailed Parameter Query Test (hybrid mode)" '{
+    "query": "What is GraphRAG",
+    "mode": "hybrid",
+    "top_k": 5,
+    "chunk_top_k": 8,
+    "max_entity_tokens": 4000,
+    "max_relation_tokens": 4000,
+    "max_total_tokens": 16000,
+    "enable_rerank": true,
+    "response_type": "Multiple Paragraphs"
+}' "hybrid"
+
+# Output test result statistics
+echo ""
+echo "=================================================="
+echo -e "${BLUE}Test Result Statistics${NC}"
+echo "=================================================="
+echo -e "Total tests: ${BLUE}$TOTAL_TESTS${NC}"
+echo -e "Passed tests: ${GREEN}$PASSED_TESTS${NC}"
+echo -e "Failed tests: ${RED}$FAILED_TESTS${NC}"
+
+if [[ $FAILED_TESTS -eq 0 ]]; then
+    echo -e "${GREEN}🎉 All tests passed! New data format adaptation successful!${NC}"
+    exit 0
+else
+    echo -e "${RED}⚠️  $FAILED_TESTS test(s) failed, please check the issues${NC}"
+    exit 1
+fi
+
+echo ""
+echo "💡 Usage Instructions:"
+echo "1. Ensure LightRAG API service is running (python -m lightrag.api.lightrag_server)"
+echo "2. Adjust BASE_URL as needed"
+echo "3. If authentication is required, add -H \"Authorization: Bearer your-token\""
+echo "4. Install jq for better JSON formatting output: brew install jq (macOS) or apt install jq (Ubuntu)"
+echo "5. Script will automatically validate new data format structure: status, message, data, metadata"

tests/api/routes/test_description_api_validation.py

@@ -0,0 +1,403 @@
+import pytest
+
+from lightrag.constants import SOURCE_IDS_LIMIT_METHOD_KEEP
+from lightrag.constants import GRAPH_FIELD_SEP
+from lightrag.operate import (
+    _handle_single_entity_extraction,
+    _merge_nodes_then_upsert,
+    _normalize_text_extraction_record_attributes,
+    _handle_single_relationship_extraction,
+)
+from lightrag import utils_graph
+from lightrag.utils import VectorStorageConsistencyError
+
+
+class DummyGraphStorage:
+    def __init__(self, node=None):
+        self.node = node
+        self.upserted_nodes = []
+
+    async def get_node(self, node_id):
+        return self.node
+
+    async def upsert_node(self, node_id, node_data):
+        self.upserted_nodes.append((node_id, node_data))
+        self.node = dict(node_data)
+
+
+class DummyVectorStorage:
+    def __init__(self):
+        self.global_config = {"workspace": "test"}
+        self.upserts = []
+        self.deletes = []
+
+    async def upsert(self, data):
+        self.upserts.append(data)
+        return None
+
+    async def delete(self, ids):
+        self.deletes.append(ids)
+        return None
+
+    async def get_by_id(self, id_):
+        return None
+
+    async def index_done_callback(self):
+        return True
+
+
+class DummyAsyncContext:
+    async def __aenter__(self):
+        return None
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+
+class DummyMergeGraphStorage:
+    def __init__(self):
+        self.nodes = {
+            "Canonical": {
+                "entity_id": "Canonical",
+                "description": "canonical desc",
+                "entity_type": "ORG",
+                "source_id": "chunk-1",
+                "file_path": "canonical.md",
+            },
+            "Alias": {
+                "entity_id": "Alias",
+                "description": "alias desc",
+                "entity_type": "ORG",
+                "source_id": "chunk-2",
+                "file_path": "alias.md",
+            },
+            "Neighbor": {
+                "entity_id": "Neighbor",
+                "description": "neighbor desc",
+                "entity_type": "ORG",
+                "source_id": "chunk-3",
+                "file_path": "neighbor.md",
+            },
+        }
+        self.edges = {
+            ("Alias", "Neighbor"): {
+                "description": "rel desc",
+                "keywords": "alias",
+                "source_id": "chunk-rel",
+                "weight": 1.0,
+                "file_path": "rel.md",
+            }
+        }
+
+    async def has_node(self, node_id):
+        return node_id in self.nodes
+
+    async def get_node(self, node_id):
+        return self.nodes[node_id]
+
+    async def upsert_node(self, node_id, node_data):
+        self.nodes[node_id] = dict(node_data)
+
+    async def get_node_edges(self, node_id):
+        results = []
+        for src, tgt in self.edges:
+            if src == node_id or tgt == node_id:
+                results.append((src, tgt))
+        return results
+
+    async def get_edge(self, src, tgt):
+        return self.edges.get((src, tgt)) or self.edges.get((tgt, src))
+
+    async def upsert_edge(self, src, tgt, edge_data):
+        self.edges[(src, tgt)] = dict(edge_data)
+
+    async def delete_node(self, node_id):
+        self.nodes.pop(node_id, None)
+        self.edges = {
+            (src, tgt): data
+            for (src, tgt), data in self.edges.items()
+            if src != node_id and tgt != node_id
+        }
+
+    async def index_done_callback(self):
+        return True
+
+
+@pytest.mark.asyncio
+async def test_merge_nodes_then_upsert_handles_missing_legacy_description():
+    graph = DummyGraphStorage(node={"source_id": "chunk-1"})
+    global_config = {
+        "source_ids_limit_method": SOURCE_IDS_LIMIT_METHOD_KEEP,
+        "max_source_ids_per_entity": 20,
+    }
+
+    result = await _merge_nodes_then_upsert(
+        entity_name="LegacyEntity",
+        nodes_data=[],
+        knowledge_graph_inst=graph,
+        entity_vdb=None,
+        global_config=global_config,
+    )
+
+    assert result["description"] == "Entity LegacyEntity"
+    assert graph.upserted_nodes[-1][1]["description"] == "Entity LegacyEntity"
+
+
+@pytest.mark.asyncio
+async def test_acreate_entity_rejects_empty_description():
+    with pytest.raises(ValueError, match="description cannot be empty"):
+        await utils_graph.acreate_entity(
+            chunk_entity_relation_graph=None,
+            entities_vdb=None,
+            relationships_vdb=None,
+            entity_name="EntityA",
+            entity_data={"description": "   "},
+        )
+
+
+@pytest.mark.asyncio
+async def test_acreate_relation_rejects_empty_description():
+    with pytest.raises(ValueError, match="description cannot be empty"):
+        await utils_graph.acreate_relation(
+            chunk_entity_relation_graph=None,
+            entities_vdb=None,
+            relationships_vdb=None,
+            source_entity="A",
+            target_entity="B",
+            relation_data={"description": ""},
+        )
+
+
+@pytest.mark.asyncio
+async def test_aedit_entity_rejects_empty_description():
+    with pytest.raises(ValueError, match="description cannot be empty"):
+        await utils_graph.aedit_entity(
+            chunk_entity_relation_graph=None,
+            entities_vdb=None,
+            relationships_vdb=None,
+            entity_name="EntityA",
+            updated_data={"description": None},
+        )
+
+
+@pytest.mark.asyncio
+async def test_aedit_relation_rejects_empty_description():
+    with pytest.raises(ValueError, match="description cannot be empty"):
+        await utils_graph.aedit_relation(
+            chunk_entity_relation_graph=None,
+            entities_vdb=None,
+            relationships_vdb=None,
+            source_entity="A",
+            target_entity="B",
+            updated_data={"description": "   "},
+        )
+
+
+@pytest.mark.asyncio
+async def test_aedit_entity_allows_updates_without_description(monkeypatch):
+    async def fake_edit_impl(*args, **kwargs):
+        return {"entity_name": "EntityA", "description": "kept", "source_id": "chunk-1"}
+
+    monkeypatch.setattr(
+        utils_graph, "get_storage_keyed_lock", lambda *a, **k: DummyAsyncContext()
+    )
+    monkeypatch.setattr(utils_graph, "_edit_entity_impl", fake_edit_impl)
+
+    result = await utils_graph.aedit_entity(
+        chunk_entity_relation_graph=None,
+        entities_vdb=DummyVectorStorage(),
+        relationships_vdb=DummyVectorStorage(),
+        entity_name="EntityA",
+        updated_data={"entity_type": "ORG"},
+    )
+
+    assert result["operation_summary"]["operation_status"] == "success"
+
+
+def test_handle_single_relationship_extraction_ignores_empty_description():
+    relation = _handle_single_relationship_extraction(
+        ["relation", "Alice", "Bob", "works_with", "   "],
+        chunk_key="chunk-1",
+        timestamp=1,
+    )
+
+    assert relation is None
+
+
+def test_mis_prefixed_relation_row_is_recovered():
+    record = _normalize_text_extraction_record_attributes(
+        ["entity", "Alice", "Acme Corp", "founded", "Alice founded Acme Corp."],
+        chunk_key="chunk-1",
+    )
+
+    relation = _handle_single_relationship_extraction(
+        record,
+        chunk_key="chunk-1",
+        timestamp=1,
+    )
+
+    assert relation is not None
+    assert relation["src_id"] == "Alice"
+    assert relation["tgt_id"] == "Acme Corp"
+
+
+def test_four_part_entity_row_remains_entity():
+    record = _normalize_text_extraction_record_attributes(
+        ["entity", "Alice", "Person", "Alice is the founder of Acme Corp."],
+        chunk_key="chunk-1",
+    )
+
+    entity = _handle_single_entity_extraction(
+        record,
+        chunk_key="chunk-1",
+        timestamp=1,
+    )
+
+    assert entity is not None
+    assert entity["entity_name"] == "Alice"
+
+
+def test_malformed_recovered_relation_still_fails():
+    record = _normalize_text_extraction_record_attributes(
+        ["entity", "Alice", "Acme Corp", "founded", "   "],
+        chunk_key="chunk-1",
+    )
+
+    relation = _handle_single_relationship_extraction(
+        record,
+        chunk_key="chunk-1",
+        timestamp=1,
+    )
+
+    assert relation is None
+
+
+def test_unrelated_five_part_prefix_remains_invalid():
+    record = _normalize_text_extraction_record_attributes(
+        ["edge", "Alice", "Acme Corp", "founded", "Alice founded Acme Corp."],
+        chunk_key="chunk-1",
+    )
+
+    relation = _handle_single_relationship_extraction(
+        record,
+        chunk_key="chunk-1",
+        timestamp=1,
+    )
+
+    assert relation is None
+
+
+@pytest.mark.asyncio
+async def test_merge_entities_preserves_file_path_in_vector_updates(monkeypatch):
+    graph = DummyMergeGraphStorage()
+    entities_vdb = DummyVectorStorage()
+    relationships_vdb = DummyVectorStorage()
+
+    async def fake_get_entity_info(*args, **kwargs):
+        return {"entity_name": "Canonical"}
+
+    monkeypatch.setattr(utils_graph, "get_entity_info", fake_get_entity_info)
+
+    await utils_graph._merge_entities_impl(
+        chunk_entity_relation_graph=graph,
+        entities_vdb=entities_vdb,
+        relationships_vdb=relationships_vdb,
+        source_entities=["Alias", "Canonical"],
+        target_entity="Canonical",
+    )
+
+    relationship_payload = relationships_vdb.upserts[-1]
+    entity_payload = entities_vdb.upserts[-1]
+
+    assert next(iter(relationship_payload.values()))["file_path"] == "rel.md"
+    assert set(
+        next(iter(entity_payload.values()))["file_path"].split(GRAPH_FIELD_SEP)
+    ) == {
+        "alias.md",
+        "canonical.md",
+    }
+
+
+@pytest.mark.asyncio
+async def test_merge_response_is_built_from_graph_without_reading_vdb():
+    # The merge response is built from graph data only. Reading the vector
+    # store is redundant (the graph is authoritative) and previously leaked a
+    # non-JSON-serializable embedding (pgvector -> numpy) into the API
+    # response, causing a 500 on /graph/entity/edit with all-PostgreSQL.
+    graph = DummyMergeGraphStorage()
+    entities_vdb = DummyVectorStorage()
+    relationships_vdb = DummyVectorStorage()
+
+    vdb_reads = []
+    inner_get_by_id = entities_vdb.get_by_id
+
+    async def _tracked_get_by_id(id_):
+        vdb_reads.append(id_)
+        return await inner_get_by_id(id_)
+
+    entities_vdb.get_by_id = _tracked_get_by_id
+
+    result = await utils_graph._merge_entities_impl(
+        chunk_entity_relation_graph=graph,
+        entities_vdb=entities_vdb,
+        relationships_vdb=relationships_vdb,
+        source_entities=["Alias", "Canonical"],
+        target_entity="Canonical",
+    )
+
+    assert "graph_data" in result
+    assert "vector_data" not in result
+    # The response did not read the entity vector store.
+    assert vdb_reads == []
+
+
+@pytest.mark.asyncio
+async def test_aedit_entity_merge_propagates_consistency_error(monkeypatch):
+    """rename-with-merge must surface VectorStorageConsistencyError, not swallow it.
+
+    The graph is updated during the merge, then the (deferred) vector-store
+    flush fails. aedit_entity used to fold every merge exception into a
+    partial-success summary and return normally, so the /graph/entity/edit
+    route reported HTTP 200 "success" and hid the consistency failure that the
+    fail-loud merge path exists to surface. It must re-raise instead.
+    """
+    graph = DummyMergeGraphStorage()
+    entities_vdb = DummyVectorStorage()
+    relationships_vdb = DummyVectorStorage()
+
+    # Deferred-embedding flush failure: upsert buffers, index_done_callback raises.
+    async def _boom():
+        raise RuntimeError("embedder down")
+
+    relationships_vdb.index_done_callback = _boom
+
+    # Single-attempt VDB wrapper (no retry delays) + a no-op storage lock.
+    async def _single_attempt(operation, **kwargs):
+        await operation()
+
+    monkeypatch.setattr(
+        utils_graph, "safe_vdb_operation_with_exception", _single_attempt
+    )
+    monkeypatch.setattr(
+        utils_graph, "get_storage_keyed_lock", lambda *a, **k: DummyAsyncContext()
+    )
+
+    async def fake_get_entity_info(*args, **kwargs):
+        return {"entity_name": "Canonical"}
+
+    monkeypatch.setattr(utils_graph, "get_entity_info", fake_get_entity_info)
+
+    with pytest.raises(VectorStorageConsistencyError) as excinfo:
+        await utils_graph.aedit_entity(
+            chunk_entity_relation_graph=graph,
+            entities_vdb=entities_vdb,
+            relationships_vdb=relationships_vdb,
+            entity_name="Alias",
+            updated_data={"entity_name": "Canonical"},
+            allow_rename=True,
+            allow_merge=True,
+        )
+
+    assert "lightrag-rebuild-vdb" in str(excinfo.value)
+    # Fail-loud happened before source deletion: 'Alias' is still in the graph.
+    assert "Alias" in graph.nodes

tests/api/routes/test_document_routes_chunking.py

@@ -0,0 +1,694 @@
+"""Tests for the `/documents/text(s)` ``chunking`` request object.
+
+Three concerns:
+
+1. **Synchronous validation**: malformed ``chunking`` is rejected at
+   request-parse time (HTTP 422 / ``ValidationError``) — never deferred to
+   the background indexing task, where the HTTP response is already sent.
+   The per-strategy typed params models do full type + value checking, not
+   just unknown-key detection.
+
+2. **``_resolve_text_chunking``**: a validated ``chunking`` config is frozen
+   into ``(process_options, chunk_options)``; ``chunk_token_size`` and the
+   strategy params land in the selected strategy's sub-dict, overriding any
+   env-derived value, while the other strategy sub-dicts are dropped (slim).
+
+3. **Route forwarding**: ``/documents/text`` and ``/documents/texts`` forward
+   ``request.chunking`` to ``pipeline_index_texts`` and return 422 (without
+   scheduling any background work) for a malformed body.
+"""
+
+import importlib
+import sys
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from pydantic import ValidationError
+
+_original_argv = sys.argv[:]
+sys.argv = [sys.argv[0]]
+_dr = importlib.import_module("lightrag.api.routers.document_routes")
+sys.argv = _original_argv
+
+TextChunkingConfig = _dr.TextChunkingConfig
+InsertTextRequest = _dr.InsertTextRequest
+_resolve_text_chunking = _dr._resolve_text_chunking
+create_document_routes = _dr.create_document_routes
+
+from lightrag.constants import (  # noqa: E402
+    PROCESS_OPTION_CHUNK_FIXED,
+    PROCESS_OPTION_CHUNK_PARAGRAH,
+    PROCESS_OPTION_CHUNK_RECURSIVE,
+    PROCESS_OPTION_CHUNK_VECTOR,
+)
+from lightrag.parser.routing import default_chunker_config  # noqa: E402
+
+pytestmark = pytest.mark.offline
+
+_ALL_STRATEGY_KEYS = {
+    "fixed_token",
+    "recursive_character",
+    "semantic_vector",
+    "paragraph_semantic",
+}
+
+
+# ---------------------------------------------------------------------------
+# 1. Synchronous validation
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "body",
+    [
+        # wrong types (strict rejects lax coercion)
+        {"strategy": "fixed_token", "params": {"chunk_token_size": True}},
+        {"strategy": "fixed_token", "params": {"chunk_token_size": "5"}},
+        {"strategy": "fixed_token", "params": {"chunk_token_size": 1.5}},
+        {"strategy": "fixed_token", "params": {"chunk_overlap_token_size": "bad"}},
+        {"strategy": "fixed_token", "params": {"split_by_character": 123}},
+        {"strategy": "fixed_token", "params": {"split_by_character_only": 1}},
+        {"strategy": "recursive_character", "params": {"separators": "abc"}},
+        {"strategy": "recursive_character", "params": {"separators": [1, 2]}},
+        # value / range
+        {"strategy": "fixed_token", "params": {"chunk_token_size": 0}},
+        {"strategy": "recursive_character", "params": {"chunk_overlap_token_size": -1}},
+        {"strategy": "semantic_vector", "params": {"buffer_size": 0}},
+        {"strategy": "semantic_vector", "params": {"buffer_size": True}},
+        {
+            "strategy": "semantic_vector",
+            "params": {"breakpoint_threshold_type": "p99"},
+        },
+        {
+            "strategy": "semantic_vector",
+            "params": {"breakpoint_threshold_amount": 0},
+        },
+        {
+            # strict float rejects strings (no lax numeric-string coercion)
+            "strategy": "semantic_vector",
+            "params": {"breakpoint_threshold_amount": "95"},
+        },
+        {
+            # strict float rejects bool (bool is an int subclass, undesirable here)
+            "strategy": "semantic_vector",
+            "params": {"breakpoint_threshold_amount": True},
+        },
+        {
+            # > 100 with an explicit percentile/gradient type is rejected at
+            # parse time (both fields present, no inheritance ambiguity).
+            "strategy": "semantic_vector",
+            "params": {
+                "breakpoint_threshold_type": "percentile",
+                "breakpoint_threshold_amount": 150,
+            },
+        },
+        {
+            # malformed regex must be compiled/rejected at parse time
+            "strategy": "semantic_vector",
+            "params": {"sentence_split_regex": "("},
+        },
+        # cross-field
+        {
+            "strategy": "fixed_token",
+            "params": {"chunk_token_size": 100, "chunk_overlap_token_size": 200},
+        },
+        # unknown / wrong-for-strategy keys
+        {"strategy": "fixed_token", "params": {"bogus": 1}},
+        {"strategy": "fixed_token", "params": {"separators": ["x"]}},
+        {"strategy": "recursive_character", "params": {"buffer_size": 1}},
+    ],
+)
+def test_chunking_config_rejects_malformed(body):
+    with pytest.raises(ValidationError):
+        TextChunkingConfig.model_validate(body)
+
+
+def test_chunking_config_defaults_to_fixed_token():
+    cfg = TextChunkingConfig.model_validate({"params": {"chunk_token_size": 500}})
+    assert cfg.strategy == "fixed_token"
+    assert cfg.params == {"chunk_token_size": 500}
+
+
+def test_chunking_config_normalizes_to_supplied_keys_only():
+    # int amount is coerced to float; only the supplied key survives.
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 95}}
+    )
+    assert cfg.params == {"breakpoint_threshold_amount": 95.0}
+
+
+def test_chunking_config_amount_in_range_for_std_deviation():
+    # standard_deviation only requires > 0 (no [0, 100] ceiling).
+    cfg = TextChunkingConfig.model_validate(
+        {
+            "strategy": "semantic_vector",
+            "params": {
+                "breakpoint_threshold_type": "standard_deviation",
+                "breakpoint_threshold_amount": 3.5,
+            },
+        }
+    )
+    assert cfg.params["breakpoint_threshold_amount"] == 3.5
+
+
+def test_chunking_config_amount_over_100_without_type_is_deferred():
+    # Type omitted -> the (0, 100] ceiling cannot be decided at parse time
+    # (the effective type may be inherited), so the model must NOT assume
+    # percentile and reject. _resolve_text_chunking applies the ceiling later.
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 150}}
+    )
+    assert cfg.params == {"breakpoint_threshold_amount": 150.0}
+
+
+def test_chunking_config_accepts_int_amount_widened_to_float():
+    # Strict float accepts an int (JSON 95) and widens it to 95.0 — the common
+    # documented threshold magnitude. (str/bool are rejected; see the
+    # rejection matrix above.) Exercised via both python and JSON validation
+    # modes so the FastAPI request path (which parses JSON) stays covered.
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 95}}
+    )
+    assert cfg.params == {"breakpoint_threshold_amount": 95.0}
+    assert isinstance(cfg.params["breakpoint_threshold_amount"], float)
+
+    cfg_json = TextChunkingConfig.model_validate_json(
+        '{"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 95}}'
+    )
+    assert cfg_json.params == {"breakpoint_threshold_amount": 95.0}
+
+
+def test_chunking_config_accepts_valid_sentence_split_regex():
+    cfg = TextChunkingConfig.model_validate(
+        {
+            "strategy": "semantic_vector",
+            "params": {"sentence_split_regex": r"(?<=[.?!])\s+"},
+        }
+    )
+    assert cfg.params == {"sentence_split_regex": r"(?<=[.?!])\s+"}
+
+
+def test_chunking_config_drops_explicit_null():
+    # Explicit null means "inherit the default" (every param field is
+    # Optional/None=inherit), so it must be dropped — not merged over the
+    # resolved default, which would later make the chunker do int(None).
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "fixed_token", "params": {"chunk_token_size": None}}
+    )
+    assert cfg.params == {}
+
+
+def test_chunking_config_keeps_real_value_drops_sibling_null():
+    cfg = TextChunkingConfig.model_validate(
+        {
+            "strategy": "fixed_token",
+            "params": {"chunk_token_size": 500, "split_by_character": None},
+        }
+    )
+    assert cfg.params == {"chunk_token_size": 500}
+
+
+def test_insert_text_request_rejects_malformed_chunking():
+    with pytest.raises(ValidationError):
+        InsertTextRequest.model_validate(
+            {
+                "text": "hi",
+                "file_source": "a.md",
+                "chunking": {
+                    "strategy": "recursive_character",
+                    "params": {"separators": "notalist"},
+                },
+            }
+        )
+
+
+# ---------------------------------------------------------------------------
+# 2. _resolve_text_chunking
+# ---------------------------------------------------------------------------
+
+
+def _stub_rag(addon_params=None):
+    return SimpleNamespace(
+        addon_params=addon_params if addon_params is not None else {}
+    )
+
+
+def test_resolve_none_keeps_default_fixed():
+    process_options, chunk_options = _resolve_text_chunking(None, _stub_rag())
+    assert process_options == PROCESS_OPTION_CHUNK_FIXED
+    assert "fixed_token" in chunk_options
+
+
+@pytest.mark.parametrize(
+    "strategy,expected_po,key",
+    [
+        ("fixed_token", PROCESS_OPTION_CHUNK_FIXED, "fixed_token"),
+        ("recursive_character", PROCESS_OPTION_CHUNK_RECURSIVE, "recursive_character"),
+        ("semantic_vector", PROCESS_OPTION_CHUNK_VECTOR, "semantic_vector"),
+        ("paragraph_semantic", PROCESS_OPTION_CHUNK_PARAGRAH, "paragraph_semantic"),
+    ],
+)
+def test_resolve_maps_strategy_and_writes_size_into_subdict(strategy, expected_po, key):
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": strategy, "params": {"chunk_token_size": 777}}
+    )
+    process_options, chunk_options = _resolve_text_chunking(cfg, _stub_rag())
+    assert process_options == expected_po
+    # chunk_token_size lands in the strategy sub-dict for ALL strategies
+    # (F included, post-cleanup) — that's where process_single_document reads it.
+    assert chunk_options[key]["chunk_token_size"] == 777
+    # slim contract: other strategies' sub-dicts are dropped
+    for other in _ALL_STRATEGY_KEYS - {key}:
+        assert other not in chunk_options
+
+
+def test_resolve_merges_strategy_params():
+    cfg = TextChunkingConfig.model_validate(
+        {
+            "strategy": "recursive_character",
+            "params": {"separators": ["A", "B"], "chunk_overlap_token_size": 0},
+        }
+    )
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag())
+    assert chunk_options["recursive_character"]["separators"] == ["A", "B"]
+    assert chunk_options["recursive_character"]["chunk_overlap_token_size"] == 0
+
+
+def test_resolve_size_overrides_env_for_recursive(monkeypatch):
+    monkeypatch.setenv("CHUNK_R_SIZE", "999")
+    addon = {"chunker": default_chunker_config()}
+    # sanity: env baked into the R sub-dict
+    assert addon["chunker"]["recursive_character"]["chunk_token_size"] == 999
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "recursive_character", "params": {"chunk_token_size": 1234}}
+    )
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
+    # API value wins over the env-derived sub-dict value.
+    assert chunk_options["recursive_character"]["chunk_token_size"] == 1234
+
+
+def test_resolve_split_by_character_only_false_overrides_env(monkeypatch):
+    # The API path can express an explicit False (a plain dict merge), unlike
+    # the ainsert positional-arg path. Prove it overrides an env-True default.
+    monkeypatch.setenv("CHUNK_F_SPLIT_BY_CHARACTER_ONLY", "true")
+    addon = {"chunker": default_chunker_config()}
+    assert addon["chunker"]["fixed_token"]["split_by_character_only"] is True
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "fixed_token", "params": {"split_by_character_only": False}}
+    )
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
+    assert chunk_options["fixed_token"]["split_by_character_only"] is False
+
+
+def test_resolve_drop_references_request_false_overrides_env_true(monkeypatch):
+    # The JSON text API can express an explicit ``drop_references=False`` that
+    # overrides ``CHUNK_P_DROP_REFERENCES=true`` (resolved via slim_chunk_options
+    # then overlaid by the request params). Mirrors the split_by_character_only
+    # override above for the paragraph-semantic switch.
+    monkeypatch.setenv("CHUNK_P_DROP_REFERENCES", "true")
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "paragraph_semantic", "params": {"drop_references": False}}
+    )
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag())
+    assert chunk_options["paragraph_semantic"]["drop_references"] is False
+
+
+def test_resolve_drop_references_env_true_without_request_param(monkeypatch):
+    # No request param → the env-true default survives into the snapshot.
+    monkeypatch.setenv("CHUNK_P_DROP_REFERENCES", "true")
+    cfg = TextChunkingConfig.model_validate({"strategy": "paragraph_semantic"})
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag())
+    assert chunk_options["paragraph_semantic"]["drop_references"] is True
+
+
+def test_resolve_rejects_size_below_inherited_overlap(monkeypatch):
+    # Overlap is inherited from addon_params (not in the request), so the
+    # request model can't catch it — _resolve_text_chunking must.
+    monkeypatch.setenv("CHUNK_F_OVERLAP_SIZE", "100")
+    addon = {"chunker": default_chunker_config()}
+    assert addon["chunker"]["fixed_token"]["chunk_overlap_token_size"] == 100
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "fixed_token", "params": {"chunk_token_size": 50}}
+    )
+    with pytest.raises(ValueError, match="chunk_overlap_token_size"):
+        _resolve_text_chunking(cfg, _stub_rag(addon))
+
+
+def test_resolve_allows_size_above_inherited_overlap(monkeypatch):
+    monkeypatch.setenv("CHUNK_F_OVERLAP_SIZE", "100")
+    addon = {"chunker": default_chunker_config()}
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "fixed_token", "params": {"chunk_token_size": 400}}
+    )
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
+    assert chunk_options["fixed_token"]["chunk_token_size"] == 400
+
+
+def test_resolve_skips_overlap_check_for_delimiter_only(monkeypatch):
+    # Delimiter-only fixed-token chunking never uses overlap, so a small
+    # chunk_token_size below the inherited overlap must NOT be rejected.
+    monkeypatch.setenv("CHUNK_F_OVERLAP_SIZE", "100")
+    addon = {"chunker": default_chunker_config()}
+    cfg = TextChunkingConfig.model_validate(
+        {
+            "strategy": "fixed_token",
+            "params": {
+                "split_by_character": "\n\n",
+                "split_by_character_only": True,
+                "chunk_token_size": 50,
+            },
+        }
+    )
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
+    assert chunk_options["fixed_token"]["chunk_token_size"] == 50
+
+
+def test_resolve_enforces_overlap_when_only_flag_without_delimiter(monkeypatch):
+    # split_by_character_only is a no-op without split_by_character: the chunker
+    # falls back to normal token windowing, which DOES use overlap — so the
+    # overlap < size check must still fire here.
+    monkeypatch.setenv("CHUNK_F_OVERLAP_SIZE", "100")
+    monkeypatch.delenv("CHUNK_F_SPLIT_BY_CHARACTER", raising=False)
+    addon = {"chunker": default_chunker_config()}
+    cfg = TextChunkingConfig.model_validate(
+        {
+            "strategy": "fixed_token",
+            "params": {"split_by_character_only": True, "chunk_token_size": 50},
+        }
+    )
+    with pytest.raises(ValueError, match="chunk_overlap_token_size"):
+        _resolve_text_chunking(cfg, _stub_rag(addon))
+
+
+def test_resolve_allows_amount_over_100_with_inherited_std_type():
+    # Request overrides only the amount; the standard_deviation type is
+    # inherited from addon_params. std/iqr have no (0, 100] ceiling, so this
+    # must NOT be rejected (the request model deferred the check here).
+    addon = {
+        "chunker": {
+            "semantic_vector": {"breakpoint_threshold_type": "standard_deviation"}
+        }
+    }
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 150}}
+    )
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
+    assert chunk_options["semantic_vector"]["breakpoint_threshold_amount"] == 150
+    assert (
+        chunk_options["semantic_vector"]["breakpoint_threshold_type"]
+        == "standard_deviation"
+    )
+
+
+def test_resolve_rejects_amount_over_100_with_inherited_percentile_type():
+    # Same partial override, but the effective (inherited) type is percentile,
+    # which feeds np.percentile and requires the (0, 100] ceiling.
+    addon = {
+        "chunker": {"semantic_vector": {"breakpoint_threshold_type": "percentile"}}
+    }
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "semantic_vector", "params": {"breakpoint_threshold_amount": 150}}
+    )
+    with pytest.raises(ValueError, match="breakpoint_threshold_amount"):
+        _resolve_text_chunking(cfg, _stub_rag(addon))
+
+
+def test_resolve_null_size_does_not_erase_inherited_default(monkeypatch):
+    # An explicit null in the request must not overwrite the resolved size
+    # with None (which would make the chunker do int(None) in the background).
+    monkeypatch.setenv("CHUNK_F_SIZE", "640")
+    addon = {"chunker": default_chunker_config()}
+    cfg = TextChunkingConfig.model_validate(
+        {"strategy": "fixed_token", "params": {"chunk_token_size": None}}
+    )
+    _, chunk_options = _resolve_text_chunking(cfg, _stub_rag(addon))
+    # null dropped by the model -> inherited CHUNK_F_SIZE survives, no None.
+    assert chunk_options["fixed_token"]["chunk_token_size"] == 640
+
+
+# ---------------------------------------------------------------------------
+# 3. Route forwarding + synchronous 422
+# ---------------------------------------------------------------------------
+
+
+class _FwdDocStatus:
+    async def get_doc_by_file_basename(self, basename):
+        return None
+
+
+class _FwdRag:
+    workspace = "chunk-fwd-test"
+    addon_params: dict = {}
+
+    def __init__(self):
+        self.doc_status = _FwdDocStatus()
+
+
+_HEADERS = {"X-API-Key": "test-key"}
+
+
+def _make_client(monkeypatch, addon_params=None):
+    """Build a TestClient whose enqueue-slot guards are no-ops and whose
+    ``pipeline_index_texts`` is a spy recording the forwarded args.
+
+    ``addon_params`` seeds the rag the routes resolve chunking against; the
+    handler calls the real ``_resolve_text_chunking`` synchronously, so the
+    effective-overlap validation runs against this snapshot.
+    """
+    captured: dict = {}
+
+    async def _spy(rag, texts, file_sources=None, track_id=None, chunking=None):
+        captured["texts"] = texts
+        captured["file_sources"] = file_sources
+        captured["chunking"] = chunking
+
+    async def _noop_reserve(rag):
+        return False
+
+    async def _noop_release(rag):
+        return None
+
+    monkeypatch.setattr(_dr, "pipeline_index_texts", _spy)
+    monkeypatch.setattr(_dr, "_reserve_enqueue_slot", _noop_reserve)
+    monkeypatch.setattr(_dr, "_release_enqueue_slot", _noop_release)
+
+    rag = _FwdRag()
+    rag.addon_params = addon_params if addon_params is not None else {}
+
+    app = FastAPI()
+    app.include_router(
+        create_document_routes(rag, SimpleNamespace(), api_key="test-key")
+    )
+    return TestClient(app), captured
+
+
+def test_insert_text_forwards_chunking(monkeypatch):
+    client, captured = _make_client(monkeypatch)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={
+            "text": "hello world",
+            "file_source": "a.md",
+            "chunking": {
+                "strategy": "recursive_character",
+                "params": {"chunk_token_size": 1000, "separators": ["X"]},
+            },
+        },
+    )
+    assert resp.status_code == 200
+    assert captured["chunking"] is not None
+    assert captured["chunking"].strategy == "recursive_character"
+    assert captured["chunking"].params == {
+        "chunk_token_size": 1000,
+        "separators": ["X"],
+    }
+
+
+def test_insert_texts_forwards_chunking(monkeypatch):
+    client, captured = _make_client(monkeypatch)
+    resp = client.post(
+        "/documents/texts",
+        headers=_HEADERS,
+        json={
+            "texts": ["one", "two"],
+            "file_sources": ["a.md", "b.md"],
+            "chunking": {"strategy": "semantic_vector", "params": {"buffer_size": 2}},
+        },
+    )
+    assert resp.status_code == 200
+    assert captured["chunking"].strategy == "semantic_vector"
+    assert captured["chunking"].params == {"buffer_size": 2}
+
+
+def test_insert_text_without_chunking_forwards_none(monkeypatch):
+    client, captured = _make_client(monkeypatch)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={"text": "hello", "file_source": "a.md"},
+    )
+    assert resp.status_code == 200
+    assert captured["chunking"] is None
+
+
+def test_insert_text_returns_422_on_malformed_chunking_without_scheduling(monkeypatch):
+    client, captured = _make_client(monkeypatch)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={
+            "text": "hello",
+            "file_source": "a.md",
+            "chunking": {
+                "strategy": "recursive_character",
+                "params": {"separators": "notalist"},
+            },
+        },
+    )
+    assert resp.status_code == 422
+    # Body validation fails before the endpoint body runs: no background
+    # indexing is scheduled, so the spy never fires.
+    assert captured == {}
+
+
+def test_insert_text_returns_422_when_size_below_inherited_overlap(monkeypatch):
+    # chunk_token_size=50 in the request, overlap=100 inherited from the
+    # rag's addon_params (not in the request). The model can't catch this;
+    # the handler's synchronous _resolve_text_chunking must, BEFORE any
+    # background work is scheduled.
+    addon = {
+        "chunker": {
+            "chunk_token_size": 1200,
+            "fixed_token": {"chunk_overlap_token_size": 100},
+        }
+    }
+    client, captured = _make_client(monkeypatch, addon_params=addon)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={
+            "text": "hello",
+            "file_source": "a.md",
+            "chunking": {"strategy": "fixed_token", "params": {"chunk_token_size": 50}},
+        },
+    )
+    assert resp.status_code == 422
+    assert "chunk_overlap_token_size" in resp.json()["detail"]
+    # Rejected synchronously: background indexing never scheduled.
+    assert captured == {}
+
+
+def test_insert_text_allows_amount_override_inheriting_std_type(monkeypatch):
+    # Reviewer scenario: deployment sets standard_deviation; a request
+    # overrides only breakpoint_threshold_amount (> 100). This must be
+    # accepted (not 422), since std has no (0, 100] ceiling.
+    addon = {
+        "chunker": {
+            "semantic_vector": {"breakpoint_threshold_type": "standard_deviation"}
+        }
+    }
+    client, captured = _make_client(monkeypatch, addon_params=addon)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={
+            "text": "hello",
+            "file_source": "a.md",
+            "chunking": {
+                "strategy": "semantic_vector",
+                "params": {"breakpoint_threshold_amount": 150},
+            },
+        },
+    )
+    assert resp.status_code == 200
+    assert captured["chunking"].params == {"breakpoint_threshold_amount": 150.0}
+
+
+def test_insert_text_rejects_amount_over_100_inheriting_percentile_type(monkeypatch):
+    addon = {
+        "chunker": {"semantic_vector": {"breakpoint_threshold_type": "percentile"}}
+    }
+    client, captured = _make_client(monkeypatch, addon_params=addon)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={
+            "text": "hello",
+            "file_source": "a.md",
+            "chunking": {
+                "strategy": "semantic_vector",
+                "params": {"breakpoint_threshold_amount": 150},
+            },
+        },
+    )
+    assert resp.status_code == 422
+    assert "breakpoint_threshold_amount" in resp.json()["detail"]
+    assert captured == {}
+
+
+def test_insert_text_rejects_malformed_sentence_split_regex(monkeypatch):
+    # Malformed regex must 422 at request parse time, before scheduling.
+    client, captured = _make_client(monkeypatch)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={
+            "text": "hello",
+            "file_source": "a.md",
+            "chunking": {
+                "strategy": "semantic_vector",
+                "params": {"sentence_split_regex": "("},
+            },
+        },
+    )
+    assert resp.status_code == 422
+    assert captured == {}
+
+
+def test_insert_text_drops_explicit_null_param(monkeypatch):
+    # "chunk_token_size": null must be treated as "inherit" (dropped), so the
+    # request succeeds and the forwarded params carry no None that would later
+    # crash the chunker with int(None).
+    client, captured = _make_client(monkeypatch)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={
+            "text": "hello",
+            "file_source": "a.md",
+            "chunking": {
+                "strategy": "fixed_token",
+                "params": {"chunk_token_size": None, "chunk_overlap_token_size": 50},
+            },
+        },
+    )
+    assert resp.status_code == 200
+    assert captured["chunking"].params == {"chunk_overlap_token_size": 50}
+
+
+def test_insert_text_allows_small_size_for_delimiter_only(monkeypatch):
+    # Paragraph splitting with a small chunk_token_size: overlap is inherited
+    # (100) but unused in delimiter-only mode, so this must succeed, not 422.
+    addon = {"chunker": {"fixed_token": {"chunk_overlap_token_size": 100}}}
+    client, captured = _make_client(monkeypatch, addon_params=addon)
+    resp = client.post(
+        "/documents/text",
+        headers=_HEADERS,
+        json={
+            "text": "hello",
+            "file_source": "a.md",
+            "chunking": {
+                "strategy": "fixed_token",
+                "params": {
+                    "split_by_character": "\n\n",
+                    "split_by_character_only": True,
+                    "chunk_token_size": 50,
+                },
+            },
+        },
+    )
+    assert resp.status_code == 200
+    assert captured["chunking"].params["chunk_token_size"] == 50

tests/api/routes/test_document_routes_paginated.py

@@ -0,0 +1,119 @@
+import importlib
+import sys
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+_original_argv = sys.argv[:]
+sys.argv = [sys.argv[0]]
+_document_routes = importlib.import_module("lightrag.api.routers.document_routes")
+_base = importlib.import_module("lightrag.base")
+sys.argv = _original_argv
+
+create_document_routes = _document_routes.create_document_routes
+DocProcessingStatus = _base.DocProcessingStatus
+DocStatus = _base.DocStatus
+DocStatusStorage = _base.DocStatusStorage
+
+pytestmark = pytest.mark.offline
+
+
+def _doc(status: DocStatus, suffix: str) -> DocProcessingStatus:
+    return DocProcessingStatus(
+        content_summary=f"{status.value} summary",
+        content_length=10,
+        file_path=f"{suffix}.pdf",
+        status=status,
+        created_at="2024-01-01T00:00:00+00:00",
+        updated_at="2024-01-01T00:00:00+00:00",
+        metadata={},
+    )
+
+
+class _FakeDocStatusStorage:
+    def __init__(self):
+        self.docs = {
+            "processed-doc": _doc(DocStatus.PROCESSED, "processed"),
+            "parsing-doc": _doc(DocStatus.PARSING, "parsing"),
+            "analyzing-doc": _doc(DocStatus.ANALYZING, "analyzing"),
+        }
+
+    async def get_docs_paginated(
+        self,
+        status_filter=None,
+        status_filters=None,
+        page=1,
+        page_size=50,
+        sort_field="updated_at",
+        sort_direction="desc",
+    ):
+        selected_statuses = DocStatusStorage.resolve_status_filter_values(
+            status_filter=status_filter,
+            status_filters=status_filters,
+        )
+        documents = [
+            (doc_id, doc)
+            for doc_id, doc in self.docs.items()
+            if selected_statuses is None or doc.status.value in selected_statuses
+        ]
+        return documents[:page_size], len(documents)
+
+    async def get_all_status_counts(self):
+        return {"processed": 1, "parsing": 1, "analyzing": 1}
+
+
+_fake_doc_status = _FakeDocStatusStorage()
+_app = FastAPI()
+_app.include_router(
+    create_document_routes(
+        SimpleNamespace(doc_status=_fake_doc_status),
+        SimpleNamespace(),
+        api_key="test-key",
+    )
+)
+_client = TestClient(_app)
+_headers = {"X-API-Key": "test-key"}
+
+
+def test_documents_paginated_accepts_status_filter():
+    response = _client.post(
+        "/documents/paginated",
+        headers=_headers,
+        json={
+            "status_filter": "processed",
+            "page": 1,
+            "page_size": 10,
+            "sort_field": "updated_at",
+            "sort_direction": "desc",
+        },
+    )
+
+    assert response.status_code == 200
+    payload = response.json()
+    assert payload["pagination"]["total_count"] == 1
+    assert [doc["id"] for doc in payload["documents"]] == ["processed-doc"]
+
+
+def test_documents_paginated_status_filters_override_status_filter():
+    response = _client.post(
+        "/documents/paginated",
+        headers=_headers,
+        json={
+            "status_filter": "processed",
+            "status_filters": ["parsing", "analyzing"],
+            "page": 1,
+            "page_size": 10,
+            "sort_field": "updated_at",
+            "sort_direction": "desc",
+        },
+    )
+
+    assert response.status_code == 200
+    payload = response.json()
+    assert payload["pagination"]["total_count"] == 2
+    assert [doc["id"] for doc in payload["documents"]] == [
+        "parsing-doc",
+        "analyzing-doc",
+    ]

tests/api/routes/test_graph_routes_pipeline_busy.py

@@ -0,0 +1,298 @@
+"""Pipeline-busy guard tests for graph mutation endpoints.
+
+These tests verify that all 7 graph-mutation endpoints refuse to operate
+with HTTP 409 while the document pipeline is busy:
+
+- POST /graph/entity/edit          (graph_routes)
+- POST /graph/relation/edit        (graph_routes)
+- POST /graph/entity/create        (graph_routes)
+- POST /graph/relation/create      (graph_routes)
+- POST /graph/entities/merge       (graph_routes)
+- DELETE /graph/entity/delete      (graph_routes)
+- DELETE /graph/relation/delete    (graph_routes)
+
+The guard logic itself lives in
+``lightrag.api.routers.document_routes.check_pipeline_busy_or_raise`` and is
+exercised both at the endpoint integration layer (via monkeypatch, no
+shared-storage dependency) and at the unit layer (against a real
+``pipeline_status`` namespace).
+"""
+
+import importlib
+import sys
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+from fastapi import FastAPI, HTTPException
+from fastapi.testclient import TestClient
+
+# Importing routers loads ``lightrag.api.config`` which parses ``sys.argv`` via
+# argparse. Stash argv so pytest's CLI flags don't trip the parser.
+_original_argv = sys.argv[:]
+sys.argv = [sys.argv[0]]
+_graph_routes = importlib.import_module("lightrag.api.routers.graph_routes")
+_document_routes = importlib.import_module("lightrag.api.routers.document_routes")
+sys.argv = _original_argv
+
+create_graph_routes = _graph_routes.create_graph_routes
+create_document_routes = _document_routes.create_document_routes
+check_pipeline_busy_or_raise = _document_routes.check_pipeline_busy_or_raise
+
+pytestmark = pytest.mark.offline
+
+_API_KEY = "test-key"
+_HEADERS = {"X-API-Key": _API_KEY}
+
+
+# ---------------------------------------------------------------------------
+# Test scaffolding
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_rag() -> SimpleNamespace:
+    """Build a minimal LightRAG stand-in with the 7 mutation methods stubbed.
+
+    Each ``AsyncMock`` returns a payload shaped enough to satisfy the
+    endpoint's response model so the idle pass-through test can verify the
+    full request path. Busy tests don't rely on these return values; the
+    guard short-circuits before they're reached.
+    """
+    return SimpleNamespace(
+        workspace="",
+        aedit_entity=AsyncMock(
+            return_value={
+                "entity_name": "Alice",
+                "description": "updated",
+                "operation_summary": {
+                    "merged": False,
+                    "merge_status": "not_attempted",
+                    "merge_error": None,
+                    "operation_status": "success",
+                    "target_entity": None,
+                    "final_entity": "Alice",
+                    "renamed": False,
+                },
+            }
+        ),
+        aedit_relation=AsyncMock(return_value={"description": "updated"}),
+        acreate_entity=AsyncMock(return_value={"entity_name": "Alice"}),
+        acreate_relation=AsyncMock(return_value={"src_id": "a", "tgt_id": "b"}),
+        amerge_entities=AsyncMock(return_value={"merged_entity": "Alice"}),
+        adelete_by_entity=AsyncMock(
+            return_value=SimpleNamespace(
+                status="success", message="deleted", doc_id="ignored"
+            )
+        ),
+        adelete_by_relation=AsyncMock(
+            return_value=SimpleNamespace(
+                status="success", message="deleted", doc_id="ignored"
+            )
+        ),
+    )
+
+
+def _build_client(rag: SimpleNamespace) -> TestClient:
+    app = FastAPI()
+    app.include_router(create_graph_routes(rag, api_key=_API_KEY))
+    app.include_router(create_document_routes(rag, SimpleNamespace(), api_key=_API_KEY))
+    return TestClient(app)
+
+
+async def _force_busy_guard(_rag) -> None:
+    """Stand-in for ``check_pipeline_busy_or_raise`` that always refuses."""
+    raise HTTPException(
+        status_code=409,
+        detail=(
+            "Pipeline is busy with another operation. "
+            "Wait for the running job to finish before editing "
+            "the knowledge graph."
+        ),
+    )
+
+
+async def _noop_guard(_rag) -> None:
+    """Stand-in for ``check_pipeline_busy_or_raise`` that always permits."""
+    return None
+
+
+def _patch_guard(monkeypatch, replacement) -> None:
+    """Replace the guard reference in BOTH consumer modules.
+
+    ``graph_routes`` re-binds the name via ``from .document_routes import ...``
+    so patching only ``document_routes`` would miss the graph endpoints.
+    """
+    monkeypatch.setattr(_graph_routes, "check_pipeline_busy_or_raise", replacement)
+    monkeypatch.setattr(_document_routes, "check_pipeline_busy_or_raise", replacement)
+
+
+# ---------------------------------------------------------------------------
+# Part A: endpoint integration -- guard refuses with 409
+# ---------------------------------------------------------------------------
+
+_ENDPOINTS = [
+    pytest.param(
+        "POST",
+        "/graph/entity/edit",
+        {"entity_name": "Alice", "updated_data": {"description": "x"}},
+        id="update_entity",
+    ),
+    pytest.param(
+        "POST",
+        "/graph/relation/edit",
+        {
+            "source_id": "Alice",
+            "target_id": "Bob",
+            "updated_data": {"description": "x"},
+        },
+        id="update_relation",
+    ),
+    pytest.param(
+        "POST",
+        "/graph/entity/create",
+        {"entity_name": "Alice", "entity_data": {"description": "x"}},
+        id="create_entity",
+    ),
+    pytest.param(
+        "POST",
+        "/graph/relation/create",
+        {
+            "source_entity": "Alice",
+            "target_entity": "Bob",
+            "relation_data": {"description": "x"},
+        },
+        id="create_relation",
+    ),
+    pytest.param(
+        "POST",
+        "/graph/entities/merge",
+        {"entities_to_change": ["Alic"], "entity_to_change_into": "Alice"},
+        id="merge_entities",
+    ),
+    pytest.param(
+        "DELETE",
+        "/graph/entity/delete",
+        {"entity_name": "Alice"},
+        id="delete_entity",
+    ),
+    pytest.param(
+        "DELETE",
+        "/graph/relation/delete",
+        {"source_entity": "Alice", "target_entity": "Bob"},
+        id="delete_relation",
+    ),
+]
+
+
+@pytest.mark.parametrize("method, path, body", _ENDPOINTS)
+def test_endpoint_refuses_with_409_when_pipeline_busy(method, path, body, monkeypatch):
+    rag = _make_mock_rag()
+    client = _build_client(rag)
+    _patch_guard(monkeypatch, _force_busy_guard)
+
+    response = client.request(method, path, json=body, headers=_HEADERS)
+
+    assert response.status_code == 409, response.text
+    payload = response.json()
+    assert "Pipeline is busy" in payload["detail"]
+    # Guard must short-circuit before the underlying mutation runs.
+    for attr in (
+        "aedit_entity",
+        "aedit_relation",
+        "acreate_entity",
+        "acreate_relation",
+        "amerge_entities",
+        "adelete_by_entity",
+        "adelete_by_relation",
+    ):
+        getattr(rag, attr).assert_not_awaited()
+
+
+def test_endpoint_passes_through_when_pipeline_idle(monkeypatch):
+    """Sanity check: with an idle guard, the request reaches ``rag.aedit_entity``."""
+    rag = _make_mock_rag()
+    client = _build_client(rag)
+    _patch_guard(monkeypatch, _noop_guard)
+
+    response = client.post(
+        "/graph/entity/edit",
+        json={"entity_name": "Alice", "updated_data": {"description": "x"}},
+        headers=_HEADERS,
+    )
+
+    assert response.status_code == 200, response.text
+    rag.aedit_entity.assert_awaited_once()
+
+
+# ---------------------------------------------------------------------------
+# Part B: helper unit -- against real pipeline_status namespace
+# ---------------------------------------------------------------------------
+
+
+async def _with_pipeline_status(action):
+    """Bootstrap pipeline_status, run ``action(pipeline_status)``, then tear down.
+
+    ``initialize_share_data`` is idempotent within a process but
+    ``finalize_share_data`` is required to release the Manager/lock state so
+    repeated calls in subsequent tests start clean.
+    """
+    from lightrag.kg.shared_storage import (
+        finalize_share_data,
+        get_namespace_data,
+        initialize_pipeline_status,
+        initialize_share_data,
+    )
+
+    initialize_share_data()
+    try:
+        await initialize_pipeline_status(workspace="")
+        pipeline_status = await get_namespace_data("pipeline_status", workspace="")
+        await action(pipeline_status)
+    finally:
+        finalize_share_data()
+
+
+async def test_helper_raises_409_when_busy_flag_set():
+    async def _do(pipeline_status):
+        pipeline_status["busy"] = True
+        rag = SimpleNamespace(workspace="")
+        with pytest.raises(HTTPException) as exc_info:
+            await check_pipeline_busy_or_raise(rag)
+        assert exc_info.value.status_code == 409
+        assert "Pipeline is busy" in exc_info.value.detail
+
+    await _with_pipeline_status(_do)
+
+
+async def test_helper_returns_silently_when_pipeline_idle():
+    async def _do(pipeline_status):
+        pipeline_status["busy"] = False
+        rag = SimpleNamespace(workspace="")
+        # Should not raise.
+        await check_pipeline_busy_or_raise(rag)
+
+    await _with_pipeline_status(_do)
+
+
+async def test_helper_is_noop_when_pipeline_status_uninitialized():
+    """When pipeline_status namespace was never bootstrapped the helper must pass.
+
+    ``get_namespace_data`` raises ``PipelineNotInitializedError`` when the
+    pipeline_status namespace is missing (share data initialized but the
+    pipeline namespace never created); the helper swallows that error so test
+    rigs without an end-to-end RAG bootstrap stay green. Mirrors the existing
+    contract of ``_acquire_destructive_busy``.
+    """
+    from lightrag.kg.shared_storage import (
+        finalize_share_data,
+        initialize_share_data,
+    )
+
+    initialize_share_data()
+    try:
+        rag = SimpleNamespace(workspace="__never_bootstrapped__")
+        # Intentionally skip ``initialize_pipeline_status``: helper should
+        # catch ``PipelineNotInitializedError`` and return silently.
+        await check_pipeline_busy_or_raise(rag)
+    finally:
+        finalize_share_data()

tests/api/routes/test_query_stream_routes.py

@@ -0,0 +1,174 @@
+"""
+Verify the /query and /query/stream endpoint response types.
+
+Ensures:
+  - /query  → application/json  (no streaming, backward-compatible)
+  - /query/stream → application/x-ndjson
+"""
+
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+from fastapi.testclient import TestClient
+
+
+_ENV_VARS_TO_ISOLATE = (
+    "LLM_BINDING",
+    "EMBEDDING_BINDING",
+    "LLM_BINDING_HOST",
+    "LLM_BINDING_API_KEY",
+    "LLM_MODEL",
+    "EMBEDDING_BINDING_HOST",
+    "EMBEDDING_BINDING_API_KEY",
+    "EMBEDDING_MODEL",
+    "LIGHTRAG_API_PREFIX",
+    "LIGHTRAG_KV_STORAGE",
+    "LIGHTRAG_VECTOR_STORAGE",
+    "LIGHTRAG_GRAPH_STORAGE",
+    "LIGHTRAG_DOC_STATUS_STORAGE",
+)
+
+
+@pytest.fixture(autouse=True)
+def _isolate_env(monkeypatch):
+    for var in _ENV_VARS_TO_ISOLATE:
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("LLM_BINDING", "ollama")
+    monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
+
+
+def _build_client():
+    original_argv = sys.argv.copy()
+    try:
+        sys.argv = ["lightrag-server"]
+        from lightrag.api.config import parse_args
+        from lightrag.api.lightrag_server import create_app
+
+        args = parse_args()
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            return TestClient(create_app(args))
+    finally:
+        sys.argv = original_argv
+
+
+class TestQueryRouteJsonOnly:
+    """The /query endpoint must stay JSON-only to preserve backward compatibility."""
+
+    def test_openapi_spec_declares_json_response(self):
+        client = _build_client()
+        response = client.get("/openapi.json")
+        assert response.status_code == 200
+        spec = response.json()
+
+        paths = spec.get("paths", {})
+        query_path = paths.get("/query", {})
+        assert query_path, "/query must be in OpenAPI paths"
+
+        post_op = query_path.get("post", {})
+        responses = post_op.get("responses", {})
+        ok_resp = responses.get("200", {})
+        content = ok_resp.get("content", {})
+
+        # The /query endpoint must declare application/json — NOT ndjson
+        assert "application/json" in content, (
+            "/query must declare application/json in OpenAPI spec"
+        )
+        assert "application/x-ndjson" not in content, (
+            "/query must NOT declare application/x-ndjson — streaming belongs to /query/stream"
+        )
+
+    def test_query_route_exists_and_accepts_post(self):
+        client = _build_client()
+        # A minimal POST to /query should reach the route (it'll 422 or 500
+        # since we don't have a real LLM, but it should NOT 404/405)
+        response = client.post("/query", json={"query": "test", "mode": "mix"})
+        assert response.status_code not in (
+            404,
+            405,
+        ), "/query route must exist and accept POST"
+
+
+class TestQueryStreamRoute:
+    """The /query/stream endpoint must serve application/x-ndjson."""
+
+    def test_openapi_spec_declares_ndjson_response(self):
+        client = _build_client()
+        response = client.get("/openapi.json")
+        assert response.status_code == 200
+        spec = response.json()
+
+        paths = spec.get("paths", {})
+        stream_path = paths.get("/query/stream", {})
+        assert stream_path, "/query/stream must be in OpenAPI paths"
+
+        post_op = stream_path.get("post", {})
+        responses = post_op.get("responses", {})
+        ok_resp = responses.get("200", {})
+        content = ok_resp.get("content", {})
+
+        # The /query/stream endpoint must declare application/x-ndjson
+        assert "application/x-ndjson" in content, (
+            "/query/stream must declare application/x-ndjson in OpenAPI spec"
+        )
+
+    def test_stream_route_exists_and_accepts_post(self):
+        client = _build_client()
+        response = client.post("/query/stream", json={"query": "test", "mode": "mix"})
+        assert response.status_code not in (
+            404,
+            405,
+        ), "/query/stream route must exist and accept POST"
+
+
+class TestQueryStreamResponseContentType:
+    """When the mock LLM returns a non-streaming result, /query/stream must
+    still set the correct Content-Type header."""
+
+    def test_stream_response_has_ndjson_content_type(self):
+        """Even without a real LLM, the streaming response must carry the
+        correct media type header."""
+
+        original_argv = sys.argv.copy()
+        try:
+            sys.argv = ["lightrag-server"]
+            from lightrag.api.config import parse_args
+            from lightrag.api.lightrag_server import create_app
+
+            args = parse_args()
+
+            mock_rag = MagicMock()
+            mock_result = {
+                "llm_response": {
+                    "is_streaming": False,
+                    "content": "test response",
+                },
+                "data": {"references": []},
+            }
+            # Return a coroutine
+            mock_rag.aquery_llm = MagicMock()
+
+            async def _fake_aquery(*a, **kw):
+                return mock_result
+
+            mock_rag.aquery_llm.side_effect = _fake_aquery
+
+            with patch("lightrag.api.lightrag_server.LightRAG", return_value=mock_rag):
+                app = create_app(args)
+
+            client = TestClient(app)
+            response = client.post(
+                "/query/stream",
+                json={
+                    "query": "test",
+                    "mode": "mix",
+                    "include_references": True,
+                },
+            )
+            content_type = response.headers.get("content-type", "")
+            assert "application/x-ndjson" in content_type, (
+                f"/query/stream must return application/x-ndjson, got: {content_type}"
+            )
+        finally:
+            sys.argv = original_argv

tests/api/routes/test_supported_extensions_markdown.py

@@ -0,0 +1,48 @@
+"""Upload allowlist exposes the native markdown suffixes (.md / .textpack).
+
+``DocumentManager.supported_extensions`` derives the uploadable suffix set
+live from the parser registry + routing: a suffix is advertised only when an
+unhinted ``x.<suffix>`` routes to an engine that actually supports it. These
+tests pin that ``.md`` and ``.textpack`` are uploadable (and routable), and
+that ``.textpack`` resolves to the native engine.
+"""
+
+import importlib
+import sys
+
+import pytest
+
+# document_routes parses argv at import time; guard it like the sibling tests.
+_original_argv = sys.argv[:]
+sys.argv = [sys.argv[0]]
+_document_routes = importlib.import_module("lightrag.api.routers.document_routes")
+_parser_routing = importlib.import_module("lightrag.parser.routing")
+sys.argv = _original_argv
+
+DocumentManager = _document_routes.DocumentManager
+resolve_file_parser_engine = _parser_routing.resolve_file_parser_engine
+
+
+@pytest.fixture
+def doc_manager(tmp_path, monkeypatch):
+    # Clear any ambient LIGHTRAG_PARSER so the default routing applies:
+    # .md -> legacy (opt-in for native), .textpack -> native.
+    monkeypatch.delenv("LIGHTRAG_PARSER", raising=False)
+    return DocumentManager(str(tmp_path))
+
+
+def test_textpack_and_md_are_uploadable(doc_manager):
+    extensions = doc_manager.supported_extensions
+    assert ".textpack" in extensions
+    assert ".md" in extensions
+
+
+def test_bare_textpack_is_supported_file_via_native(doc_manager, monkeypatch):
+    monkeypatch.delenv("LIGHTRAG_PARSER", raising=False)
+    assert doc_manager.is_supported_file("note.textpack") is True
+    assert resolve_file_parser_engine("note.textpack") == "native"
+
+
+def test_md_is_supported_file(doc_manager, monkeypatch):
+    monkeypatch.delenv("LIGHTRAG_PARSER", raising=False)
+    assert doc_manager.is_supported_file("doc.md") is True

tests/api/test_cors.py

@@ -0,0 +1,187 @@
+"""
+Tests for CORS middleware configuration on the LightRAG API server.
+
+LightRAG authenticates exclusively via request headers (Authorization Bearer
+token and X-API-Key) and never via cookies or other ambient credentials. The
+Fetch spec forbids pairing the wildcard origin "*" with credentialed responses
+("Access-Control-Allow-Origin: *" together with
+"Access-Control-Allow-Credentials: true"). These tests pin the resulting
+behavior:
+
+- When CORS_ORIGINS is the wildcard "*" (the default), credentials are disabled
+  and the server responds with a clean "Access-Control-Allow-Origin: *" and no
+  "Access-Control-Allow-Credentials" header.
+- When CORS_ORIGINS is an explicit allowlist, the matching origin is reflected
+  and credentials are enabled.
+"""
+
+import sys
+from unittest.mock import patch, MagicMock
+
+import pytest
+from fastapi.testclient import TestClient
+
+
+# Env vars that the project's `.env` may have populated (via load_dotenv at
+# import time of lightrag.api.config). Tests must be hermetic and not depend on
+# developer-local .env values, so we clear/override anything that affects
+# parse_args() / create_app().
+_ENV_VARS_TO_ISOLATE = (
+    "LLM_BINDING",
+    "EMBEDDING_BINDING",
+    "LLM_BINDING_HOST",
+    "LLM_BINDING_API_KEY",
+    "LLM_MODEL",
+    "EMBEDDING_BINDING_HOST",
+    "EMBEDDING_BINDING_API_KEY",
+    "EMBEDDING_MODEL",
+    "CORS_ORIGINS",
+    "LIGHTRAG_API_PREFIX",
+    "LIGHTRAG_KV_STORAGE",
+    "LIGHTRAG_VECTOR_STORAGE",
+    "LIGHTRAG_GRAPH_STORAGE",
+    "LIGHTRAG_DOC_STATUS_STORAGE",
+)
+
+
+@pytest.fixture(autouse=True)
+def _isolate_env(monkeypatch):
+    """Isolate tests from developer-local .env pollution and global config state.
+
+    Clear env vars that affect parse_args()/create_app(), then set the minimal
+    viable defaults (ollama bindings) so create_app's binding validation passes
+    without touching real services. The global config singleton in
+    lightrag.api.config is reset before and after each test so neither developer
+    .env values nor other tests' parsed args leak into our CORS assertions.
+    """
+    for var in _ENV_VARS_TO_ISOLATE:
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("LLM_BINDING", "ollama")
+    monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
+
+    import lightrag.api.config as config
+
+    config._global_args = None
+    config._initialized = False
+    yield
+    config._global_args = None
+    config._initialized = False
+
+
+def _build_client(cors_origins, monkeypatch):
+    """Create a TestClient for an app configured with the given CORS_ORIGINS."""
+    from lightrag.api.config import parse_args, initialize_config
+
+    monkeypatch.setenv("CORS_ORIGINS", cors_origins)
+
+    original_argv = sys.argv.copy()
+    try:
+        sys.argv = ["lightrag-server"]
+        args = parse_args()
+    finally:
+        sys.argv = original_argv
+
+    # create_app / get_cors_origins read the module-level global_args proxy
+    # (not the passed args), which otherwise lazily re-parses sys.argv — by then
+    # polluted with pytest's argv. Pin the parsed args explicitly.
+    initialize_config(args, force=True)
+
+    with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+        mock_rag.return_value = MagicMock()
+        from lightrag.api.lightrag_server import create_app
+
+        app = create_app(args)
+    return TestClient(app)
+
+
+def _preflight(client, origin):
+    """Issue a CORS preflight request and return the response."""
+    return client.options(
+        "/query",
+        headers={
+            "Origin": origin,
+            "Access-Control-Request-Method": "POST",
+            "Access-Control-Request-Headers": "authorization,content-type",
+        },
+    )
+
+
+class TestWildcardCorsDisablesCredentials:
+    """With CORS_ORIGINS="*" the config must stay spec-compliant."""
+
+    def test_preflight_wildcard_no_credentials(self, monkeypatch):
+        client = _build_client("*", monkeypatch)
+        resp = _preflight(client, "https://evil.example.com")
+
+        assert resp.headers["access-control-allow-origin"] == "*"
+        # Wildcard origin must NOT be paired with credentials.
+        assert "access-control-allow-credentials" not in resp.headers
+
+    def test_simple_request_wildcard_no_credentials(self, monkeypatch):
+        client = _build_client("*", monkeypatch)
+        # /health is whitelisted (no auth) so this returns a real response with
+        # CORS headers attached by the middleware.
+        resp = client.get("/health", headers={"Origin": "https://evil.example.com"})
+
+        assert resp.headers.get("access-control-allow-origin") == "*"
+        assert "access-control-allow-credentials" not in resp.headers
+
+    def test_wildcard_mixed_with_explicit_origin_no_credentials(self, monkeypatch):
+        # Starlette treats any list containing "*" as allow-all, so a mixed config
+        # is still allow-all and must NOT enable credentials.
+        client = _build_client("*,https://app.example.com", monkeypatch)
+        resp = _preflight(client, "https://evil.example.com")
+
+        assert "access-control-allow-credentials" not in resp.headers
+
+    def test_wildcard_trailing_comma_no_credentials(self, monkeypatch):
+        # A trailing comma yields ["*", ""]; the empty entry is dropped and the
+        # wildcard still disables credentials.
+        client = _build_client("*,", monkeypatch)
+        resp = _preflight(client, "https://evil.example.com")
+
+        assert "access-control-allow-credentials" not in resp.headers
+
+
+class TestEmptyCorsFailsClosed:
+    """An explicitly empty CORS_ORIGINS disables cross-origin access (fail closed)."""
+
+    def test_empty_string_allows_no_origin(self, monkeypatch):
+        # CORS_ORIGINS= is an explicit "disable cross-origin" config and must not
+        # silently widen to "*".
+        client = _build_client("", monkeypatch)
+        resp = _preflight(client, "https://evil.example.com")
+
+        assert "access-control-allow-origin" not in resp.headers
+        assert "access-control-allow-credentials" not in resp.headers
+
+    def test_only_commas_allows_no_origin(self, monkeypatch):
+        # A value with no real origins (e.g. ",,") also fails closed.
+        client = _build_client(",,", monkeypatch)
+        resp = _preflight(client, "https://evil.example.com")
+
+        assert "access-control-allow-origin" not in resp.headers
+        assert "access-control-allow-credentials" not in resp.headers
+
+
+class TestExplicitAllowlistEnablesCredentials:
+    """An explicit origin allowlist reflects the origin and allows credentials."""
+
+    def test_preflight_allowed_origin_reflected_with_credentials(self, monkeypatch):
+        client = _build_client(
+            "https://app.example.com,https://dash.example.com", monkeypatch
+        )
+        resp = _preflight(client, "https://app.example.com")
+
+        assert resp.headers["access-control-allow-origin"] == "https://app.example.com"
+        assert resp.headers["access-control-allow-credentials"] == "true"
+
+    def test_preflight_disallowed_origin_not_reflected(self, monkeypatch):
+        client = _build_client("https://app.example.com", monkeypatch)
+        resp = _preflight(client, "https://evil.example.com")
+
+        # A non-allowlisted origin is never echoed back.
+        assert (
+            resp.headers.get("access-control-allow-origin")
+            != "https://evil.example.com"
+        )

tests/api/test_health_auth.py

@@ -0,0 +1,272 @@
+"""Tests for credential-gated configuration disclosure on ``GET /health``.
+
+Issue #3294: ``/health`` is whitelisted by default so it answers unauthenticated
+liveness probes (HTTP 200). It must therefore reveal sensitive runtime
+configuration (filesystem paths, LLM/embedding provider + model + host, storage
+backends, queue status, ...) ONLY to authenticated callers. These tests pin
+that behavior across the three authentication modes:
+
+- fully open (no AUTH_ACCOUNTS, no API key): everything is open anyway, so the
+  full configuration is returned to every caller.
+- password auth (AUTH_ACCOUNTS): only a valid non-guest JWT (or a configured
+  API key) unlocks the configuration; anonymous probes get liveness only.
+- API-key-only: only a valid X-API-Key unlocks the configuration.
+
+In every case ``/health`` returns HTTP 200 so external liveness probes keep
+working.
+"""
+
+import sys
+from unittest.mock import AsyncMock
+
+import pytest
+from fastapi import APIRouter
+from fastapi.testclient import TestClient
+
+# Fields that must NEVER appear in an unauthenticated response.
+_SENSITIVE_TOP_LEVEL = ("working_directory", "input_directory", "configuration")
+# Liveness fields that must always be present (safe; already exposed by the
+# unauthenticated /auth-status endpoint or pure liveness signals).
+_LIVENESS_FIELDS = ("status", "auth_mode", "core_version", "api_version")
+
+
+_ENV_VARS_TO_ISOLATE = (
+    "LLM_BINDING",
+    "EMBEDDING_BINDING",
+    "AUTH_ACCOUNTS",
+    "TOKEN_SECRET",
+    "LIGHTRAG_API_KEY",
+    "WHITELIST_PATHS",
+    "LIGHTRAG_API_PREFIX",
+)
+
+
+@pytest.fixture(autouse=True)
+def _isolate_env(monkeypatch):
+    """Keep tests hermetic from developer-local .env and global config state."""
+    for var in _ENV_VARS_TO_ISOLATE:
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("AUTH_ACCOUNTS", "")
+    monkeypatch.setenv("LIGHTRAG_API_KEY", "")
+    monkeypatch.setenv("TOKEN_SECRET", "")
+    monkeypatch.setenv("WHITELIST_PATHS", "/health,/api/*")
+    monkeypatch.setenv("LLM_BINDING", "ollama")
+    monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
+
+    import lightrag.api.config as config
+
+    config._global_args = None
+    config._initialized = False
+    yield
+    config._global_args = None
+    config._initialized = False
+
+
+class _FakeLightRAG:
+    """Minimal stand-in implementing the async surface /health touches."""
+
+    def __init__(self, *_args, **_kwargs):
+        pass
+
+    def register_role_llm_builder(self, _builder):
+        return None
+
+    def set_role_llm_metadata(self, _role, **_metadata):
+        return None
+
+    def get_llm_role_config(self):
+        return {}
+
+    async def get_llm_queue_status(self, include_base=True):
+        return {}
+
+    async def get_embedding_queue_status(self):
+        return {}
+
+    async def get_rerank_queue_status(self):
+        return {}
+
+
+class _FakeOllamaAPI:
+    def __init__(self, *_args, **_kwargs):
+        self.router = APIRouter()
+
+
+def _build_client(monkeypatch, *, api_key=None):
+    """Build a /health-capable TestClient with all backend I/O mocked out."""
+    from lightrag.api.config import parse_args, initialize_config
+
+    original_argv = sys.argv.copy()
+    try:
+        sys.argv = ["lightrag-server"]
+        args = parse_args()
+    finally:
+        sys.argv = original_argv
+    if api_key is not None:
+        args.key = api_key
+    initialize_config(args, force=True)
+
+    import lightrag.api.lightrag_server as lightrag_server
+
+    monkeypatch.setattr(lightrag_server, "LightRAG", _FakeLightRAG)
+    monkeypatch.setattr(lightrag_server, "check_frontend_build", lambda: (True, False))
+    monkeypatch.setattr(
+        lightrag_server, "create_document_routes", lambda *_a, **_k: APIRouter()
+    )
+    monkeypatch.setattr(
+        lightrag_server, "create_query_routes", lambda *_a, **_k: APIRouter()
+    )
+    monkeypatch.setattr(
+        lightrag_server, "create_graph_routes", lambda *_a, **_k: APIRouter()
+    )
+    monkeypatch.setattr(lightrag_server, "OllamaAPI", _FakeOllamaAPI)
+    monkeypatch.setattr(
+        lightrag_server, "get_namespace_data", AsyncMock(return_value={"busy": False})
+    )
+    monkeypatch.setattr(lightrag_server, "get_default_workspace", lambda: "default")
+    monkeypatch.setattr(
+        lightrag_server,
+        "cleanup_keyed_lock",
+        lambda: {"cleanup_performed": {}, "current_status": {}},
+    )
+
+    app = lightrag_server.create_app(args)
+    return TestClient(app)
+
+
+def _set_auth_mode(monkeypatch, *, auth_configured):
+    """Override the module-level auth flags the /health gate reads at runtime.
+
+    Also pin a whitelist that exempts /health so combined_auth keeps returning
+    200 for anonymous callers (the gate, not combined_auth, hides the config).
+    """
+    import lightrag.api.utils_api as utils_api
+
+    monkeypatch.setattr(utils_api, "auth_configured", auth_configured)
+    monkeypatch.setattr(
+        utils_api, "whitelist_patterns", [("/health", False), ("/api", True)]
+    )
+
+
+def _assert_liveness_only(body):
+    for field in _LIVENESS_FIELDS:
+        assert field in body, f"liveness field {field!r} missing"
+    for field in _SENSITIVE_TOP_LEVEL:
+        assert field not in body, f"sensitive field {field!r} leaked"
+
+
+def _assert_full_config(body):
+    assert "configuration" in body
+    assert "working_directory" in body
+    assert "llm_binding" in body["configuration"]
+
+
+# --------------------------------------------------------------------------- #
+# Fully open mode: everything is open, so config is returned to everyone.
+# --------------------------------------------------------------------------- #
+def test_open_mode_returns_full_config_to_anonymous(monkeypatch):
+    client = _build_client(monkeypatch)
+    _set_auth_mode(monkeypatch, auth_configured=False)
+
+    resp = client.get("/health")
+
+    assert resp.status_code == 200
+    _assert_full_config(resp.json())
+
+
+# --------------------------------------------------------------------------- #
+# Password auth: anonymous gets liveness only; a valid token unlocks config.
+# --------------------------------------------------------------------------- #
+def test_password_mode_anonymous_gets_liveness_only(monkeypatch):
+    client = _build_client(monkeypatch)
+    _set_auth_mode(monkeypatch, auth_configured=True)
+
+    resp = client.get("/health")
+
+    assert resp.status_code == 200  # liveness probe must stay green
+    _assert_liveness_only(resp.json())
+
+
+def test_password_mode_valid_token_unlocks_config(monkeypatch):
+    import lightrag.api.utils_api as utils_api
+
+    client = _build_client(monkeypatch)
+    _set_auth_mode(monkeypatch, auth_configured=True)
+    monkeypatch.setattr(
+        utils_api.auth_handler,
+        "validate_token",
+        lambda token: (
+            {"username": "admin", "role": "user"}
+            if token == "valid-user-token"
+            else (_ for _ in ()).throw(ValueError("bad token"))
+        ),
+    )
+
+    resp = client.get("/health", headers={"Authorization": "Bearer valid-user-token"})
+
+    assert resp.status_code == 200
+    _assert_full_config(resp.json())
+
+
+def test_password_mode_guest_token_stays_liveness_only(monkeypatch):
+    import lightrag.api.utils_api as utils_api
+
+    client = _build_client(monkeypatch)
+    _set_auth_mode(monkeypatch, auth_configured=True)
+    monkeypatch.setattr(
+        utils_api.auth_handler,
+        "validate_token",
+        lambda token: {"username": "guest", "role": "guest"},
+    )
+
+    resp = client.get("/health", headers={"Authorization": "Bearer guest-token"})
+
+    assert resp.status_code == 200
+    _assert_liveness_only(resp.json())
+
+
+# --------------------------------------------------------------------------- #
+# API-key-only mode: only a valid X-API-Key unlocks the configuration.
+# --------------------------------------------------------------------------- #
+def test_api_key_mode_anonymous_gets_liveness_only(monkeypatch):
+    client = _build_client(monkeypatch, api_key="secret-key")
+    _set_auth_mode(monkeypatch, auth_configured=False)
+
+    resp = client.get("/health")
+
+    assert resp.status_code == 200
+    _assert_liveness_only(resp.json())
+
+
+def test_api_key_mode_valid_key_unlocks_config(monkeypatch):
+    client = _build_client(monkeypatch, api_key="secret-key")
+    _set_auth_mode(monkeypatch, auth_configured=False)
+
+    resp = client.get("/health", headers={"X-API-Key": "secret-key"})
+
+    assert resp.status_code == 200
+    _assert_full_config(resp.json())
+
+
+# --------------------------------------------------------------------------- #
+# Combined mode (AUTH_ACCOUNTS + LIGHTRAG_API_KEY): either a valid JWT or a
+# valid X-API-Key unlocks the configuration; an anonymous probe gets liveness.
+# --------------------------------------------------------------------------- #
+def test_combined_mode_anonymous_gets_liveness_only(monkeypatch):
+    client = _build_client(monkeypatch, api_key="secret-key")
+    _set_auth_mode(monkeypatch, auth_configured=True)
+
+    resp = client.get("/health")
+
+    assert resp.status_code == 200
+    _assert_liveness_only(resp.json())
+
+
+def test_combined_mode_valid_api_key_unlocks_config(monkeypatch):
+    client = _build_client(monkeypatch, api_key="secret-key")
+    _set_auth_mode(monkeypatch, auth_configured=True)
+
+    resp = client.get("/health", headers={"X-API-Key": "secret-key"})
+
+    assert resp.status_code == 200
+    _assert_full_config(resp.json())

tests/api/test_lightrag_ollama_chat.py

@@ -0,0 +1,875 @@
+"""
+LightRAG Ollama Compatibility Interface Test Script
+
+This script tests the LightRAG's Ollama compatibility interface, including:
+1. Basic functionality tests (streaming and non-streaming responses)
+2. Query mode tests (local, global, naive, hybrid)
+3. Error handling tests (including streaming and non-streaming scenarios)
+
+All responses use the JSON Lines format, complying with the Ollama API specification.
+"""
+
+import pytest
+import requests
+import json
+import argparse
+import time
+from typing import Dict, Any, Optional, List, Callable
+from dataclasses import dataclass, asdict
+from datetime import datetime
+from pathlib import Path
+from enum import Enum, auto
+
+
+class ErrorCode(Enum):
+    """Error codes for MCP errors"""
+
+    InvalidRequest = auto()
+    InternalError = auto()
+
+
+class McpError(Exception):
+    """Base exception class for MCP errors"""
+
+    def __init__(self, code: ErrorCode, message: str):
+        self.code = code
+        self.message = message
+        super().__init__(message)
+
+
+DEFAULT_CONFIG = {
+    "server": {
+        "host": "localhost",
+        "port": 9621,
+        "model": "lightrag:latest",
+        "timeout": 300,
+        "max_retries": 1,
+        "retry_delay": 1,
+    },
+    "test_cases": {
+        "basic": {"query": "唐僧有几个徒弟"},
+        "generate": {"query": "电视剧西游记导演是谁"},
+    },
+}
+
+# Example conversation history for testing
+EXAMPLE_CONVERSATION = [
+    {"role": "user", "content": "你好"},
+    {"role": "assistant", "content": "你好!我是一个AI助手,很高兴为你服务。"},
+    {"role": "user", "content": "Who are you?"},
+    {"role": "assistant", "content": "I'm a Knowledge base query assistant."},
+]
+
+
+class OutputControl:
+    """Output control class, manages the verbosity of test output"""
+
+    _verbose: bool = False
+
+    @classmethod
+    def set_verbose(cls, verbose: bool) -> None:
+        cls._verbose = verbose
+
+    @classmethod
+    def is_verbose(cls) -> bool:
+        return cls._verbose
+
+
+@dataclass
+class ExecutionResult:
+    """Test execution result data class"""
+
+    name: str
+    success: bool
+    duration: float
+    error: Optional[str] = None
+    timestamp: str = ""
+
+    def __post_init__(self):
+        if not self.timestamp:
+            self.timestamp = datetime.now().isoformat()
+
+
+class ExecutionStats:
+    """Test execution statistics"""
+
+    def __init__(self):
+        self.results: List[ExecutionResult] = []
+        self.start_time = datetime.now()
+
+    def add_result(self, result: ExecutionResult):
+        self.results.append(result)
+
+    def export_results(self, path: str = "test_results.json"):
+        """Export test results to a JSON file
+        Args:
+            path: Output file path
+        """
+        results_data = {
+            "start_time": self.start_time.isoformat(),
+            "end_time": datetime.now().isoformat(),
+            "results": [asdict(r) for r in self.results],
+            "summary": {
+                "total": len(self.results),
+                "passed": sum(1 for r in self.results if r.success),
+                "failed": sum(1 for r in self.results if not r.success),
+                "total_duration": sum(r.duration for r in self.results),
+            },
+        }
+
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(results_data, f, ensure_ascii=False, indent=2)
+        print(f"\nTest results saved to: {path}")
+
+    def print_summary(self):
+        total = len(self.results)
+        passed = sum(1 for r in self.results if r.success)
+        failed = total - passed
+        duration = sum(r.duration for r in self.results)
+
+        print("\n=== Test Summary ===")
+        print(f"Start time: {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}")
+        print(f"Total duration: {duration:.2f} seconds")
+        print(f"Total tests: {total}")
+        print(f"Passed: {passed}")
+        print(f"Failed: {failed}")
+
+        if failed > 0:
+            print("\nFailed tests:")
+            for result in self.results:
+                if not result.success:
+                    print(f"- {result.name}: {result.error}")
+
+
+def make_request(
+    url: str, data: Dict[str, Any], stream: bool = False, check_status: bool = True
+) -> requests.Response:
+    """Send an HTTP request with retry mechanism
+    Args:
+        url: Request URL
+        data: Request data
+        stream: Whether to use streaming response
+        check_status: Whether to check HTTP status code (default: True)
+    Returns:
+        requests.Response: Response object
+
+    Raises:
+        requests.exceptions.RequestException: Request failed after all retries
+        requests.exceptions.HTTPError: HTTP status code is not 200 (when check_status is True)
+    """
+    server_config = CONFIG["server"]
+    max_retries = server_config["max_retries"]
+    retry_delay = server_config["retry_delay"]
+    timeout = server_config["timeout"]
+
+    for attempt in range(max_retries):
+        try:
+            response = requests.post(url, json=data, stream=stream, timeout=timeout)
+            if check_status and response.status_code != 200:
+                response.raise_for_status()
+            return response
+        except requests.exceptions.RequestException as e:
+            if attempt == max_retries - 1:  # Last retry
+                raise
+            print(f"\nRequest failed, retrying in {retry_delay} seconds: {str(e)}")
+            time.sleep(retry_delay)
+
+
+def load_config() -> Dict[str, Any]:
+    """Load configuration file
+
+    First try to load from config.json in the current directory,
+    if it doesn't exist, use the default configuration
+    Returns:
+        Configuration dictionary
+    """
+    config_path = Path("config.json")
+    if config_path.exists():
+        with open(config_path, "r", encoding="utf-8") as f:
+            return json.load(f)
+    return DEFAULT_CONFIG
+
+
+def print_json_response(data: Dict[str, Any], title: str = "", indent: int = 2) -> None:
+    """Format and print JSON response data
+    Args:
+        data: Data dictionary to print
+        title: Title to print
+        indent: Number of spaces for JSON indentation
+    """
+    if OutputControl.is_verbose():
+        if title:
+            print(f"\n=== {title} ===")
+        print(json.dumps(data, ensure_ascii=False, indent=indent))
+
+
+# Global configuration
+CONFIG = load_config()
+
+
+def get_base_url(endpoint: str = "chat") -> str:
+    """Return the base URL for specified endpoint
+    Args:
+        endpoint: API endpoint name (chat or generate)
+    Returns:
+        Complete URL for the endpoint
+    """
+    server = CONFIG["server"]
+    return f"http://{server['host']}:{server['port']}/api/{endpoint}"
+
+
+def create_chat_request_data(
+    content: str,
+    stream: bool = False,
+    model: str = None,
+    conversation_history: List[Dict[str, str]] = None,
+) -> Dict[str, Any]:
+    """Create chat request data
+    Args:
+        content: User message content
+        stream: Whether to use streaming response
+        model: Model name
+        conversation_history: List of previous conversation messages
+    Returns:
+        Dictionary containing complete chat request data
+    """
+    messages = conversation_history or []
+    messages.append({"role": "user", "content": content})
+
+    return {
+        "model": model or CONFIG["server"]["model"],
+        "messages": messages,
+        "stream": stream,
+    }
+
+
+def create_generate_request_data(
+    prompt: str,
+    system: str = None,
+    stream: bool = False,
+    model: str = None,
+    options: Dict[str, Any] = None,
+) -> Dict[str, Any]:
+    """Create generate request data
+    Args:
+        prompt: Generation prompt
+        system: System prompt
+        stream: Whether to use streaming response
+        model: Model name
+        options: Additional options
+    Returns:
+        Dictionary containing complete generate request data
+    """
+    data = {
+        "model": model or CONFIG["server"]["model"],
+        "prompt": prompt,
+        "stream": stream,
+    }
+    if system:
+        data["system"] = system
+    if options:
+        data["options"] = options
+    return data
+
+
+# Global test statistics
+STATS = ExecutionStats()
+
+
+def run_test(func: Callable, name: str) -> None:
+    """Run a test and record the results
+    Args:
+        func: Test function
+        name: Test name
+    """
+    start_time = time.time()
+    try:
+        func()
+        duration = time.time() - start_time
+        STATS.add_result(ExecutionResult(name, True, duration))
+    except Exception as e:
+        duration = time.time() - start_time
+        STATS.add_result(ExecutionResult(name, False, duration, str(e)))
+        raise
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_non_stream_chat() -> None:
+    """Test non-streaming call to /api/chat endpoint"""
+    url = get_base_url()
+
+    # Send request with conversation history
+    data = create_chat_request_data(
+        CONFIG["test_cases"]["basic"]["query"],
+        stream=False,
+        conversation_history=EXAMPLE_CONVERSATION,
+    )
+    response = make_request(url, data)
+
+    # Print response
+    if OutputControl.is_verbose():
+        print("\n=== Non-streaming call response ===")
+    response_json = response.json()
+
+    # Print response content
+    print_json_response(
+        {"model": response_json["model"], "message": response_json["message"]},
+        "Response content",
+    )
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_stream_chat() -> None:
+    """Test streaming call to /api/chat endpoint
+
+    Use JSON Lines format to process streaming responses, each line is a complete JSON object.
+    Response format:
+    {
+        "model": "lightrag:latest",
+        "created_at": "2024-01-15T00:00:00Z",
+        "message": {
+            "role": "assistant",
+            "content": "Partial response content",
+            "images": null
+        },
+        "done": false
+    }
+
+    The last message will contain performance statistics, with done set to true.
+    """
+    url = get_base_url()
+
+    # Send request with conversation history
+    data = create_chat_request_data(
+        CONFIG["test_cases"]["basic"]["query"],
+        stream=True,
+        conversation_history=EXAMPLE_CONVERSATION,
+    )
+    response = make_request(url, data, stream=True)
+
+    if OutputControl.is_verbose():
+        print("\n=== Streaming call response ===")
+    output_buffer = []
+    try:
+        for line in response.iter_lines():
+            if line:  # Skip empty lines
+                try:
+                    # Decode and parse JSON
+                    data = json.loads(line.decode("utf-8"))
+                    if data.get("done", True):  # If it's the completion marker
+                        if (
+                            "total_duration" in data
+                        ):  # Final performance statistics message
+                            # print_json_response(data, "Performance statistics")
+                            break
+                    else:  # Normal content message
+                        message = data.get("message", {})
+                        content = message.get("content", "")
+                        if content:  # Only collect non-empty content
+                            output_buffer.append(content)
+                            print(
+                                content, end="", flush=True
+                            )  # Print content in real-time
+                except json.JSONDecodeError:
+                    print("Error decoding JSON from response line")
+    finally:
+        response.close()  # Ensure the response connection is closed
+
+    # Print a newline
+    print()
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_query_modes() -> None:
+    """Test different query mode prefixes
+
+    Supported query modes:
+    - /local: Local retrieval mode, searches only in highly relevant documents
+    - /global: Global retrieval mode, searches across all documents
+    - /naive: Naive mode, does not use any optimization strategies
+    - /hybrid: Hybrid mode (default), combines multiple strategies
+    - /mix: Mix mode
+
+    Each mode will return responses in the same format, but with different retrieval strategies.
+    """
+    url = get_base_url()
+    modes = ["local", "global", "naive", "hybrid", "mix"]
+
+    for mode in modes:
+        if OutputControl.is_verbose():
+            print(f"\n=== Testing /{mode} mode ===")
+        data = create_chat_request_data(
+            f"/{mode} {CONFIG['test_cases']['basic']['query']}", stream=False
+        )
+
+        # Send request
+        response = make_request(url, data)
+        response_json = response.json()
+
+        # Print response content
+        print_json_response(
+            {"model": response_json["model"], "message": response_json["message"]}
+        )
+
+
+def create_error_test_data(error_type: str) -> Dict[str, Any]:
+    """Create request data for error testing
+    Args:
+        error_type: Error type, supported:
+            - empty_messages: Empty message list
+            - invalid_role: Invalid role field
+            - missing_content: Missing content field
+
+    Returns:
+        Request dictionary containing error data
+    """
+    error_data = {
+        "empty_messages": {"model": "lightrag:latest", "messages": [], "stream": True},
+        "invalid_role": {
+            "model": "lightrag:latest",
+            "messages": [{"invalid_role": "user", "content": "Test message"}],
+            "stream": True,
+        },
+        "missing_content": {
+            "model": "lightrag:latest",
+            "messages": [{"role": "user"}],
+            "stream": True,
+        },
+    }
+    return error_data.get(error_type, error_data["empty_messages"])
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_stream_error_handling() -> None:
+    """Test error handling for streaming responses
+
+    Test scenarios:
+    1. Empty message list
+    2. Message format error (missing required fields)
+
+    Error responses should be returned immediately without establishing a streaming connection.
+    The status code should be 4xx, and detailed error information should be returned.
+    """
+    url = get_base_url()
+
+    if OutputControl.is_verbose():
+        print("\n=== Testing streaming response error handling ===")
+
+    # Test empty message list
+    if OutputControl.is_verbose():
+        print("\n--- Testing empty message list (streaming) ---")
+    data = create_error_test_data("empty_messages")
+    response = make_request(url, data, stream=True, check_status=False)
+    print(f"Status code: {response.status_code}")
+    if response.status_code != 200:
+        print_json_response(response.json(), "Error message")
+    response.close()
+
+    # Test invalid role field
+    if OutputControl.is_verbose():
+        print("\n--- Testing invalid role field (streaming) ---")
+    data = create_error_test_data("invalid_role")
+    response = make_request(url, data, stream=True, check_status=False)
+    print(f"Status code: {response.status_code}")
+    if response.status_code != 200:
+        print_json_response(response.json(), "Error message")
+    response.close()
+
+    # Test missing content field
+    if OutputControl.is_verbose():
+        print("\n--- Testing missing content field (streaming) ---")
+    data = create_error_test_data("missing_content")
+    response = make_request(url, data, stream=True, check_status=False)
+    print(f"Status code: {response.status_code}")
+    if response.status_code != 200:
+        print_json_response(response.json(), "Error message")
+    response.close()
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_error_handling() -> None:
+    """Test error handling for non-streaming responses
+
+    Test scenarios:
+    1. Empty message list
+    2. Message format error (missing required fields)
+
+    Error response format:
+    {
+        "detail": "Error description"
+    }
+
+    All errors should return appropriate HTTP status codes and clear error messages.
+    """
+    url = get_base_url()
+
+    if OutputControl.is_verbose():
+        print("\n=== Testing error handling ===")
+
+    # Test empty message list
+    if OutputControl.is_verbose():
+        print("\n--- Testing empty message list ---")
+    data = create_error_test_data("empty_messages")
+    data["stream"] = False  # Change to non-streaming mode
+    response = make_request(url, data, check_status=False)
+    print(f"Status code: {response.status_code}")
+    print_json_response(response.json(), "Error message")
+
+    # Test invalid role field
+    if OutputControl.is_verbose():
+        print("\n--- Testing invalid role field ---")
+    data = create_error_test_data("invalid_role")
+    data["stream"] = False  # Change to non-streaming mode
+    response = make_request(url, data, check_status=False)
+    print(f"Status code: {response.status_code}")
+    print_json_response(response.json(), "Error message")
+
+    # Test missing content field
+    if OutputControl.is_verbose():
+        print("\n--- Testing missing content field ---")
+    data = create_error_test_data("missing_content")
+    data["stream"] = False  # Change to non-streaming mode
+    response = make_request(url, data, check_status=False)
+    print(f"Status code: {response.status_code}")
+    print_json_response(response.json(), "Error message")
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_non_stream_generate() -> None:
+    """Test non-streaming call to /api/generate endpoint"""
+    url = get_base_url("generate")
+    data = create_generate_request_data(
+        CONFIG["test_cases"]["generate"]["query"], stream=False
+    )
+
+    # Send request
+    response = make_request(url, data)
+
+    # Print response
+    if OutputControl.is_verbose():
+        print("\n=== Non-streaming generate response ===")
+    response_json = response.json()
+
+    # Print response content
+    print(json.dumps(response_json, ensure_ascii=False, indent=2))
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_stream_generate() -> None:
+    """Test streaming call to /api/generate endpoint"""
+    url = get_base_url("generate")
+    data = create_generate_request_data(
+        CONFIG["test_cases"]["generate"]["query"], stream=True
+    )
+
+    # Send request and get streaming response
+    response = make_request(url, data, stream=True)
+
+    if OutputControl.is_verbose():
+        print("\n=== Streaming generate response ===")
+    output_buffer = []
+    try:
+        for line in response.iter_lines():
+            if line:  # Skip empty lines
+                try:
+                    # Decode and parse JSON
+                    data = json.loads(line.decode("utf-8"))
+                    if data.get("done", True):  # If it's the completion marker
+                        if (
+                            "total_duration" in data
+                        ):  # Final performance statistics message
+                            break
+                    else:  # Normal content message
+                        content = data.get("response", "")
+                        if content:  # Only collect non-empty content
+                            output_buffer.append(content)
+                            print(
+                                content, end="", flush=True
+                            )  # Print content in real-time
+                except json.JSONDecodeError:
+                    print("Error decoding JSON from response line")
+    finally:
+        response.close()  # Ensure the response connection is closed
+
+    # Print a newline
+    print()
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_generate_with_system() -> None:
+    """Test generate with system prompt"""
+    url = get_base_url("generate")
+    data = create_generate_request_data(
+        CONFIG["test_cases"]["generate"]["query"],
+        system="你是一个知识渊博的助手",
+        stream=False,
+    )
+
+    # Send request
+    response = make_request(url, data)
+
+    # Print response
+    if OutputControl.is_verbose():
+        print("\n=== Generate with system prompt response ===")
+    response_json = response.json()
+
+    # Print response content
+    print_json_response(
+        {
+            "model": response_json["model"],
+            "response": response_json["response"],
+            "done": response_json["done"],
+        },
+        "Response content",
+    )
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_generate_error_handling() -> None:
+    """Test error handling for generate endpoint"""
+    url = get_base_url("generate")
+
+    # Test empty prompt
+    if OutputControl.is_verbose():
+        print("\n=== Testing empty prompt ===")
+    data = create_generate_request_data("", stream=False)
+    response = make_request(url, data, check_status=False)
+    print(f"Status code: {response.status_code}")
+    print_json_response(response.json(), "Error message")
+
+    # Test invalid options
+    if OutputControl.is_verbose():
+        print("\n=== Testing invalid options ===")
+    data = create_generate_request_data(
+        CONFIG["test_cases"]["basic"]["query"],
+        options={"invalid_option": "value"},
+        stream=False,
+    )
+    response = make_request(url, data, check_status=False)
+    print(f"Status code: {response.status_code}")
+    print_json_response(response.json(), "Error message")
+
+
+@pytest.mark.integration
+@pytest.mark.requires_api
+def test_generate_concurrent() -> None:
+    """Test concurrent generate requests"""
+    import asyncio
+    import aiohttp
+    from contextlib import asynccontextmanager
+
+    @asynccontextmanager
+    async def get_session():
+        async with aiohttp.ClientSession() as session:
+            yield session
+
+    async def make_request(session, prompt: str, request_id: int):
+        url = get_base_url("generate")
+        data = create_generate_request_data(prompt, stream=False)
+        try:
+            async with session.post(url, json=data) as response:
+                if response.status != 200:
+                    error_msg = (
+                        f"Request {request_id} failed with status {response.status}"
+                    )
+                    if OutputControl.is_verbose():
+                        print(f"\n{error_msg}")
+                    raise McpError(ErrorCode.InternalError, error_msg)
+                result = await response.json()
+                if "error" in result:
+                    error_msg = (
+                        f"Request {request_id} returned error: {result['error']}"
+                    )
+                    if OutputControl.is_verbose():
+                        print(f"\n{error_msg}")
+                    raise McpError(ErrorCode.InternalError, error_msg)
+                return result
+        except Exception as e:
+            error_msg = f"Request {request_id} failed: {str(e)}"
+            if OutputControl.is_verbose():
+                print(f"\n{error_msg}")
+            raise McpError(ErrorCode.InternalError, error_msg)
+
+    async def run_concurrent_requests():
+        prompts = ["第一个问题", "第二个问题", "第三个问题", "第四个问题", "第五个问题"]
+
+        async with get_session() as session:
+            tasks = [
+                make_request(session, prompt, i + 1) for i, prompt in enumerate(prompts)
+            ]
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+
+            success_results = []
+            error_messages = []
+
+            for i, result in enumerate(results):
+                if isinstance(result, Exception):
+                    error_messages.append(f"Request {i + 1} failed: {str(result)}")
+                else:
+                    success_results.append((i + 1, result))
+
+            if error_messages:
+                for req_id, result in success_results:
+                    if OutputControl.is_verbose():
+                        print(f"\nRequest {req_id} succeeded:")
+                        print_json_response(result)
+
+                error_summary = "\n".join(error_messages)
+                raise McpError(
+                    ErrorCode.InternalError,
+                    f"Some concurrent requests failed:\n{error_summary}",
+                )
+
+            return results
+
+    if OutputControl.is_verbose():
+        print("\n=== Testing concurrent generate requests ===")
+
+    # Run concurrent requests
+    try:
+        results = asyncio.run(run_concurrent_requests())
+        # all success, print out results
+        for i, result in enumerate(results, 1):
+            print(f"\nRequest {i} result:")
+            print_json_response(result)
+    except McpError:
+        # error message already printed
+        raise
+
+
+def get_test_cases() -> Dict[str, Callable]:
+    """Get all available test cases
+    Returns:
+        A dictionary mapping test names to test functions
+    """
+    return {
+        "non_stream": test_non_stream_chat,
+        "stream": test_stream_chat,
+        "modes": test_query_modes,
+        "errors": test_error_handling,
+        "stream_errors": test_stream_error_handling,
+        "non_stream_generate": test_non_stream_generate,
+        "stream_generate": test_stream_generate,
+        "generate_with_system": test_generate_with_system,
+        "generate_errors": test_generate_error_handling,
+        "generate_concurrent": test_generate_concurrent,
+    }
+
+
+def create_default_config():
+    """Create a default configuration file"""
+    config_path = Path("config.json")
+    if not config_path.exists():
+        with open(config_path, "w", encoding="utf-8") as f:
+            json.dump(DEFAULT_CONFIG, f, ensure_ascii=False, indent=2)
+        print(f"Default configuration file created: {config_path}")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command line arguments"""
+    parser = argparse.ArgumentParser(
+        description="LightRAG Ollama Compatibility Interface Testing",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Configuration file (config.json):
+  {
+    "server": {
+      "host": "localhost",      # Server address
+      "port": 9621,            # Server port
+      "model": "lightrag:latest" # Default model name
+    },
+    "test_cases": {
+      "basic": {
+        "query": "Test query",      # Basic query text
+        "stream_query": "Stream query" # Stream query text
+      }
+    }
+  }
+""",
+    )
+    parser.add_argument(
+        "-q",
+        "--quiet",
+        action="store_true",
+        help="Silent mode, only display test result summary",
+    )
+    parser.add_argument(
+        "-a",
+        "--ask",
+        type=str,
+        help="Specify query content, which will override the query settings in the configuration file",
+    )
+    parser.add_argument(
+        "--init-config", action="store_true", help="Create default configuration file"
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="",
+        help="Test result output file path, default is not to output to a file",
+    )
+    parser.add_argument(
+        "--tests",
+        nargs="+",
+        choices=list(get_test_cases().keys()) + ["all"],
+        default=["all"],
+        help="Test cases to run, options: %(choices)s. Use 'all' to run all tests （except error tests)",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    # Set output mode
+    OutputControl.set_verbose(not args.quiet)
+
+    # If query content is specified, update the configuration
+    if args.ask:
+        CONFIG["test_cases"]["basic"]["query"] = args.ask
+
+    # If specified to create a configuration file
+    if args.init_config:
+        create_default_config()
+        exit(0)
+
+    test_cases = get_test_cases()
+
+    try:
+        if "all" in args.tests:
+            # Run all tests except error handling tests
+            if OutputControl.is_verbose():
+                print("\n【Chat API Tests】")
+            run_test(test_non_stream_chat, "Non-streaming Chat Test")
+            run_test(test_stream_chat, "Streaming Chat Test")
+            run_test(test_query_modes, "Chat Query Mode Test")
+
+            if OutputControl.is_verbose():
+                print("\n【Generate API Tests】")
+            run_test(test_non_stream_generate, "Non-streaming Generate Test")
+            run_test(test_stream_generate, "Streaming Generate Test")
+            run_test(test_generate_with_system, "Generate with System Prompt Test")
+            run_test(test_generate_concurrent, "Generate Concurrent Test")
+        else:
+            # Run specified tests
+            for test_name in args.tests:
+                if OutputControl.is_verbose():
+                    print(f"\n【Running Test: {test_name}】")
+                run_test(test_cases[test_name], test_name)
+    except Exception as e:
+        print(f"\nAn error occurred: {str(e)}")
+    finally:
+        # Print test statistics
+        STATS.print_summary()
+        # If an output file path is specified, export the results
+        if args.output:
+            STATS.export_results(args.output)

tests/api/test_path_prefixes.py

@@ -0,0 +1,655 @@
+"""
+Integration tests for API and WebUI path prefix support via root_path.
+
+With the root_path approach, routes always stay at their natural paths
+(/docs, /health, /query, /documents/...). The api_prefix is passed to
+FastAPI's root_path parameter, which controls the servers URL in the
+OpenAPI spec for correct reverse proxy operation.
+"""
+
+import os
+import sys
+from unittest.mock import patch, MagicMock
+from fastapi.testclient import TestClient
+
+import pytest
+
+
+# Env vars that the project's `.env` may have populated (via load_dotenv at
+# import time of lightrag.api.config). Tests must be hermetic and not depend
+# on developer-local .env values, so we clear/override anything that affects
+# parse_args() / create_app().
+_ENV_VARS_TO_ISOLATE = (
+    "LLM_BINDING",
+    "EMBEDDING_BINDING",
+    "LLM_BINDING_HOST",
+    "LLM_BINDING_API_KEY",
+    "LLM_MODEL",
+    "EMBEDDING_BINDING_HOST",
+    "EMBEDDING_BINDING_API_KEY",
+    "EMBEDDING_MODEL",
+    "LIGHTRAG_API_PREFIX",
+    "LIGHTRAG_KV_STORAGE",
+    "LIGHTRAG_VECTOR_STORAGE",
+    "LIGHTRAG_GRAPH_STORAGE",
+    "LIGHTRAG_DOC_STATUS_STORAGE",
+)
+
+
+@pytest.fixture(autouse=True)
+def _isolate_env(monkeypatch):
+    """Isolate tests from developer-local .env pollution.
+
+    The lightrag.api.config module loads .env at import time, which can leave
+    bindings/hosts/keys in os.environ that mismatch what these tests assume.
+    Clear them, then set the minimal viable defaults (ollama bindings) so
+    create_app's binding validation passes without touching real services.
+    """
+    for var in _ENV_VARS_TO_ISOLATE:
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("LLM_BINDING", "ollama")
+    monkeypatch.setenv("EMBEDDING_BINDING", "ollama")
+
+
+@pytest.fixture
+def mock_args_api_prefix():
+    """Create mock args with API prefix."""
+    from lightrag.api.config import parse_args
+
+    original_argv = sys.argv.copy()
+    try:
+        sys.argv = ["lightrag-server", "--api-prefix", "/test-api"]
+        args = parse_args()
+        yield args
+    finally:
+        sys.argv = original_argv
+
+
+@pytest.fixture
+def mock_args_no_prefix():
+    """Create mock args without API prefix."""
+    from lightrag.api.config import parse_args
+
+    original_argv = sys.argv.copy()
+    try:
+        sys.argv = ["lightrag-server"]
+        args = parse_args()
+        yield args
+    finally:
+        sys.argv = original_argv
+
+
+class TestRootPathConfiguration:
+    """Test that root_path is set correctly on the FastAPI app."""
+
+    def test_root_path_set_when_prefix_provided(self, mock_args_api_prefix):
+        """Test app.root_path reflects api_prefix."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_api_prefix)
+            assert app.root_path == "/test-api"
+
+    def test_root_path_none_when_no_prefix(self, mock_args_no_prefix):
+        """Test app.root_path is not set when no prefix is configured."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_no_prefix)
+            # When no prefix, root_path is None (not passed to FastAPI)
+            # FastAPI stores None as-is, which means no root_path injection
+            assert not app.root_path
+
+
+class TestRoutesAtNaturalPaths:
+    """Test that routes stay at their natural paths regardless of root_path."""
+
+    def test_routes_accessible_at_both_paths_with_prefix(self, mock_args_api_prefix):
+        """With root_path, routes work at both prefixed and natural paths.
+
+        FastAPI injects root_path into the ASGI scope, and Starlette strips
+        it from the path before matching. So /test-api/docs and /docs both work.
+        """
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_api_prefix)
+            client = TestClient(app)
+
+            # Natural path works
+            response = client.get("/docs")
+            assert response.status_code == 200
+
+            response = client.get("/openapi.json")
+            assert response.status_code == 200
+
+            # Prefixed path also works (FastAPI strips root_path from scope)
+            response = client.get("/test-api/docs")
+            assert response.status_code == 200
+
+            response = client.get("/test-api/openapi.json")
+            assert response.status_code == 200
+
+    def test_document_routes_at_natural_path(self, mock_args_api_prefix):
+        """Test document routes are at /documents/ (their router-level prefix)."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_api_prefix)
+            client = TestClient(app)
+
+            response = client.post(
+                "/documents/paginated",
+                json={},
+                headers={"Authorization": "Bearer test"},
+            )
+            # The route is mounted; the mocked LightRAG may cause 401/422/500,
+            # but a missing route (404) or wrong method (405) means routing
+            # itself broke and is what we want to catch here.
+            assert response.status_code not in (404, 405)
+
+    def test_routes_accessible_at_root_no_prefix(self, mock_args_no_prefix):
+        """Test routes are at root when no prefix is set (default)."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_no_prefix)
+            client = TestClient(app)
+
+            # API docs accessible at root
+            response = client.get("/docs")
+            assert response.status_code == 200
+
+            # openapi.json at root
+            response = client.get("/openapi.json")
+            assert response.status_code == 200
+
+            # Prefixed paths return 404 when no root_path is configured
+            response = client.get("/test-api/docs")
+            assert response.status_code == 404
+
+
+class TestOpenAPISpecIntegration:
+    """Test that OpenAPI spec uses root_path for servers URL."""
+
+    def test_openapi_spec_has_servers_url_with_prefix(self, mock_args_api_prefix):
+        """Test OpenAPI spec servers URL includes the prefix via root_path."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_api_prefix)
+            client = TestClient(app)
+
+            # OpenAPI JSON is served at the natural path
+            response = client.get("/openapi.json")
+            assert response.status_code == 200
+            spec = response.json()
+
+            # Servers URL should include the prefix
+            servers = spec.get("servers", [])
+            assert len(servers) > 0, (
+                "OpenAPI spec should have servers entry when root_path is set"
+            )
+            assert servers[0].get("url") == "/test-api", (
+                f"Expected servers URL to be exactly /test-api, got: {servers[0].get('url')}"
+            )
+
+    def test_openapi_spec_no_servers_without_prefix(self, mock_args_no_prefix):
+        """Test OpenAPI spec has no servers entry when no root_path."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_no_prefix)
+            client = TestClient(app)
+
+            response = client.get("/openapi.json")
+            assert response.status_code == 200
+            spec = response.json()
+
+            # No servers when root_path is None/empty
+            assert "servers" not in spec or spec["servers"] is None
+
+    def test_openapi_spec_paths_at_natural_paths(self, mock_args_api_prefix):
+        """Test OpenAPI spec paths are at natural paths (not prefixed)."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_api_prefix)
+            client = TestClient(app)
+
+            response = client.get("/openapi.json")
+            assert response.status_code == 200
+            spec = response.json()
+            paths = spec.get("paths", {})
+
+            # Paths should be at natural paths
+            for path in paths:
+                if path == "/":
+                    continue
+                assert not path.startswith("/test-api/"), (
+                    f"Path {path} should not be prefixed with /test-api/ in root_path mode"
+                )
+
+
+class TestWebUIPrefixIntegration:
+    """Test that the WebUI is served at the expected (fixed) /webui path,
+    composed with `root_path` when an API prefix is set."""
+
+    def test_webui_at_prefixed_path(self, mock_args_api_prefix):
+        """With root_path="/test-api" the WebUI lives at /test-api/webui/
+        because FastAPI injects root_path into the ASGI scope."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_api_prefix)
+            client = TestClient(app)
+
+            response = client.get("/test-api/webui/")
+            assert response.status_code in [200, 307]
+
+    def test_webui_without_api_prefix(self, mock_args_no_prefix):
+        """Without an API prefix the WebUI is served at /webui/."""
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            from lightrag.api.lightrag_server import create_app
+
+            app = create_app(mock_args_no_prefix)
+            client = TestClient(app)
+
+            response = client.get("/webui/")
+            assert response.status_code in [200, 307]
+
+
+class TestEnvironmentVariables:
+    """Test that environment variables are read correctly."""
+
+    def test_env_api_prefix(self):
+        """Test LIGHTRAG_API_PREFIX environment variable."""
+        from lightrag.api.config import get_env_value
+
+        os.environ["LIGHTRAG_API_PREFIX"] = "unit-test-back/api"
+        try:
+            value = get_env_value("LIGHTRAG_API_PREFIX", "")
+            assert value == "unit-test-back/api"
+        finally:
+            del os.environ["LIGHTRAG_API_PREFIX"]
+
+
+class TestPathNormalization:
+    """User input for `--api-prefix` may contain trailing slashes, a missing
+    leading slash, or be just '/'. create_app must canonicalize these before
+    passing to FastAPI's `root_path`, which doesn't accept arbitrary strings."""
+
+    def _build(self, *cli_args):
+        # sys.argv must be the lightrag-server form *before* lightrag_server is
+        # imported, because importing lightrag.api.utils_api evaluates
+        # `global_args.whitelist_paths` at module top level, which triggers
+        # parse_args() against whatever sys.argv currently holds.
+        original_argv = sys.argv.copy()
+        try:
+            sys.argv = ["lightrag-server", *cli_args]
+            from lightrag.api.config import parse_args
+            from lightrag.api.lightrag_server import create_app
+
+            args = parse_args()
+            with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+                mock_rag.return_value = MagicMock()
+                return create_app(args)
+        finally:
+            sys.argv = original_argv
+
+    def test_api_prefix_slash_only_treated_as_empty(self):
+        """`--api-prefix /` is degenerate; must collapse to no prefix."""
+        app = self._build("--api-prefix", "/")
+        assert not app.root_path
+
+    def test_api_prefix_trailing_slash_stripped(self):
+        """Trailing slash on api_prefix is stripped to keep OpenAPI servers
+        URL clean and avoid double-slash artifacts."""
+        app = self._build("--api-prefix", "/api/v1/")
+        assert app.root_path == "/api/v1"
+
+    def test_api_prefix_missing_leading_slash_added(self):
+        app = self._build("--api-prefix", "api/v1")
+        assert app.root_path == "/api/v1"
+
+
+class TestRuntimeConfigInjection:
+    """End-to-end tests for the WebUI runtime-config injection.
+
+    The browser-visible URL prefixes are no longer baked into the bundle.
+    Instead, the server replaces a placeholder comment in index.html with
+    a `<script>window.__LIGHTRAG_CONFIG__ = {...}</script>` snippet on
+    every HTML response, so one build can serve any reverse-proxy mount.
+
+    These tests stage a minimal index.html in a tmp dir, patch
+    `lightrag_server.__file__` so both `check_frontend_build()` and the
+    static-files mount resolve to it, then drive the app via TestClient
+    and assert that the body contains the expected injected JSON.
+    """
+
+    PLACEHOLDER = "<!-- __LIGHTRAG_RUNTIME_CONFIG__ -->"
+
+    def _stage_index_html(self, tmp_path, *, with_placeholder=True):
+        """Mirror what Vite emits: a tiny index.html with the runtime-config
+        placeholder in <head> plus a hashed asset reference.
+
+        with_placeholder=False simulates a stale build that pre-dates this
+        feature — the server should serve it untouched, not crash.
+        """
+        webui_dir = tmp_path / "webui"
+        webui_dir.mkdir()
+        placeholder = self.PLACEHOLDER if with_placeholder else ""
+        (webui_dir / "index.html").write_text(
+            "<!doctype html><html><head>"
+            f"{placeholder}"
+            '<script type="module" crossorigin src="./assets/index-X.js"></script>'
+            '<link rel="stylesheet" href="./assets/index-X.css">'
+            "</head><body><div id=root></div></body></html>",
+            encoding="utf-8",
+        )
+
+    def _build_app(self, tmp_path, monkeypatch, *cli_args):
+        # Force benign argv before the (potentially fresh) module import —
+        # see TestPathNormalization._build for the rationale.
+        monkeypatch.setattr(sys, "argv", ["lightrag-server", *cli_args])
+        from lightrag.api.config import parse_args
+        from lightrag.api import lightrag_server
+        from lightrag.api.lightrag_server import create_app
+
+        # Redirect both check_frontend_build() and the StaticFiles mount to
+        # our staged tmp directory.
+        monkeypatch.setattr(
+            lightrag_server, "__file__", str(tmp_path / "lightrag_server.py")
+        )
+
+        args = parse_args()
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            return create_app(args)
+
+    def test_injection_populates_window_config_with_prefix(self, tmp_path, monkeypatch):
+        """With api_prefix=/site01, the injected script must carry both the
+        api prefix and the composed webui prefix the browser will see."""
+        self._stage_index_html(tmp_path)
+        app = self._build_app(tmp_path, monkeypatch, "--api-prefix", "/site01")
+        client = TestClient(app)
+
+        response = client.get("/site01/webui/")
+        assert response.status_code == 200
+        body = response.text
+
+        # Placeholder must be gone and replaced with the runtime config.
+        assert self.PLACEHOLDER not in body
+        assert "window.__LIGHTRAG_CONFIG__" in body
+        assert '"apiPrefix": "/site01"' in body or '"apiPrefix":"/site01"' in body
+        assert (
+            '"webuiPrefix": "/site01/webui/"' in body
+            or '"webuiPrefix":"/site01/webui/"' in body
+        )
+
+    def test_injection_default_prefixes_when_unconfigured(self, tmp_path, monkeypatch):
+        """No CLI flags → empty api prefix and the default webui mount.
+        The injected JSON must reflect this so the SPA falls through to
+        same-origin requests."""
+        self._stage_index_html(tmp_path)
+        app = self._build_app(tmp_path, monkeypatch)
+        client = TestClient(app)
+
+        response = client.get("/webui/")
+        assert response.status_code == 200
+        body = response.text
+
+        assert '"apiPrefix": ""' in body or '"apiPrefix":""' in body
+        assert '"webuiPrefix": "/webui/"' in body or '"webuiPrefix":"/webui/"' in body
+
+    def test_missing_placeholder_serves_original_html(self, tmp_path, monkeypatch):
+        """Older builds without the placeholder must still serve cleanly —
+        no 500, no partial replacement, just the original HTML. Avoids
+        breaking anyone whose pre-built bundle is in use during an upgrade."""
+        self._stage_index_html(tmp_path, with_placeholder=False)
+        app = self._build_app(tmp_path, monkeypatch)
+        client = TestClient(app)
+
+        response = client.get("/webui/")
+        assert response.status_code == 200
+        # No placeholder was present, so no injected script either.
+        assert "window.__LIGHTRAG_CONFIG__" not in response.text
+
+    def test_injection_idempotent_across_requests(self, tmp_path, monkeypatch):
+        """Each request reads the file fresh; the placeholder must be
+        present in the *file* even after replies (we don't mutate it)."""
+        self._stage_index_html(tmp_path)
+        app = self._build_app(tmp_path, monkeypatch, "--api-prefix", "/abc")
+        client = TestClient(app)
+
+        first = client.get("/abc/webui/").text
+        second = client.get("/abc/webui/").text
+        assert first == second
+        # Source file untouched.
+        on_disk = (tmp_path / "webui" / "index.html").read_text(encoding="utf-8")
+        assert self.PLACEHOLDER in on_disk
+
+    def test_html_response_keeps_no_cache_headers(self, tmp_path, monkeypatch):
+        """Injection must not regress the existing no-cache behaviour for
+        HTML — otherwise an updated runtime config could be cached client-
+        side and never picked up."""
+        self._stage_index_html(tmp_path)
+        app = self._build_app(tmp_path, monkeypatch, "--api-prefix", "/x")
+        client = TestClient(app)
+
+        response = client.get("/x/webui/")
+        assert response.status_code == 200
+        cache_control = response.headers.get("cache-control", "")
+        assert "no-cache" in cache_control
+        assert "no-store" in cache_control
+
+
+class TestUvicornRootPathSemantics:
+    """Lock in the deployment contract that both proxy-strip and verbatim
+    forwarding work through FastAPI's app-level ``root_path`` plus a
+    ``_RootPathNormalizationMiddleware`` — without ever passing
+    ``root_path`` through to uvicorn/gunicorn.
+
+    Background:
+
+      - uvicorn builds ``scope["path"] = uvicorn.root_path + <incoming http
+        path>`` and ``scope["root_path"] = uvicorn.root_path``
+        (uvicorn/protocols/http/h11_impl.py).
+      - FastAPI's ``__call__`` overrides ``scope["root_path"]`` from
+        ``app.root_path`` but does NOT touch ``scope["path"]``
+        (fastapi/applications.py:1131-1134).
+      - Starlette's ``Mount.matches`` mutates the child scope's root_path
+        to ``original + matched_path`` and again leaves ``scope["path"]``
+        untouched (starlette/routing.py:401-432). The inner sub-app
+        (e.g. StaticFiles) then sees ``scope["path"]`` and
+        ``scope["root_path"]`` that do not overlap, and 404s.
+
+    Concrete failure mode without the middleware (proxy strips /site01,
+    backend sees ``/webui/``):
+
+        outer get_route_path:  /webui/ does not start with /site01 → /webui/
+        Mount.matches:         path_regex "^/webui/(?P<path>.*)$" matches
+                               child scope.root_path = /site01/webui
+                               child scope.path      = /webui/  (unchanged)
+        StaticFiles.get_path:  /webui/ does not start with /site01/webui →
+                               returns /webui/ → looked up as filename
+                               "webui" inside the webui static dir → 404
+
+    ``_RootPathNormalizationMiddleware`` runs before routing and prepends
+    ``root_path`` to ``scope["path"]`` whenever the latter does not already
+    start with it. This makes both modes converge to the canonical ASGI
+    form (path always contains root_path) without requiring uvicorn's
+    --root-path — which would break verbatim mode by double-prefixing.
+    """
+
+    @staticmethod
+    async def _call_with_scope(app, http_path, *, uvicorn_root_path=""):
+        """Drive the ASGI app the way uvicorn would.
+
+        Simulates uvicorn's scope construction so we can catch the
+        path-doubling bug that TestClient hides.
+        """
+        full_path = uvicorn_root_path + http_path
+        scope = {
+            "type": "http",
+            "asgi": {"version": "3.0", "spec_version": "2.3"},
+            "http_version": "1.1",
+            "method": "GET",
+            "scheme": "http",
+            "server": ("testserver", 80),
+            "client": ("testclient", 50000),
+            "root_path": uvicorn_root_path,
+            "path": full_path,
+            "raw_path": full_path.encode("ascii"),
+            "query_string": b"",
+            "headers": [],
+            "state": {},
+        }
+        status = {"code": None}
+
+        async def receive():
+            return {"type": "http.request", "body": b"", "more_body": False}
+
+        async def send(message):
+            if message["type"] == "http.response.start":
+                status["code"] = message["status"]
+
+        await app(scope, receive, send)
+        return status["code"]
+
+    def _build_app_with_prefix(self, prefix):
+        original_argv = sys.argv.copy()
+        try:
+            sys.argv = ["lightrag-server", "--api-prefix", prefix]
+            from lightrag.api.config import parse_args
+            from lightrag.api.lightrag_server import create_app
+
+            args = parse_args()
+            with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+                mock_rag.return_value = MagicMock()
+                return create_app(args)
+        finally:
+            sys.argv = original_argv
+
+    @pytest.mark.asyncio
+    async def test_route_strip_mode_matches(self):
+        """Plain Route, proxy-strip mode: backend receives /openapi.json."""
+        app = self._build_app_with_prefix("/site01")
+        status = await self._call_with_scope(app, "/openapi.json")
+        assert status == 200
+
+    @pytest.mark.asyncio
+    async def test_route_verbatim_mode_matches(self):
+        """Plain Route, verbatim mode: backend receives /site01/openapi.json.
+
+        Locks in the contract that PR #3128's original fix broke: setting
+        uvicorn root_path would have made this case 404 by doubling the
+        prefix in scope["path"].
+        """
+        app = self._build_app_with_prefix("/site01")
+        status = await self._call_with_scope(app, "/site01/openapi.json")
+        assert status == 200
+
+    @pytest.mark.asyncio
+    async def test_mount_strip_mode_matches(self):
+        """WebUI Mount, proxy-strip mode: backend receives /webui/.
+
+        This is the bug the middleware fixes. Without normalization,
+        StaticFiles.get_path sees path=/webui/ and root_path=/site01/webui
+        (mutated by Mount.matches) and serves the literal "webui" filename
+        lookup → 404. With normalization, scope.path becomes
+        /site01/webui/ before routing and the lookup resolves to
+        index.html.
+        """
+        app = self._build_app_with_prefix("/site01")
+        status = await self._call_with_scope(app, "/webui/")
+        assert status == 200
+
+    @pytest.mark.asyncio
+    async def test_mount_verbatim_mode_matches(self):
+        """WebUI Mount, verbatim mode: backend receives /site01/webui/.
+
+        Already canonical; the middleware is a no-op for this case. The
+        test guards against accidentally regressing verbatim while fixing
+        strip — symmetric to test_route_verbatim_mode_matches but exercises
+        the Mount path with its nested ``get_route_path`` resolution.
+        """
+        app = self._build_app_with_prefix("/site01")
+        status = await self._call_with_scope(app, "/site01/webui/")
+        assert status == 200
+
+    @pytest.mark.asyncio
+    async def test_simulated_uvicorn_root_path_breaks_verbatim(self):
+        """Document the failure mode that the revert prevents.
+
+        If uvicorn's root_path were also "/site01", uvicorn would build
+        scope.path = "/site01" + "/site01/openapi.json". Starlette strips
+        one prefix; the remaining "/site01/openapi.json" has no route.
+        """
+        app = self._build_app_with_prefix("/site01")
+        status = await self._call_with_scope(
+            app, "/site01/openapi.json", uvicorn_root_path="/site01"
+        )
+        assert status == 404
+
+    def test_launcher_does_not_pass_root_path_to_uvicorn(self, monkeypatch, tmp_path):
+        """Guard against re-adding root_path to the uvicorn launcher kwargs.
+
+        Mocks uvicorn.run and exercises lightrag_server.main() far enough
+        to capture the config dict, then asserts root_path is absent.
+        """
+        monkeypatch.setenv("LIGHTRAG_API_PREFIX", "/site01")
+        monkeypatch.setattr(sys, "argv", ["lightrag-server"])
+
+        from lightrag.api import lightrag_server
+
+        captured = {}
+
+        def fake_run(**kwargs):
+            captured.update(kwargs)
+
+        monkeypatch.setattr(lightrag_server, "uvicorn", MagicMock(run=fake_run))
+        monkeypatch.setattr(lightrag_server, "check_env_file", lambda: True)
+        monkeypatch.setattr(
+            lightrag_server, "check_and_install_dependencies", lambda: None
+        )
+        monkeypatch.setattr(lightrag_server, "configure_logging", lambda: None)
+        monkeypatch.setattr(lightrag_server, "update_uvicorn_mode_config", lambda: None)
+        monkeypatch.setattr(lightrag_server, "display_splash_screen", lambda *_: None)
+        with patch("lightrag.api.lightrag_server.LightRAG") as mock_rag:
+            mock_rag.return_value = MagicMock()
+            # Re-parse args under the patched env so global_args picks up the prefix.
+            from lightrag.api.config import parse_args, global_args as _ga
+
+            new_args = parse_args()
+            for attr in vars(new_args):
+                setattr(_ga, attr, getattr(new_args, attr))
+
+            lightrag_server.main()
+
+        assert "root_path" not in captured, (
+            "uvicorn_config must not include root_path; rely on FastAPI's "
+            "app-level root_path only — see TestUvicornRootPathSemantics docstring."
+        )
+
+    def test_gunicorn_uses_upstream_uvicorn_worker(self):
+        """Symmetric guard for the multi-worker launcher.
+
+        gunicorn_config.worker_class must remain the upstream
+        ``uvicorn.workers.UvicornWorker`` — a custom subclass injecting
+        root_path via CONFIG_KWARGS would re-introduce the same
+        path-doubling regression in worker processes.
+        """
+        from lightrag.api import gunicorn_config
+
+        assert gunicorn_config.worker_class == "uvicorn.workers.UvicornWorker"

tests/chunker/test_chunker_recursive_character.py

@@ -0,0 +1,254 @@
+"""Unit tests for ``chunking_by_recursive_character`` (process_options=R)."""
+
+import pytest
+
+pytest.importorskip("langchain_text_splitters")
+from langchain_text_splitters import RecursiveCharacterTextSplitter  # noqa: E402
+
+from lightrag.chunker import chunking_by_recursive_character  # noqa: E402
+from lightrag.chunker.recursive_character import _split_text_with_spans  # noqa: E402
+from lightrag.utils import Tokenizer, TokenizerInterface  # noqa: E402
+
+
+class _CharTokenizer(TokenizerInterface):
+    """1 char ≈ 1 token; lets assertions reason in terms of input length."""
+
+    def encode(self, content: str):
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens):
+        return "".join(chr(t) for t in tokens)
+
+
+def _tok() -> Tokenizer:
+    return Tokenizer("char-tokenizer", _CharTokenizer())
+
+
+@pytest.mark.offline
+def test_empty_input_returns_empty_list():
+    chunks = chunking_by_recursive_character(_tok(), "")
+    assert chunks == []
+
+
+@pytest.mark.offline
+def test_short_input_single_chunk():
+    body = "Para A.\n\nPara B."
+    chunks = chunking_by_recursive_character(_tok(), body, chunk_token_size=1000)
+
+    assert len(chunks) == 1
+    assert chunks[0]["content"] == body
+    assert chunks[0]["_source_span"] == {"start": 0, "end": len(body)}
+    assert chunks[0]["chunk_order_index"] == 0
+    assert chunks[0]["tokens"] == len(body)
+
+
+@pytest.mark.offline
+def test_paragraph_separator_used_first():
+    """``\\n\\n`` is the first separator in the default cascade — three
+    paragraphs that each fit under the cap should split exactly there."""
+    body = "Alpha section.\n\nBeta section.\n\nGamma section."
+    chunks = chunking_by_recursive_character(
+        _tok(),
+        body,
+        chunk_token_size=20,
+        chunk_overlap_token_size=0,
+    )
+
+    assert [c["chunk_order_index"] for c in chunks] == list(range(len(chunks)))
+    assert all(c["content"].strip() for c in chunks)
+    # Reconstructed (joined with the splitter's separator semantics) must
+    # at least contain each original paragraph as a substring.
+    joined = "\n\n".join(c["content"] for c in chunks)
+    for para in ("Alpha section.", "Beta section.", "Gamma section."):
+        assert para in joined
+
+
+@pytest.mark.offline
+def test_token_field_matches_tokenizer_encode_length():
+    chunks = chunking_by_recursive_character(
+        _tok(),
+        "X" * 50 + "\n\n" + "Y" * 50,
+        chunk_token_size=40,
+        chunk_overlap_token_size=5,
+    )
+    tok = _tok()
+    for c in chunks:
+        assert c["tokens"] == len(tok.encode(c["content"]))
+
+
+@pytest.mark.offline
+def test_custom_separators_are_honored():
+    body = "alpha|beta|gamma|delta"
+    chunks = chunking_by_recursive_character(
+        _tok(),
+        body,
+        chunk_token_size=10,
+        chunk_overlap_token_size=0,
+        separators=["|", ""],
+    )
+    contents = [c["content"] for c in chunks]
+    # With "|" as the primary separator and a 10-token cap, each 5-char
+    # token name must land in its own chunk.
+    assert any("alpha" in c for c in contents)
+    assert any("delta" in c for c in contents)
+    # Every chunk fits the cap.
+    for c in chunks:
+        assert c["tokens"] <= 10
+
+
+@pytest.mark.offline
+def test_recursive_chunks_carry_exact_source_spans_with_overlap():
+    body = "Alpha section.\n\nBeta section.\n\nGamma section."
+    chunks = chunking_by_recursive_character(
+        _tok(),
+        body,
+        chunk_token_size=22,
+        chunk_overlap_token_size=6,
+    )
+
+    assert len(chunks) >= 2
+    for chunk in chunks:
+        span = chunk["_source_span"]
+        assert body[span["start"] : span["end"]] == chunk["content"]
+
+
+@pytest.mark.offline
+def test_long_unique_text_keeps_every_source_span():
+    """Every chunk of a long, non-repeating document must carry a span.
+
+    Regression for the LangChain ``add_start_index`` unit mismatch: with a
+    token-based length function its search cursor overshot each chunk's true
+    start, dropping ``_source_span`` (and failing span-first backfill) on
+    roughly half of a unique document's chunks.
+    """
+    body = " ".join(f"word{i}" for i in range(800))
+    chunks = chunking_by_recursive_character(
+        _tok(),
+        body,
+        chunk_token_size=50,
+        chunk_overlap_token_size=10,
+    )
+
+    assert len(chunks) > 5
+    for chunk in chunks:
+        span = chunk["_source_span"]
+        assert body[span["start"] : span["end"]] == chunk["content"]
+
+
+@pytest.mark.offline
+def test_repeated_text_spans_tile_whole_document():
+    """Spans over heavily repeated text must tile the document, not cluster.
+
+    Regression for the repeated-content ambiguity: a naive forward ``find``
+    matches every identical overlapping chunk to the nearest repetition, so all
+    spans collapse into the document head and the tail gets no provenance. The
+    offset-aware splitter mirror must instead advance spans across the whole text.
+    """
+    unit = "ABCDE fghij. "
+    body = unit * 60
+    chunks = chunking_by_recursive_character(
+        _tok(),
+        body,
+        chunk_token_size=40,
+        chunk_overlap_token_size=10,
+    )
+
+    spans = [c["_source_span"] for c in chunks]
+    assert len(spans) == len(chunks)  # no chunk lost its span
+
+    # Each span is exact and starts are monotonically non-decreasing.
+    for chunk, span in zip(chunks, spans):
+        assert body[span["start"] : span["end"]] == chunk["content"]
+    starts = [s["start"] for s in spans]
+    assert starts == sorted(starts)
+
+    # Spans reach the document tail rather than clustering at the head. Under
+    # the old nearest-repetition behaviour the furthest end stalled near the
+    # first few hundred chars; tiling reaches within one ``unit`` of the end.
+    assert max(s["end"] for s in spans) >= len(body) - len(unit)
+
+
+@pytest.mark.offline
+def test_span_mirror_matches_langchain_split_text():
+    """Drift guard for the offset-aware mirror of LangChain's splitter."""
+    body = (
+        "Alpha repeats.\n\n"
+        "ff hh aa hh ee\n"
+        "ff hh aa hh ee\n\n"
+        "Tail repeats. Tail repeats. Tail repeats."
+    )
+    tok = _tok()
+
+    def length_function(text: str) -> int:
+        return len(tok.encode(text))
+
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=36,
+        chunk_overlap=14,
+        length_function=length_function,
+        strip_whitespace=True,
+    )
+
+    mirrored = _split_text_with_spans(
+        body,
+        base_offset=0,
+        separators=splitter._separators,
+        chunk_size=splitter._chunk_size,
+        chunk_overlap=splitter._chunk_overlap,
+        length_function=length_function,
+        keep_separator=splitter._keep_separator,
+        is_separator_regex=splitter._is_separator_regex,
+        strip_whitespace=splitter._strip_whitespace,
+    )
+
+    assert [piece for piece, _, _ in mirrored] == splitter.split_text(body)
+    for piece, start, end in mirrored:
+        assert body[start:end] == piece
+
+
+@pytest.mark.offline
+def test_chunk_repeating_an_earlier_block_maps_to_its_own_block():
+    """A chunk whose text also appears in an earlier, different block must map
+    to *its own* occurrence, not the earlier copy.
+
+    Regression for cross-block duplicate provenance: merged text ``"ff aa\\n\\naa"``
+    splits into ``["ff aa", "aa"]``. The second chunk ``"aa"`` is block 2, so its
+    span must point at the *second* ``"aa"`` (offset 7), not the ``"aa"`` inside
+    ``"ff aa"`` (offset 3). The earlier locator anchored on a predicted offset
+    and snapped to the nearest copy, silently emitting the wrong — but
+    legal-looking — span that strict backfill would trust.
+    """
+    merged = "ff aa\n\naa"
+    chunks = chunking_by_recursive_character(
+        _tok(),
+        merged,
+        chunk_token_size=5,
+        chunk_overlap_token_size=1,
+    )
+
+    assert [c["content"] for c in chunks] == ["ff aa", "aa"]
+    assert chunks[0]["_source_span"] == {"start": 0, "end": 5}
+    # The second "aa" lives in block 2 at offset 7, after the "\n\n" separator —
+    # not the "aa" embedded in "ff aa" at offset 3.
+    assert chunks[1]["_source_span"] == {"start": 7, "end": 9}
+    assert merged[7:9] == "aa"
+
+
+@pytest.mark.offline
+def test_repeated_block_window_does_not_shift_to_previous_duplicate():
+    """A repeated multi-block chunk must not slide back into the prior duplicate."""
+    block = "ff hh aa hh ee"
+    merged = "\n\n".join([block] * 4)
+    chunks = chunking_by_recursive_character(
+        _tok(),
+        merged,
+        chunk_token_size=36,
+        chunk_overlap_token_size=14,
+    )
+
+    assert [c["content"] for c in chunks] == [
+        f"{block}\n\n{block}",
+        f"{block}\n\n{block}",
+    ]
+    assert chunks[0]["_source_span"] == {"start": 0, "end": 30}
+    assert chunks[1]["_source_span"] == {"start": 32, "end": 62}

tests/chunker/test_chunker_semantic_vector.py

@@ -0,0 +1,232 @@
+"""Unit tests for ``chunking_by_semantic_vector`` (process_options=V)."""
+
+import asyncio
+import logging
+
+import numpy as np
+import pytest
+
+pytest.importorskip("langchain_experimental")
+
+from lightrag.chunker import chunking_by_semantic_vector  # noqa: E402
+from lightrag.utils import EmbeddingFunc, Tokenizer, TokenizerInterface  # noqa: E402
+
+
+class _CharTokenizer(TokenizerInterface):
+    """1 char ≈ 1 token."""
+
+    def encode(self, content: str):
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens):
+        return "".join(chr(t) for t in tokens)
+
+
+def _tok() -> Tokenizer:
+    return Tokenizer("char-tokenizer", _CharTokenizer())
+
+
+def _make_deterministic_embedding(dim: int = 8) -> EmbeddingFunc:
+    """A toy async embedding func that hashes each input text into a
+    stable unit vector — enough to drive SemanticChunker without needing
+    a real model."""
+
+    async def _embed(texts, **kwargs):
+        rng = np.random.default_rng(seed=0)
+        # Use a simple hash → seeded rng to get reproducible vectors per text.
+        rows = []
+        for text in texts:
+            seed = abs(hash(text)) % (2**32)
+            rng = np.random.default_rng(seed=seed)
+            vec = rng.normal(size=dim).astype(np.float32)
+            vec /= np.linalg.norm(vec) or 1.0
+            rows.append(vec)
+        return np.vstack(rows)
+
+    return EmbeddingFunc(embedding_dim=dim, max_token_size=4096, func=_embed)
+
+
+@pytest.mark.offline
+def test_v_chunker_runs_with_stub_embedding():
+    """Async chunker should split a multi-sentence body into ≥1 chunk
+    when given a working embedding func."""
+    body = (
+        "Quantum mechanics describes nature at small scales. "
+        "It contradicts classical intuition. "
+        "Bread is baked from flour. "
+        "Sourdough requires a long fermentation. "
+    )
+
+    async def _run():
+        chunks = await chunking_by_semantic_vector(
+            _tok(),
+            body,
+            chunk_token_size=200,
+            embedding_func=_make_deterministic_embedding(),
+        )
+        return chunks
+
+    chunks = asyncio.run(_run())
+
+    assert len(chunks) >= 1
+    # Each chunk dict has the canonical schema.
+    assert all({"tokens", "content", "chunk_order_index"} <= set(c) for c in chunks)
+    # chunk_order_index is contiguous starting at 0.
+    assert [c["chunk_order_index"] for c in chunks] == list(range(len(chunks)))
+    # No empty content rows.
+    assert all(c["content"].strip() for c in chunks)
+    for chunk in chunks:
+        span = chunk["_source_span"]
+        assert body[span["start"] : span["end"]] == chunk["content"]
+
+
+class _ListHandler(logging.Handler):
+    def __init__(self) -> None:
+        super().__init__()
+        self.records: list[logging.LogRecord] = []
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self.records.append(record)
+
+
+@pytest.mark.offline
+def test_v_chunker_falls_back_to_recursive_when_no_embedding():
+    """When ``embedding_func`` is None, V must log a warning and route
+    to chunking_by_recursive_character (R) — V's only differentiator
+    is embeddings, so without them R is the closest neighbour."""
+    body = "Para A.\n\nPara B for fallback test.\n\nPara C."
+
+    lightrag_logger = logging.getLogger("lightrag")
+    handler = _ListHandler()
+    handler.setLevel(logging.WARNING)
+    lightrag_logger.addHandler(handler)
+    try:
+
+        async def _run():
+            return await chunking_by_semantic_vector(
+                _tok(),
+                body,
+                chunk_token_size=20,
+                embedding_func=None,
+            )
+
+        chunks = asyncio.run(_run())
+    finally:
+        lightrag_logger.removeHandler(handler)
+
+    assert len(chunks) >= 1
+    assert any(
+        "embedding_func is None" in rec.getMessage()
+        for rec in handler.records
+        if rec.levelno == logging.WARNING
+    )
+
+
+@pytest.mark.offline
+def test_v_chunker_empty_input_returns_empty_list():
+    async def _run():
+        return await chunking_by_semantic_vector(_tok(), "")
+
+    assert asyncio.run(_run()) == []
+
+
+@pytest.mark.offline
+def test_v_chunker_spans_repeated_sentences_and_number_of_chunks():
+    body = (
+        "Repeat sentence.\n"
+        "Repeat sentence. "
+        "A distant topic appears. "
+        "Another distant topic appears."
+    )
+
+    async def _run():
+        return await chunking_by_semantic_vector(
+            _tok(),
+            body,
+            chunk_token_size=400,
+            embedding_func=_make_deterministic_embedding(),
+            number_of_chunks=2,
+        )
+
+    chunks = asyncio.run(_run())
+
+    assert chunks
+    for chunk in chunks:
+        span = chunk["_source_span"]
+        assert body[span["start"] : span["end"]] == chunk["content"]
+
+
+@pytest.mark.offline
+def test_v_oversized_group_resplit_preserves_child_source_spans():
+    body = " ".join(f"word{i:02d}" for i in range(30)) + "."
+
+    async def _run():
+        return await chunking_by_semantic_vector(
+            _tok(),
+            body,
+            chunk_token_size=35,
+            embedding_func=_make_deterministic_embedding(),
+        )
+
+    chunks = asyncio.run(_run())
+
+    assert len(chunks) > 1
+    for chunk in chunks:
+        span = chunk["_source_span"]
+        assert body[span["start"] : span["end"]] == chunk["content"]
+        assert chunk["tokens"] <= 35
+
+
+@pytest.mark.offline
+def test_semantic_groups_mirror_matches_upstream_split_text():
+    """Drift guard: ``_semantic_groups_with_spans`` re-implements the private body
+    of ``SemanticChunker.split_text`` to keep verbatim source spans. If an upstream
+    langchain-experimental release changes that grouping (or a private member it
+    relies on), this catches it — the mirror must yield the same group count and the
+    same content modulo whitespace (it keeps the verbatim slice; upstream reflows
+    sentences with single spaces)."""
+    from langchain_core.embeddings import Embeddings
+    from langchain_experimental.text_splitter import SemanticChunker
+
+    from lightrag.chunker.semantic_vector import _semantic_groups_with_spans
+    from lightrag.constants import DEFAULT_SENTENCE_SPLIT_REGEX
+
+    class _SyncEmbeddings(Embeddings):
+        """Deterministic per-text vectors, mirroring _make_deterministic_embedding
+        but synchronous so SemanticChunker can be driven directly."""
+
+        def embed_documents(self, texts):
+            rows = []
+            for text in texts:
+                seed = abs(hash(text)) % (2**32)
+                rng = np.random.default_rng(seed=seed)
+                vec = rng.normal(size=8).astype(np.float64)
+                vec /= np.linalg.norm(vec) or 1.0
+                rows.append(vec.tolist())
+            return rows
+
+        def embed_query(self, text):
+            return self.embed_documents([text])[0]
+
+    text = (
+        "Quantum mechanics describes nature at small scales. "
+        "It contradicts classical intuition. "
+        "Bread is baked from flour. "
+        "Sourdough requires a long fermentation."
+    )
+    splitter = SemanticChunker(
+        _SyncEmbeddings(),
+        buffer_size=1,
+        breakpoint_threshold_type="percentile",
+        sentence_split_regex=DEFAULT_SENTENCE_SPLIT_REGEX,
+    )
+
+    # Same process → identical deterministic embeddings on both calls → identical
+    # distances/breakpoints, so any divergence is a real grouping drift.
+    upstream = splitter.split_text(text)
+    mirror = _semantic_groups_with_spans(splitter, text)
+
+    assert len(mirror) == len(upstream)
+    for (grp_text, start, end), up in zip(mirror, upstream):
+        assert grp_text == text[start:end]  # mirror keeps the verbatim source slice
+        assert "".join(grp_text.split()) == "".join(up.split())

tests/chunker/test_chunking_raw_lightrag_parity.py

@@ -0,0 +1,715 @@
+"""F-chunking parity between raw and lightrag formats.
+
+After the F-chunking unification, ``apipeline_process_enqueue_documents``
+strips the ``{{LRdoc}}`` marker from lightrag-format content and feeds the
+result into the same ``chunking_func`` used by raw documents.  These tests
+guard the contract end-to-end:
+
+* T1: identical input text produces identical chunking inputs whether it
+  arrives as raw or as an already-parsed lightrag full_docs row pulled
+  back through the parse queue on resume (ReuseParser) — the only
+  production path for lightrag rows now that the ``docs_format="lightrag"``
+  enqueue entrypoint is removed (the in-run parse → chunk path is T6).
+* T2: after a pending_parse document is parsed by the native engine,
+  ``full_docs.content`` carries the *full* merged text with the
+  ``{{LRdoc}}`` marker (written by ``_persist_parsed_full_docs``), while
+  ``doc_status`` reports the bare body length / summary (no marker
+  leakage).
+* T3: ``extraction_meta["parse_format"]`` (surfaced via
+  ``doc_status.metadata``) is ``"lightrag"`` for natively-parsed docs —
+  previously a structured-parse fallback always tagged ``raw`` and
+  silently mislabelled the persisted record.
+* T4: a raw document whose body coincidentally *looks* like structured
+  JSONL is still tokenised as plain text — guards against re-introducing
+  dropped structured-format detection in the raw path.
+* T5: ``process_options`` selecting R/V/P logs the deferred-strategy
+  warning and falls back to fixed-token chunking.
+* T6: a ``pending_parse`` document that resolves to lightrag at parse
+  time ends up with a real ``content_summary`` after PROCESSED — the
+  ANALYZING transition refreshes the summary from the parsed body so
+  pending-parse rows no longer carry the empty enqueue-time placeholder
+  through to the user-facing list APIs.
+* T7: a raw document whose body *literally* starts with ``{{LRdoc}}``
+  is chunked verbatim — guards against accidental re-introduction of an
+  unconditional ``strip_lightrag_doc_prefix`` at the chunking boundary
+  (which would silently drop the user's first 9 characters).
+"""
+
+import asyncio
+import json
+import logging
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from lightrag import LightRAG, ROLES, RoleLLMConfig
+from lightrag.constants import (
+    FULL_DOCS_FORMAT_LIGHTRAG,
+    FULL_DOCS_FORMAT_PENDING_PARSE,
+    LIGHTRAG_DOC_CONTENT_PREFIX,
+)
+from lightrag.utils import (
+    EmbeddingFunc,
+    Tokenizer,
+    compute_mdhash_id,
+    get_content_summary,
+)
+from lightrag.utils_pipeline import make_lightrag_doc_content
+
+
+# ---------------------------------------------------------------------------
+# Shared fixtures (mirrors the harness used by test_pipeline_release_closure)
+# ---------------------------------------------------------------------------
+
+
+class _SimpleTokenizerImpl:
+    """Char-level tokenizer so 1 char ≈ 1 token; keeps assertions readable."""
+
+    def encode(self, content: str) -> list[int]:
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens: list[int]) -> str:
+        return "".join(chr(t) for t in tokens)
+
+
+async def _mock_embedding(texts: list[str]) -> np.ndarray:
+    return np.random.rand(len(texts), 32)
+
+
+async def _mock_llm(prompt, **kwargs):
+    return '{"name":"x","summary":"s","detail_description":"d"}'
+
+
+_ROLE_FIELD_SUFFIXES = (
+    ("_llm_model_func", "func"),
+    ("_llm_model_kwargs", "kwargs"),
+    ("_llm_model_max_async", "max_async"),
+    ("_llm_timeout", "timeout"),
+)
+
+
+def _new_rag(tmp_path: Path, **kwargs) -> LightRAG:
+    role_configs: dict[str, RoleLLMConfig] = {}
+    for spec in ROLES:
+        bucket = {}
+        for suffix, target in _ROLE_FIELD_SUFFIXES:
+            key = f"{spec.name}{suffix}"
+            if key in kwargs:
+                bucket[target] = kwargs.pop(key)
+        if bucket:
+            role_configs[spec.name] = RoleLLMConfig(**bucket)
+    if role_configs:
+        kwargs["role_llm_configs"] = role_configs
+
+    return LightRAG(
+        working_dir=str(tmp_path),
+        workspace=f"chunking-parity-{tmp_path.name}",
+        llm_model_func=_mock_llm,
+        embedding_func=EmbeddingFunc(
+            embedding_dim=32,
+            max_token_size=4096,
+            func=_mock_embedding,
+        ),
+        tokenizer=Tokenizer("mock-tokenizer", _SimpleTokenizerImpl()),
+        **kwargs,
+    )
+
+
+def _attach_chunking_spy(rag: LightRAG) -> dict:
+    """Replace ``rag.chunking_func`` with a recording wrapper.
+
+    Returns a dict whose ``input`` key receives the second positional arg
+    (the content string) at every chunking call.  The original chunker
+    runs normally so the pipeline reaches PROCESSED.
+    """
+    captured: dict = {"input": None, "calls": 0}
+    real = rag.chunking_func
+
+    def _spy(tokenizer, content, *args, **kwargs):
+        captured["input"] = content
+        captured["calls"] += 1
+        return real(tokenizer, content, *args, **kwargs)
+
+    rag.chunking_func = _spy
+    return captured
+
+
+# ---------------------------------------------------------------------------
+# T1 — parity: raw vs lightrag produce identical chunking input
+# ---------------------------------------------------------------------------
+
+
+async def _seed_lightrag_row(rag: LightRAG, doc_id: str, *, body: str, file_path: str):
+    """Seed an already-parsed lightrag full_docs row + a PENDING doc_status.
+
+    Mirrors how lightrag-format rows exist in production: written by the
+    parsers (``_persist_parsed_full_docs``) and pulled back through the
+    parse queue (→ ReuseParser) only on resume/retry.
+    """
+    from datetime import datetime, timezone
+
+    await rag.full_docs.upsert(
+        {
+            doc_id: {
+                "content": make_lightrag_doc_content(body),
+                "file_path": file_path,
+                "parse_format": FULL_DOCS_FORMAT_LIGHTRAG,
+            }
+        }
+    )
+    now = datetime.now(timezone.utc).isoformat()
+    await rag.doc_status.upsert(
+        {
+            doc_id: {
+                "status": "pending",
+                "content_summary": get_content_summary(body),
+                "content_length": len(body),
+                "file_path": file_path,
+                "created_at": now,
+                "updated_at": now,
+                "track_id": "track-lr",
+            }
+        }
+    )
+
+
+@pytest.mark.offline
+def test_chunking_input_parity_raw_vs_lightrag(tmp_path):
+    """Same body text must reach ``chunking_func`` with byte-identical input
+    whether it arrives as raw or as a resumed lightrag full_docs row."""
+
+    paragraphs = [
+        "Alpha paragraph with enough words to make it look real.",
+        "Beta paragraph extends the body so chunking has substance.",
+        "Gamma paragraph closes the document with a few more sentences.",
+    ]
+    expected_merged = "\n\n".join(paragraphs)
+
+    async def _run():
+        # ---- RAW path ----
+        rag_raw = _new_rag(tmp_path / "raw")
+        await rag_raw.initialize_storages()
+        spy_raw = _attach_chunking_spy(rag_raw)
+        try:
+            await rag_raw.apipeline_enqueue_documents(
+                expected_merged,
+                file_paths="parity_raw.txt",
+                track_id="track-raw",
+            )
+            await rag_raw.apipeline_process_enqueue_documents()
+        finally:
+            await rag_raw.finalize_storages()
+
+        # ---- LIGHTRAG path (resume: seeded parsed row → ReuseParser) ----
+        rag_lr = _new_rag(tmp_path / "lr")
+        await rag_lr.initialize_storages()
+        spy_lr = _attach_chunking_spy(rag_lr)
+        try:
+            await _seed_lightrag_row(
+                rag_lr,
+                "doc-parity-lr",
+                body=expected_merged,
+                file_path="parity.lightrag",
+            )
+            await rag_lr.apipeline_process_enqueue_documents()
+        finally:
+            await rag_lr.finalize_storages()
+
+        assert spy_raw["calls"] >= 1, "raw doc never reached chunking_func"
+        assert spy_lr["calls"] >= 1, "lightrag doc never reached chunking_func"
+        assert spy_lr["input"] == spy_raw["input"] == expected_merged, (
+            "chunking_func received different inputs for raw vs lightrag; "
+            f"raw={spy_raw['input']!r}\nlr={spy_lr['input']!r}"
+        )
+        assert not spy_lr["input"].startswith(LIGHTRAG_DOC_CONTENT_PREFIX), (
+            "{{LRdoc}} marker leaked into chunking_func input"
+        )
+
+    asyncio.run(_run())
+
+
+# ---------------------------------------------------------------------------
+# T2 — full_docs.content carries full text; doc_status reports bare body
+# ---------------------------------------------------------------------------
+
+
+def _stub_docx_blocks(monkeypatch, body_paragraphs: list[str]) -> None:
+    """Stub the docx extractor so native parse yields deterministic blocks;
+    the adapter still writes the canonical .blocks.jsonl + sidecars and
+    persists the lightrag-format full_docs row."""
+
+    def _stub_extract(file_path, drawing_context=None, **kwargs):
+        return [
+            {
+                "uuid": f"para-{i}",
+                "uuid_end": f"para-{i}",
+                "heading": "",
+                "content": para,
+                "type": "text",
+                "parent_headings": [],
+                "level": 0,
+                "table_chunk_role": "none",
+            }
+            for i, para in enumerate(body_paragraphs)
+        ]
+
+    monkeypatch.setattr(
+        "lightrag.parser.docx.parse_document.extract_docx_blocks",
+        _stub_extract,
+    )
+
+
+@pytest.mark.offline
+def test_full_docs_content_carries_full_merged_text(tmp_path, monkeypatch):
+    """After the native engine parses a pending_parse upload,
+    ``_persist_parsed_full_docs`` stores the *full* merged text with the
+    ``{{LRdoc}}`` marker while doc_status reports bare-body semantics."""
+
+    body = "x" * 5000  # single paragraph, 5000 chars
+
+    async def _run():
+        input_dir = tmp_path / "input"
+        input_dir.mkdir()
+        monkeypatch.setenv("INPUT_DIR", str(input_dir))
+        (input_dir / "big.docx").write_bytes(b"fake docx bytes")
+        _stub_docx_blocks(monkeypatch, [body])
+
+        rag = _new_rag(tmp_path / "work")
+        await rag.initialize_storages()
+        try:
+            await rag.apipeline_enqueue_documents(
+                "",
+                file_paths="big.docx",
+                docs_format=FULL_DOCS_FORMAT_PENDING_PARSE,
+                parse_engine="native",
+                track_id="track-big",
+            )
+            await rag.apipeline_process_enqueue_documents()
+
+            doc_id = compute_mdhash_id("big.docx", prefix="doc-")
+            full_doc = await rag.full_docs.get_by_id(doc_id)
+            assert full_doc is not None
+            # full_docs preserves the marker AND the full merged text.
+            assert full_doc["content"] == LIGHTRAG_DOC_CONTENT_PREFIX + body
+            assert full_doc.get("parse_format") == FULL_DOCS_FORMAT_LIGHTRAG
+            assert full_doc.get("sidecar_location"), (
+                "native parse must record the sidecar_location on the "
+                "lightrag full_docs row"
+            )
+
+            # doc_status reports body-length semantics (no marker leakage).
+            status_doc = await rag.doc_status.get_by_id(doc_id)
+            assert status_doc is not None
+            length = (
+                status_doc.get("content_length")
+                if isinstance(status_doc, dict)
+                else getattr(status_doc, "content_length", None)
+            )
+            summary = (
+                status_doc.get("content_summary")
+                if isinstance(status_doc, dict)
+                else getattr(status_doc, "content_summary", "")
+            )
+            assert length == 5000, f"content_length should match body, got {length}"
+            assert not summary.startswith(LIGHTRAG_DOC_CONTENT_PREFIX)
+        finally:
+            await rag.finalize_storages()
+
+    asyncio.run(_run())
+
+
+# ---------------------------------------------------------------------------
+# T3 — extraction_meta.parse_format reflects persisted format (regression guard)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.offline
+def test_extraction_meta_records_lightrag_parse_format(tmp_path, monkeypatch):
+    """Before the unification, a structured-parse fallback tagged
+    ``extraction_meta.parse_format = raw`` for lightrag docs, silently
+    mislabelling them in ``doc_status.metadata``.  Assert the tag now
+    reflects the persisted format end-to-end (pending_parse upload →
+    native parse → lightrag full_docs row)."""
+
+    paragraphs = ["Body paragraph for parse_format tagging test."]
+
+    async def _run():
+        input_dir = tmp_path / "input"
+        input_dir.mkdir()
+        monkeypatch.setenv("INPUT_DIR", str(input_dir))
+        (input_dir / "tag.docx").write_bytes(b"fake docx bytes")
+        _stub_docx_blocks(monkeypatch, paragraphs)
+
+        rag = _new_rag(tmp_path / "work")
+        await rag.initialize_storages()
+
+        try:
+            await rag.apipeline_enqueue_documents(
+                "",
+                file_paths="tag.docx",
+                docs_format=FULL_DOCS_FORMAT_PENDING_PARSE,
+                parse_engine="native",
+                track_id="track-tag",
+            )
+            await rag.apipeline_process_enqueue_documents()
+
+            doc_id = compute_mdhash_id("tag.docx", prefix="doc-")
+            status_doc = await rag.doc_status.get_by_id(doc_id)
+            assert status_doc is not None
+            metadata = (
+                status_doc.get("metadata")
+                if isinstance(status_doc, dict)
+                else getattr(status_doc, "metadata", None)
+            )
+            assert isinstance(metadata, dict), (
+                f"doc_status.metadata should be a dict, got {type(metadata)!r}"
+            )
+            assert metadata.get("parse_format") == FULL_DOCS_FORMAT_LIGHTRAG, (
+                f"doc_status.metadata.parse_format="
+                f"{metadata.get('parse_format')!r}; "
+                f"expected {FULL_DOCS_FORMAT_LIGHTRAG!r} so the multimodal "
+                f"sidecar merge path opens"
+            )
+        finally:
+            await rag.finalize_storages()
+
+    asyncio.run(_run())
+
+
+# ---------------------------------------------------------------------------
+# T4 — JSONL-shaped raw text is still treated as plain text
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.offline
+def test_jsonl_shaped_raw_text_chunks_as_plain_text(tmp_path):
+    """A raw document whose body coincidentally resembles structured JSONL
+    must be tokenised plainly — guarding against accidental
+    re-introduction of removed structured-format detection."""
+
+    # No trailing newline — sanitize_text_for_encoding strips trailing
+    # whitespace on raw enqueue, and that pre-chunking cleanup is unrelated
+    # to structured-format detection.
+    pseudo_jsonl = (
+        json.dumps({"type": "meta", "format_version": "1.0"})
+        + "\n"
+        + json.dumps(
+            {
+                "type": "text",
+                "chunk_id": "c0",
+                "chunk_order_index": 0,
+                "content": "fake structured line",
+            }
+        )
+    )
+
+    async def _run():
+        rag = _new_rag(tmp_path)
+        await rag.initialize_storages()
+        spy = _attach_chunking_spy(rag)
+        try:
+            await rag.apipeline_enqueue_documents(
+                pseudo_jsonl,
+                file_paths="pseudo.txt",
+                track_id="track-pseudo",
+            )
+            await rag.apipeline_process_enqueue_documents()
+        finally:
+            await rag.finalize_storages()
+
+        # The full pseudo-jsonl text reaches chunking_func; nothing parses
+        # it as JSONL and hijacks the chunks list.
+        assert spy["input"] == pseudo_jsonl
+
+    asyncio.run(_run())
+
+
+# ---------------------------------------------------------------------------
+# T5 — R/V/P process_options trigger the deferred-strategy warning
+# ---------------------------------------------------------------------------
+
+
+class _ListHandler(logging.Handler):
+    """Capture log records into an in-memory list.
+
+    The ``lightrag`` logger has ``propagate = False`` so pytest's caplog
+    fixture cannot intercept its records via the root logger; this handler
+    attaches directly to the logger we care about.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.records: list[logging.LogRecord] = []
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self.records.append(record)
+
+
+@pytest.mark.offline
+def test_explicit_R_dispatches_to_recursive_character(tmp_path, monkeypatch):
+    """``process_options=R`` must invoke
+    :func:`chunking_by_recursive_character` (the new file-chunker
+    contract) rather than the legacy ``chunking_func``.
+
+    Verifies the explicit-selector dispatch contract:
+      1. ``chunking_by_recursive_character`` runs at least once.
+      2. The legacy ``chunking_func`` is bypassed entirely.
+      3. The deprecated "R/V not yet implemented" warning no longer
+         appears (now that R has a real implementation).
+    """
+
+    pytest.importorskip("langchain_text_splitters")
+
+    import lightrag.chunker as chunker_pkg
+    from lightrag.chunker import chunking_by_recursive_character as real_r
+
+    captured = {"calls": 0}
+
+    def _r_spy(*args, **kwargs):
+        captured["calls"] += 1
+        return real_r(*args, **kwargs)
+
+    # The dispatcher does ``from lightrag.chunker import …`` inside the
+    # function body, which re-resolves the name from the package each
+    # call — patching the package attribute is enough to intercept it.
+    monkeypatch.setattr(chunker_pkg, "chunking_by_recursive_character", _r_spy)
+
+    async def _run():
+        rag = _new_rag(tmp_path)
+        await rag.initialize_storages()
+        legacy_spy = _attach_chunking_spy(rag)
+
+        lightrag_logger = logging.getLogger("lightrag")
+        list_handler = _ListHandler()
+        list_handler.setLevel(logging.WARNING)
+        lightrag_logger.addHandler(list_handler)
+        try:
+            await rag.apipeline_enqueue_documents(
+                "Body paragraph one.\n\nBody paragraph two for R dispatch test.",
+                file_paths="rs.[native-R].txt",
+                track_id="track-rs",
+                process_options="R",
+            )
+            await rag.apipeline_process_enqueue_documents()
+        finally:
+            lightrag_logger.removeHandler(list_handler)
+            await rag.finalize_storages()
+
+        assert captured["calls"] >= 1, "R must route to chunking_by_recursive_character"
+        assert legacy_spy["calls"] == 0, (
+            "explicit process_options selector must bypass legacy "
+            "chunking_func; got "
+            f"{legacy_spy['calls']} calls"
+        )
+        warning_messages = [
+            rec.getMessage()
+            for rec in list_handler.records
+            if rec.levelno == logging.WARNING
+        ]
+        assert not any(
+            "R/V strategies are not yet implemented" in msg for msg in warning_messages
+        ), (
+            "deprecated 'not yet implemented' warning must be gone now "
+            f"that R is wired up; saw: {warning_messages!r}"
+        )
+
+    asyncio.run(_run())
+
+
+@pytest.mark.offline
+def test_explicit_V_dispatches_to_semantic_vector(tmp_path, monkeypatch):
+    """``process_options=V`` must invoke
+    :func:`chunking_by_semantic_vector` and bypass the legacy
+    ``chunking_func``.  The test installs a stub embedding (the spy
+    short-circuits before the real LangChain SemanticChunker runs) so
+    the assertion is purely about dispatch routing, not chunk quality.
+    """
+
+    pytest.importorskip("langchain_experimental")
+
+    import lightrag.chunker as chunker_pkg
+
+    captured = {"calls": 0}
+
+    async def _v_spy(*args, **kwargs):
+        # Short-circuit: skip langchain SemanticChunker entirely and
+        # return one synthetic chunk.  We're only verifying that the
+        # dispatcher routed here with the right keyword args.
+        captured["calls"] += 1
+        captured["embedding_func"] = kwargs.get("embedding_func")
+        captured["chunk_token_size"] = args[2] if len(args) > 2 else None
+        return [
+            {"tokens": 5, "content": "stub", "chunk_order_index": 0},
+        ]
+
+    monkeypatch.setattr(chunker_pkg, "chunking_by_semantic_vector", _v_spy)
+
+    async def _run():
+        rag = _new_rag(tmp_path)
+        await rag.initialize_storages()
+        legacy_spy = _attach_chunking_spy(rag)
+        try:
+            await rag.apipeline_enqueue_documents(
+                "Body for V dispatch test. Sentence one. Sentence two.",
+                file_paths="vs.[native-V].txt",
+                track_id="track-vs",
+                process_options="V",
+            )
+            await rag.apipeline_process_enqueue_documents()
+        finally:
+            await rag.finalize_storages()
+
+        assert captured["calls"] >= 1, "V must route to chunking_by_semantic_vector"
+        assert captured.get("embedding_func") is rag.embedding_func, (
+            "dispatcher must hand the LightRAG embedding_func to the V chunker"
+        )
+        assert legacy_spy["calls"] == 0, (
+            "explicit process_options selector must bypass legacy chunking_func"
+        )
+
+    asyncio.run(_run())
+
+
+# ---------------------------------------------------------------------------
+# T6 — pending_parse → lightrag summary is populated after PROCESSED
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.offline
+def test_pending_parse_lightrag_summary_populated_after_processed(
+    tmp_path, monkeypatch
+):
+    """A document enqueued as ``pending_parse`` has empty content at
+    enqueue time, so ``content_summary`` starts empty.  After
+    ``parse_native`` produces ``.blocks.jsonl`` and the state machine
+    moves through ANALYZING → PROCESSING → PROCESSED, the summary must
+    reflect the parsed body — not the enqueue-time placeholder."""
+
+    body_paragraphs = [
+        "Pending-parse summary regression body paragraph one.",
+        "Body paragraph two carries enough text for a meaningful preview.",
+        "Body paragraph three closes the document.",
+    ]
+
+    async def _run():
+        input_dir = tmp_path / "input"
+        input_dir.mkdir()
+        monkeypatch.setenv("INPUT_DIR", str(input_dir))
+
+        source_path = input_dir / "summary.docx"
+        source_path.write_bytes(b"fake docx bytes")
+        _stub_docx_blocks(monkeypatch, body_paragraphs)
+
+        rag = _new_rag(tmp_path / "work")
+        await rag.initialize_storages()
+        try:
+            await rag.apipeline_enqueue_documents(
+                "",
+                file_paths="summary.docx",
+                docs_format=FULL_DOCS_FORMAT_PENDING_PARSE,
+                parse_engine="native",
+                track_id="track-summary",
+            )
+
+            doc_id = compute_mdhash_id("summary.docx", prefix="doc-")
+
+            pending = await rag.doc_status.get_by_id(doc_id)
+            assert pending is not None
+            pending_summary = (
+                pending.get("content_summary")
+                if isinstance(pending, dict)
+                else getattr(pending, "content_summary", "")
+            )
+            # At enqueue time pending_parse content is "" so summary is empty.
+            assert pending_summary == "", (
+                f"pending_parse should start with empty summary, got "
+                f"{pending_summary!r}"
+            )
+
+            await rag.apipeline_process_enqueue_documents()
+
+            final = await rag.doc_status.get_by_id(doc_id)
+            assert final is not None
+            final_summary = (
+                final.get("content_summary")
+                if isinstance(final, dict)
+                else getattr(final, "content_summary", "")
+            )
+            final_length = (
+                final.get("content_length")
+                if isinstance(final, dict)
+                else getattr(final, "content_length", 0)
+            )
+
+            assert final_summary, (
+                "content_summary still empty after PROCESSED; ANALYZING "
+                "refresh did not propagate"
+            )
+            assert not final_summary.startswith(LIGHTRAG_DOC_CONTENT_PREFIX), (
+                f"{{LRdoc}} marker leaked into doc_status summary: {final_summary!r}"
+            )
+            # The parser stub produces these paragraphs verbatim; the
+            # blocks.jsonl writer joins them with a blank line, so the
+            # summary must be a prefix of that merged text.
+            merged_text = "\n\n".join(body_paragraphs)
+            assert final_summary == get_content_summary(merged_text), (
+                f"summary should match get_content_summary(merged_text); "
+                f"got {final_summary!r} vs "
+                f"{get_content_summary(merged_text)!r}"
+            )
+            assert final_length == len(merged_text), (
+                f"content_length should equal len(merged_text)={len(merged_text)}, "
+                f"got {final_length}"
+            )
+        finally:
+            await rag.finalize_storages()
+
+    asyncio.run(_run())
+
+
+# ---------------------------------------------------------------------------
+# T7 — raw text starting with {{LRdoc}} must not be stripped at chunking
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.offline
+def test_raw_text_starting_with_marker_chunked_verbatim(tmp_path):
+    """A raw document whose body literally begins with ``{{LRdoc}}`` is a
+    legitimate user input — the chunking branch must not strip those 9
+    characters.  ``strip_lightrag_doc_prefix`` is a lightrag-only contract
+    enforced by ``parse_native``; raw paths return ``content_data["content"]``
+    verbatim, so chunking must hand the body to ``chunking_func`` unchanged."""
+
+    body_with_marker = LIGHTRAG_DOC_CONTENT_PREFIX + (
+        "literal-marker-prefix raw document body that should survive "
+        "the chunking boundary intact."
+    )
+
+    async def _run():
+        rag = _new_rag(tmp_path)
+        await rag.initialize_storages()
+        spy = _attach_chunking_spy(rag)
+        try:
+            await rag.apipeline_enqueue_documents(
+                body_with_marker,
+                file_paths="marker_raw.txt",
+                track_id="track-marker",
+            )
+            await rag.apipeline_process_enqueue_documents()
+        finally:
+            await rag.finalize_storages()
+
+        assert spy["calls"] >= 1, "raw doc never reached chunking_func"
+        # The full body — including the literal {{LRdoc}} prefix — must
+        # reach chunking_func; nothing in the chunking branch should
+        # treat the marker as a stripping signal for raw content.
+        assert spy["input"] == body_with_marker, (
+            "chunking_func received corrupted input: "
+            f"got {spy['input']!r}, expected {body_with_marker!r}"
+        )
+        assert spy["input"].startswith(LIGHTRAG_DOC_CONTENT_PREFIX), (
+            "literal marker prefix lost at chunking boundary"
+        )
+
+    asyncio.run(_run())

tests/chunker/test_overlap_validation.py

@@ -0,0 +1,113 @@
+"""
+Test for overlap_tokens validation to prevent infinite loop.
+
+This test validates the fix for the bug where overlap_tokens >= max_tokens
+causes an infinite loop in the chunking function.
+"""
+
+from lightrag.rerank import chunk_documents_for_rerank
+
+
+class TestOverlapValidation:
+    """Test suite for overlap_tokens validation"""
+
+    def test_overlap_greater_than_max_tokens(self):
+        """Test that overlap_tokens > max_tokens is clamped and doesn't hang"""
+        documents = [" ".join([f"word{i}" for i in range(100)])]
+
+        # This should clamp overlap_tokens to 29 (max_tokens - 1)
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=30, overlap_tokens=32
+        )
+
+        # Should complete without hanging
+        assert len(chunked_docs) > 0
+        assert all(idx == 0 for idx in doc_indices)
+
+    def test_overlap_equal_to_max_tokens(self):
+        """Test that overlap_tokens == max_tokens is clamped and doesn't hang"""
+        documents = [" ".join([f"word{i}" for i in range(100)])]
+
+        # This should clamp overlap_tokens to 29 (max_tokens - 1)
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=30, overlap_tokens=30
+        )
+
+        # Should complete without hanging
+        assert len(chunked_docs) > 0
+        assert all(idx == 0 for idx in doc_indices)
+
+    def test_overlap_slightly_less_than_max_tokens(self):
+        """Test that overlap_tokens < max_tokens works normally"""
+        documents = [" ".join([f"word{i}" for i in range(100)])]
+
+        # This should work without clamping
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=30, overlap_tokens=29
+        )
+
+        # Should complete successfully
+        assert len(chunked_docs) > 0
+        assert all(idx == 0 for idx in doc_indices)
+
+    def test_small_max_tokens_with_large_overlap(self):
+        """Test edge case with very small max_tokens"""
+        documents = [" ".join([f"word{i}" for i in range(50)])]
+
+        # max_tokens=5, overlap_tokens=10 should clamp to 4
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=5, overlap_tokens=10
+        )
+
+        # Should complete without hanging
+        assert len(chunked_docs) > 0
+        assert all(idx == 0 for idx in doc_indices)
+
+    def test_multiple_documents_with_invalid_overlap(self):
+        """Test multiple documents with overlap_tokens >= max_tokens"""
+        documents = [
+            " ".join([f"word{i}" for i in range(50)]),
+            "short document",
+            " ".join([f"word{i}" for i in range(75)]),
+        ]
+
+        # overlap_tokens > max_tokens
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=25, overlap_tokens=30
+        )
+
+        # Should complete successfully and chunk the long documents
+        assert len(chunked_docs) >= len(documents)
+        # Short document should not be chunked
+        assert "short document" in chunked_docs
+
+    def test_normal_operation_unaffected(self):
+        """Test that normal cases continue to work correctly"""
+        documents = [
+            " ".join([f"word{i}" for i in range(100)]),
+            "short doc",
+        ]
+
+        # Normal case: overlap_tokens (10) < max_tokens (50)
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=50, overlap_tokens=10
+        )
+
+        # Long document should be chunked, short one should not
+        assert len(chunked_docs) > 2  # At least 3 chunks (2 from long doc + 1 short)
+        assert "short doc" in chunked_docs
+        # Verify doc_indices maps correctly
+        assert doc_indices[-1] == 1  # Last chunk is from second document
+
+    def test_edge_case_max_tokens_one(self):
+        """Test edge case where max_tokens=1"""
+        documents = [" ".join([f"word{i}" for i in range(20)])]
+
+        # max_tokens=1, overlap_tokens=5 should clamp to 0
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=1, overlap_tokens=5
+        )
+
+        # Should complete without hanging
+        assert len(chunked_docs) > 0
+        assert all(idx == 0 for idx in doc_indices)

tests/chunker/test_paragraph_semantic_drop_references.py

@@ -0,0 +1,310 @@
+"""Tests for the paragraph-semantic ``drop_references`` option (chunking=P).
+
+A reference block is dropped only when it BOTH sits within the last
+``references_tail_n`` content blocks AND its heading matches a reference prefix.
+The switch is the only knob that flows through ``chunk_options``; the tail
+window and heading prefixes are read live from env at chunk time (verified here
+by mutating env between calls, proving they are NOT snapshotted).
+
+Assertions check the *content* of the emitted chunks (each block carries a
+distinctive marker) rather than heading names, because LevelMerge may merge
+small blocks and rewrite the surviving heading.
+"""
+
+import json
+import logging
+
+import pytest
+
+from lightrag.chunker.paragraph_semantic import (
+    _format_dropped_headings,
+    _is_reference_heading,
+    chunking_by_paragraph_semantic,
+)
+from lightrag.constants import DEFAULT_P_REFERENCES_HEADINGS
+from lightrag.utils import Tokenizer, TokenizerInterface, logger as _lightrag_logger
+
+
+class _CharTokenizer(TokenizerInterface):
+    """1:1 character-to-token mapping — keeps math obvious in assertions."""
+
+    def encode(self, content: str):
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens):
+        return "".join(chr(t) for t in tokens)
+
+
+def _make_tokenizer() -> Tokenizer:
+    return Tokenizer(model_name="char", tokenizer=_CharTokenizer())
+
+
+def _row(heading: str, content: str, *, level: int = 1) -> dict:
+    return {
+        "type": "content",
+        "blockid": heading or "blk",
+        "format": "plain_text",
+        "content": content,
+        "heading": heading,
+        "parent_headings": [],
+        "level": level,
+        "session_type": "body",
+        "table_slice": "none",
+        "positions": [],
+    }
+
+
+def _write_blocks_jsonl(tmp_path, rows: list[dict]) -> str:
+    path = tmp_path / "doc.blocks.jsonl"
+    path.write_text(
+        "\n".join(json.dumps(row, ensure_ascii=False) for row in rows),
+        encoding="utf-8",
+    )
+    return str(path)
+
+
+def _all_content(chunks: list[dict]) -> str:
+    return "\n".join(c["content"] for c in chunks)
+
+
+# --------------------------------------------------------------------------- #
+# _is_reference_heading (word boundary vs. plain CJK prefix)
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    "heading,expected",
+    [
+        ("References", True),
+        ("references", True),  # case-insensitive
+        ("REFERENCES", True),
+        ("Bibliography", True),
+        ("References [1-50]", True),  # word boundary: next char is non-alnum
+        ("参考文献", True),
+        ("参考文献列表", True),  # CJK: plain prefix, no word boundary
+        ("Referenced work", False),  # ASCII word boundary excludes this
+        ("Related Work", False),
+        ("", False),
+    ],
+)
+def test_is_reference_heading(heading, expected):
+    assert _is_reference_heading(heading, DEFAULT_P_REFERENCES_HEADINGS) is expected
+
+
+# --------------------------------------------------------------------------- #
+# _format_dropped_headings (length-bounded log rendering)
+# --------------------------------------------------------------------------- #
+
+
+def test_format_dropped_headings_short_list():
+    assert _format_dropped_headings(["References"]) == "'References'"
+    assert (
+        _format_dropped_headings(["References", "Bibliography"])
+        == "'References', 'Bibliography'"
+    )
+
+
+def test_format_dropped_headings_truncates_long_heading():
+    out = _format_dropped_headings(["X" * 200], max_each=60)
+    # 60 kept chars + an ellipsis, all within one repr'd string.
+    assert "X" * 60 + "…" in out
+    assert "X" * 61 not in out
+
+
+def test_format_dropped_headings_caps_item_count():
+    out = _format_dropped_headings([f"H{i}" for i in range(12)], max_items=5)
+    assert out.count("'H") == 5  # only 5 listed
+    assert "(+7 more)" in out
+
+
+# --------------------------------------------------------------------------- #
+# Filtering behaviour (assert on content markers, robust to LevelMerge)
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.offline
+def test_drops_trailing_reference_block(tmp_path):
+    tokenizer = _make_tokenizer()
+    rows = [
+        _row("Introduction", "INTRO_MARKER intro body"),
+        _row("Method", "METHOD_MARKER method body"),
+        _row("References", "REF_MARKER [1] Foo. [2] Bar."),
+    ]
+    blocks_path = _write_blocks_jsonl(tmp_path, rows)
+
+    body = _all_content(
+        chunking_by_paragraph_semantic(
+            tokenizer, "", 2000, blocks_path=blocks_path, drop_references=True
+        )
+    )
+    assert "REF_MARKER" not in body
+    assert "INTRO_MARKER" in body and "METHOD_MARKER" in body
+
+
+@pytest.mark.offline
+def test_default_keeps_reference_block(tmp_path):
+    tokenizer = _make_tokenizer()
+    rows = [
+        _row("Introduction", "INTRO_MARKER intro body"),
+        _row("References", "REF_MARKER [1] Foo."),
+    ]
+    blocks_path = _write_blocks_jsonl(tmp_path, rows)
+
+    # drop_references defaults to False — references survive.
+    body = _all_content(
+        chunking_by_paragraph_semantic(tokenizer, "", 2000, blocks_path=blocks_path)
+    )
+    assert "REF_MARKER" in body
+
+
+@pytest.mark.offline
+def test_mid_document_reference_section_outside_window_kept(tmp_path):
+    tokenizer = _make_tokenizer()
+    # A "References" heading mid-document, with two non-reference blocks after
+    # it, so it falls outside the last-2 window and must NOT be dropped.
+    rows = [
+        _row("Introduction", "INTRO_MARKER intro body"),
+        _row("References", "REF_MARKER discusses references"),
+        _row("Method", "METHOD_MARKER method body"),
+        _row("Results", "RESULTS_MARKER results body"),
+    ]
+    blocks_path = _write_blocks_jsonl(tmp_path, rows)
+
+    body = _all_content(
+        chunking_by_paragraph_semantic(
+            tokenizer,
+            "",
+            2000,
+            blocks_path=blocks_path,
+            drop_references=True,
+            references_tail_n=2,
+        )
+    )
+    assert "REF_MARKER" in body
+
+
+@pytest.mark.offline
+def test_word_boundary_keeps_referenced(tmp_path):
+    tokenizer = _make_tokenizer()
+    rows = [
+        _row("Introduction", "INTRO_MARKER intro body"),
+        _row("参考文献列表", "CJKREF_MARKER [1] Foo."),
+        _row("Referenced datasets", "DATASET_MARKER dataset details"),
+    ]
+    blocks_path = _write_blocks_jsonl(tmp_path, rows)
+
+    body = _all_content(
+        chunking_by_paragraph_semantic(
+            tokenizer, "", 2000, blocks_path=blocks_path, drop_references=True
+        )
+    )
+    # CJK prefix block dropped; "Referenced ..." kept (ASCII word boundary).
+    assert "CJKREF_MARKER" not in body
+    assert "DATASET_MARKER" in body
+
+
+@pytest.mark.offline
+def test_all_blocks_match_keeps_to_avoid_empty(tmp_path):
+    tokenizer = _make_tokenizer()
+    rows = [
+        _row("References", "REF1_MARKER [1] Foo."),
+        _row("Bibliography", "REF2_MARKER [2] Bar."),
+    ]
+    blocks_path = _write_blocks_jsonl(tmp_path, rows)
+
+    # Both blocks match within the last-2 window; dropping all would leave no
+    # content, so the chunker keeps them and warns instead.
+    chunks = chunking_by_paragraph_semantic(
+        tokenizer, "", 2000, blocks_path=blocks_path, drop_references=True
+    )
+    body = _all_content(chunks)
+    assert chunks  # not an empty document
+    assert "REF1_MARKER" in body and "REF2_MARKER" in body
+
+
+# --------------------------------------------------------------------------- #
+# Detection knobs are read LIVE from env (not snapshotted)
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.offline
+def test_tail_n_read_live_from_env(tmp_path, monkeypatch):
+    tokenizer = _make_tokenizer()
+    rows = [
+        _row("Introduction", "INTRO_MARKER intro body"),
+        _row("References", "REF_MARKER [1] Foo."),
+        _row("Appendix", "APPENDIX_MARKER appendix body"),
+    ]
+    blocks_path = _write_blocks_jsonl(tmp_path, rows)
+
+    # tail_n=1: only the last block ("Appendix") is in the window → References
+    # is outside it and survives.
+    monkeypatch.setenv("CHUNK_P_REFERENCES_TAIL_N", "1")
+    body = _all_content(
+        chunking_by_paragraph_semantic(
+            tokenizer, "", 2000, blocks_path=blocks_path, drop_references=True
+        )
+    )
+    assert "REF_MARKER" in body
+
+    # Widen the window to 2 → References now falls inside and is dropped. Same
+    # call, only the env changed: proves the knob is read live, not snapshotted.
+    monkeypatch.setenv("CHUNK_P_REFERENCES_TAIL_N", "2")
+    body = _all_content(
+        chunking_by_paragraph_semantic(
+            tokenizer, "", 2000, blocks_path=blocks_path, drop_references=True
+        )
+    )
+    assert "REF_MARKER" not in body
+
+
+@pytest.mark.offline
+def test_drop_references_emits_info_log(tmp_path, caplog):
+    tokenizer = _make_tokenizer()
+    rows = [
+        _row("Introduction", "INTRO_MARKER intro body"),
+        _row("References", "REF_MARKER [1] Foo."),
+    ]
+    blocks_path = _write_blocks_jsonl(tmp_path, rows)
+
+    # The lightrag logger sets propagate=False, so caplog can't see it by
+    # default — enable propagation for the duration of the call.
+    original_propagate = _lightrag_logger.propagate
+    _lightrag_logger.propagate = True
+    try:
+        with caplog.at_level(logging.INFO, logger=_lightrag_logger.name):
+            chunking_by_paragraph_semantic(
+                tokenizer,
+                "",
+                2000,
+                blocks_path=blocks_path,
+                drop_references=True,
+                doc_id="doc-xyz",
+            )
+    finally:
+        _lightrag_logger.propagate = original_propagate
+
+    info_msgs = [r.getMessage() for r in caplog.records if r.levelno == logging.INFO]
+    # Log names the dropped heading and attributes it to the doc_id.
+    assert any("References" in m and "doc_id: doc-xyz" in m for m in info_msgs), (
+        info_msgs
+    )
+
+
+@pytest.mark.offline
+def test_headings_read_live_from_env(tmp_path, monkeypatch):
+    tokenizer = _make_tokenizer()
+    rows = [
+        _row("Introduction", "INTRO_MARKER intro body"),
+        _row("Literature", "LIT_MARKER [1] Foo. [2] Bar."),
+    ]
+    blocks_path = _write_blocks_jsonl(tmp_path, rows)
+
+    # Custom prefix list makes "Literature" a reference heading.
+    monkeypatch.setenv("CHUNK_P_REFERENCES_HEADINGS", "Literature|参考文献")
+    body = _all_content(
+        chunking_by_paragraph_semantic(
+            tokenizer, "", 2000, blocks_path=blocks_path, drop_references=True
+        )
+    )
+    assert "LIT_MARKER" not in body

tests/chunker/test_paragraph_semantic_merge_and_fallback.py

@@ -0,0 +1,766 @@
+"""Regression tests for paragraph-semantic LevelMerge merging and the top-level R fallback."""
+
+import pytest
+
+from lightrag.chunker.paragraph_semantic import (
+    _glue_heading_only_blocks,
+    _is_heading_only,
+    _merge_small_blocks,
+    chunking_by_paragraph_semantic,
+)
+from lightrag.utils import Tokenizer, TokenizerInterface
+
+
+class _CharTokenizer(TokenizerInterface):
+    """1:1 character-to-token mapping — keeps math obvious in assertions."""
+
+    def encode(self, content: str):
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens):
+        return "".join(chr(t) for t in tokens)
+
+
+def _make_tokenizer() -> Tokenizer:
+    return Tokenizer(model_name="char", tokenizer=_CharTokenizer())
+
+
+def _make_block(text: str, *, tokenizer: Tokenizer, level: int = 1) -> dict:
+    return {
+        "heading": "H",
+        "parent_headings": [],
+        "level": level,
+        "paragraphs": [{"text": text, "is_table": False}],
+        "content": text,
+        "tokens": len(tokenizer.encode(text)),
+        "table_chunk_role": "none",
+    }
+
+
+def _pblock(
+    text: str,
+    *,
+    heading: str,
+    parent_headings: list[str],
+    level: int,
+    tokenizer: Tokenizer,
+    role: str = "none",
+) -> dict:
+    """Block carrying an explicit heading path — for parent-path gate tests."""
+    return {
+        "heading": heading,
+        "parent_headings": list(parent_headings),
+        "level": level,
+        "paragraphs": [{"text": text, "is_table": False}],
+        "content": text,
+        "tokens": len(tokenizer.encode(text)),
+        "table_chunk_role": role,
+    }
+
+
+@pytest.mark.offline
+def test_tail_absorption_rejects_when_separator_pushes_over_cap():
+    # Tail absorption joins blocks with ``"\n\n"`` but the original
+    # predicate only summed per-block tokens. With cur=99 and tail=1
+    # the raw sum equals target_max=100, but the actual joined
+    # ``"x"*99 + "\n\n" + "y"*1`` measures 102 tokens — the absorbed
+    # block silently overflowed before the fix re-measured the joined
+    # content.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _make_block("x" * 99, tokenizer=tokenizer),
+        _make_block("y" * 1, tokenizer=tokenizer),
+    ]
+
+    merged = _merge_small_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=100,
+        target_ideal=80,
+        small_tail_threshold=12,
+    )
+
+    assert all(b["tokens"] <= 100 for b in merged), [b["tokens"] for b in merged]
+
+
+@pytest.mark.offline
+def test_tail_absorption_still_fires_when_joined_size_fits():
+    # Sanity check: when the joined content (including separators)
+    # genuinely fits target_max, absorption still happens. cur=80 +
+    # "\n\n" (2 tokens) + tail=1 = 83 ≤ 100.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _make_block("x" * 80, tokenizer=tokenizer),
+        _make_block("y" * 1, tokenizer=tokenizer),
+    ]
+
+    merged = _merge_small_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=100,
+        target_ideal=80,
+        small_tail_threshold=12,
+    )
+
+    assert len(merged) == 1
+    assert merged[0]["tokens"] == 83
+    assert merged[0]["content"] == "x" * 80 + "\n\n" + "y" * 1
+
+
+@pytest.mark.offline
+def test_phase_a_merges_same_parent_path_siblings():
+    # True siblings under one parent (identical parent_headings) merge.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _pblock(
+            "a" * 40,
+            heading="2.4.1",
+            parent_headings=["2", "2.4"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+        _pblock(
+            "b" * 40,
+            heading="2.4.2",
+            parent_headings=["2", "2.4"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+    ]
+    merged = _merge_small_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=200,
+        target_ideal=150,
+        small_tail_threshold=12,
+    )
+    assert len(merged) == 1
+
+
+@pytest.mark.offline
+def test_phase_a_keeps_different_parent_path_siblings_separate():
+    # Same level but different parents (2.4.x vs 2.5.x) must NOT merge —
+    # the anti-cross-topic-pollution guarantee (§9.1 #4).
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _pblock(
+            "a" * 40,
+            heading="2.4.1",
+            parent_headings=["2", "2.4"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+        _pblock(
+            "b" * 40,
+            heading="2.5.1",
+            parent_headings=["2", "2.5"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+    ]
+    merged = _merge_small_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=200,
+        target_ideal=150,
+        small_tail_threshold=12,
+    )
+    assert len(merged) == 2
+
+
+@pytest.mark.offline
+def test_phase_b_shallow_absorbs_descendant_deeper():
+    # Cross-level absorption is allowed when the deep block is nested under the
+    # shallow one: 2.4 (parents [2]) absorbs its child 2.4.1 (parents [2, 2.4]).
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _pblock(
+            "a" * 40,
+            heading="2.4",
+            parent_headings=["2"],
+            level=2,
+            tokenizer=tokenizer,
+        ),
+        _pblock(
+            "b" * 40,
+            heading="2.4.1",
+            parent_headings=["2", "2.4"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+    ]
+    merged = _merge_small_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=200,
+        target_ideal=150,
+        small_tail_threshold=12,
+    )
+    assert len(merged) == 1
+
+
+@pytest.mark.offline
+def test_phase_b_refuses_nondescendant_deeper():
+    # 2.4 must NOT absorb a deeper block from a different branch (2.5.1) even
+    # though it is shallower — that would be cross-topic pollution.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _pblock(
+            "a" * 40,
+            heading="2.4",
+            parent_headings=["2"],
+            level=2,
+            tokenizer=tokenizer,
+        ),
+        _pblock(
+            "b" * 40,
+            heading="2.5.1",
+            parent_headings=["2", "2.5"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+    ]
+    merged = _merge_small_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=200,
+        target_ideal=150,
+        small_tail_threshold=12,
+    )
+    assert len(merged) == 2
+
+
+@pytest.mark.offline
+def test_tail_absorption_stops_at_divergent_parent_path():
+    # A saturated block absorbs the same-parent sliver that follows but stops
+    # the run at the first block whose parent path diverges.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _pblock(
+            "a" * 160,
+            heading="2.4.1",
+            parent_headings=["2", "2.4"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+        _pblock(
+            "b" * 5,
+            heading="2.4.2",
+            parent_headings=["2", "2.4"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+        _pblock(
+            "c" * 5,
+            heading="2.5.1",
+            parent_headings=["2", "2.5"],
+            level=3,
+            tokenizer=tokenizer,
+        ),
+    ]
+    merged = _merge_small_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=200,
+        target_ideal=150,
+        small_tail_threshold=50,
+    )
+    assert len(merged) == 2
+    assert "b" * 5 in merged[0]["content"]  # same-parent sliver absorbed
+    assert merged[1]["content"] == "c" * 5  # divergent-parent block untouched
+
+
+@pytest.mark.offline
+def test_paragraph_semantic_fallback_passes_configured_recursive_overlap(monkeypatch):
+    # When ``blocks_path`` is missing, paragraph-semantic chunking
+    # delegates to ``chunking_by_recursive_character``. P now permits
+    # overlap for long text under one JSONL row, so the fallback must
+    # pass through the configured overlap rather than forcing zero.
+    captured: dict[str, object] = {}
+
+    def fake_chunker(
+        tokenizer,
+        content,
+        chunk_token_size: int = 1200,
+        *,
+        chunk_overlap_token_size: int = 100,
+        separators=None,
+    ):
+        captured["chunk_overlap_token_size"] = chunk_overlap_token_size
+        captured["chunk_token_size"] = chunk_token_size
+        return [
+            {
+                "tokens": len(tokenizer.encode(content)),
+                "content": content,
+                "chunk_order_index": 0,
+            }
+        ]
+
+    import lightrag.chunker.recursive_character as rc_mod
+
+    monkeypatch.setattr(rc_mod, "chunking_by_recursive_character", fake_chunker)
+
+    tokenizer = _make_tokenizer()
+    chunking_by_paragraph_semantic(
+        tokenizer,
+        "fallback corpus",
+        chunk_token_size=500,
+        blocks_path=None,
+        chunk_overlap_token_size=37,
+    )
+
+    assert captured.get("chunk_overlap_token_size") == 37, (
+        "P→R fallback must pass the configured chunk_overlap_token_size"
+    )
+    assert captured.get("chunk_token_size") == 500
+
+
+# ---------------------------------------------------------------------------
+# HeadingGlue — body-less heading glue (forward into child / backward into prev).
+# ---------------------------------------------------------------------------
+
+
+def _hblock(
+    content: str,
+    *,
+    heading: str,
+    level: int,
+    tokenizer: Tokenizer,
+    table_chunk_role: str = "none",
+) -> dict:
+    """Build a block whose ``content`` keeps the markdown heading line(s).
+
+    Unlike ``_make_block`` (heading-less ``content``), heading-only detection
+    needs the ``#``-prefixed heading line preserved verbatim in ``content``.
+    """
+    return {
+        "heading": heading,
+        "parent_headings": [],
+        "level": level,
+        "paragraphs": [
+            {"text": line, "is_table": False}
+            for line in content.split("\n")
+            if line.strip()
+        ],
+        "content": content,
+        "tokens": len(tokenizer.encode(content)),
+        "table_chunk_role": table_chunk_role,
+        "blockids": [],
+    }
+
+
+@pytest.mark.offline
+def test_is_heading_only_detection():
+    tokenizer = _make_tokenizer()
+    assert _is_heading_only(
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer)
+    )
+    # A glued accumulation of bare headings is still heading-only.
+    assert _is_heading_only(
+        _hblock("## 2.4\n\n### 2.4.1", heading="2.4", level=2, tokenizer=tokenizer)
+    )
+    # Heading + body is NOT heading-only.
+    assert not _is_heading_only(
+        _hblock("## 2.3\nbody text", heading="2.3", level=2, tokenizer=tokenizer)
+    )
+    # Preamble (no heading) is excluded by the heading guard.
+    assert not _is_heading_only(
+        _hblock("preamble text", heading="", level=1, tokenizer=tokenizer)
+    )
+
+
+@pytest.mark.offline
+def test_heading_only_glues_forward_into_deeper_child():
+    # `## 2.4` (heading-only) must bond with its deeper child `### 2.4.1`,
+    # NOT get appended to the previous same-level block `## 2.3`.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.3\n" + "a" * 40, heading="2.3", level=2, tokenizer=tokenizer),
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock(
+            "### 2.4.1\n" + "b" * 40, heading="2.4.1", level=3, tokenizer=tokenizer
+        ),
+    ]
+
+    out = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+
+    assert len(out) == 2
+    # 2.3 stays untouched — the lone heading was NOT glued onto its tail.
+    assert out[0]["heading"] == "2.3"
+    assert "## 2.4" not in out[0]["content"]
+    # The bonded group keeps the shallower parent identity (2.4 / level 2)
+    # but carries the child content.
+    assert out[1]["heading"] == "2.4"
+    assert out[1]["level"] == 2
+    assert "## 2.4" in out[1]["content"]
+    assert "### 2.4.1" in out[1]["content"]
+
+
+@pytest.mark.offline
+def test_heading_only_glue_respects_target_max_when_child_near_cap():
+    # The child fits target_max on its own, but prepending the heading-only
+    # parent line would tip the bonded block over the hard cap. The pre-pass
+    # must re-split so every emitted piece stays within target_max, while the
+    # parent heading still rides with the first piece (never detached).
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock(
+            "### 2.4.1\n" + "b" * 86, heading="2.4.1", level=3, tokenizer=tokenizer
+        ),
+    ]
+    # child alone = 10 + 86 = 96 ≤ 100; bonded = 6 + 2 + 96 = 104 > 100.
+    assert blocks[1]["tokens"] <= 100
+
+    out = _glue_heading_only_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=100,
+        target_ideal=75,
+        chunk_overlap_token_size=0,
+    )
+
+    assert len(out) >= 2
+    assert all(b["tokens"] <= 100 for b in out), [b["tokens"] for b in out]
+    # Parent heading is not detached — it leads the first emitted piece.
+    assert "## 2.4" in out[0]["content"]
+
+
+@pytest.mark.offline
+def test_heading_only_cap_split_does_not_orphan_when_body_has_no_anchor():
+    # Regression: child is near the cap and its body is ONE long paragraph
+    # (> _MAX_ANCHOR_CANDIDATE_LENGTH chars), so the only anchor candidate in
+    # the glued block is the child heading at index 1. The naive
+    # split-the-whole-block path sliced off `[## 2.4]` alone — a heading-only
+    # orphan that LevelMerge then re-absorbs backward, recreating the separation.
+    # The prefix-aware re-split must keep the heading with real body content.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock(
+            "### 2.4.1\n" + "b" * 110, heading="2.4.1", level=3, tokenizer=tokenizer
+        ),
+    ]
+    # child alone = 10 + 110 = 120 (near cap); bonded = 6 + 2 + 120 = 128 > 120.
+    assert blocks[1]["tokens"] <= 120
+
+    out = _glue_heading_only_blocks(
+        blocks,
+        tokenizer=tokenizer,
+        target_max=120,
+        target_ideal=90,
+        chunk_overlap_token_size=0,
+    )
+
+    assert all(b["tokens"] <= 120 for b in out), [b["tokens"] for b in out]
+    # No piece is a heading-only orphan.
+    assert not any(_is_heading_only(b) for b in out)
+    # The heading lines ride with real body content in the first piece.
+    assert "## 2.4" in out[0]["content"]
+    assert "### 2.4.1" in out[0]["content"]
+    assert "b" in out[0]["content"]
+
+
+@pytest.mark.offline
+def test_heading_only_cap_split_does_not_shrink_later_body_chunks():
+    # When a glued block must be re-split, only the FIRST piece carries the
+    # heading prefix. The body must split at the FULL target_max so later
+    # body-only chunks keep the full budget; only the first piece reserves room
+    # for the prefix. (Earlier code split the whole body at the reduced budget,
+    # over-fragmenting every later chunk to the leftover first-chunk budget.)
+    tokenizer = _make_tokenizer()
+    # Large heading prefix -> a small leftover budget if wrongly applied to all.
+    parent = "## " + "P" * 47  # 50 tokens
+    body = "\n".join("y" * 28 for _ in range(6))
+    blocks = [
+        _hblock(parent, heading="P", level=2, tokenizer=tokenizer),
+        _hblock("### c\n" + body, heading="c", level=3, tokenizer=tokenizer),
+    ]
+    target_max = 100
+    # prefix = "## P*47" + "### c" = 50 + 1 + 5 = 56 tokens; sep = 1.
+    reduced_max = target_max - 56 - 1  # = 43, the over-shrunk budget to beat.
+
+    out = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=target_max, target_ideal=75
+    )
+
+    # Cap still honoured everywhere.
+    assert all(b["tokens"] <= target_max for b in out), [b["tokens"] for b in out]
+    # The prefix rides with the first piece.
+    assert out[0]["content"].startswith(parent)
+    # Body-only pieces (everything after the first) keep the FULL budget — at
+    # least one exceeds the reduced prefix budget, proving they were not shrunk.
+    assert max(b["tokens"] for b in out[1:]) > reduced_max, [
+        b["tokens"] for b in out[1:]
+    ]
+
+
+@pytest.mark.offline
+def test_heading_only_cap_split_handles_prefix_larger_than_cap():
+    # Degenerate: the heading prefix alone exceeds target_max (a very long
+    # title, or a tiny chunk_token_size). There is no room to keep it whole, so
+    # the fused heading+body paragraph is itself character-split — every emitted
+    # piece must still honour the hard cap (no over-cap chunk escapes).
+    tokenizer = _make_tokenizer()
+    parent = "## " + "P" * 40  # 43 tokens, alone already > target_max below
+    blocks = [
+        _hblock(parent, heading="P", level=2, tokenizer=tokenizer),
+        _hblock("### c\n" + "y" * 20, heading="c", level=3, tokenizer=tokenizer),
+    ]
+    target_max = 30
+
+    out = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=target_max, target_ideal=22
+    )
+
+    # The hard cap is enforced on every piece (the buggy path emitted a ~51-token
+    # chunk because reduced_max clamped to 1 and the prefix was prepended whole).
+    assert all(b["tokens"] <= target_max for b in out), [b["tokens"] for b in out]
+    # Content is not dropped: heading text and body both survive across pieces.
+    joined = "".join(b["content"] for b in out)
+    assert "P" in joined and "y" in joined
+
+
+@pytest.mark.offline
+def test_heading_only_chain_collapses_to_shallowest_identity():
+    # `# 2` -> `## 2.4` -> `### 2.4.1` (body) collapses into one block whose
+    # identity is the shallowest heading (level 1).
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("# 2", heading="2", level=1, tokenizer=tokenizer),
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock(
+            "### 2.4.1\n" + "c" * 30, heading="2.4.1", level=3, tokenizer=tokenizer
+        ),
+    ]
+
+    out = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+
+    assert len(out) == 1
+    assert out[0]["heading"] == "2"
+    assert out[0]["level"] == 1
+    content = out[0]["content"]
+    assert "# 2" in content and "## 2.4" in content and "### 2.4.1" in content
+
+
+@pytest.mark.offline
+def test_heading_only_no_glue_when_next_not_deeper():
+    # Next block is same level -> no forced forward glue; left for LevelMerge.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock("## 2.5\nbody", heading="2.5", level=2, tokenizer=tokenizer),
+    ]
+
+    out = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+
+    assert len(out) == 2
+
+
+@pytest.mark.offline
+def test_heading_only_no_glue_into_middle_or_last_table_slice():
+    # A deeper next block that is a `middle`/`last` table slice must not absorb
+    # the heading-only block (only `none` and `first` are glue targets — see
+    # test_heading_only_glues_into_first_table_slice for the `first` case).
+    tokenizer = _make_tokenizer()
+    for role in ("middle", "last"):
+        blocks = [
+            _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+            _hblock(
+                '<table id="t" format="json">[[1]]</table>',
+                heading="2.4.1",
+                level=3,
+                tokenizer=tokenizer,
+                table_chunk_role=role,
+            ),
+        ]
+
+        out = _glue_heading_only_blocks(
+            blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+        )
+
+        assert len(out) == 2, role
+
+
+@pytest.mark.offline
+def test_heading_only_group_stays_separate_when_prev_is_saturated():
+    # Rule 1 end-to-end: `## 2.3` already reached target_ideal AND the bonded
+    # `2.4 + 2.4.1` group exceeds small_tail_threshold, so neither peer merging
+    # nor tail absorption pulls it backward — it stays its own chunk, with 2.4
+    # bonded to 2.4.1 (not to 2.3). (A group below small_tail_threshold could
+    # still be tail-absorbed into a saturated 2.3, which is acceptable since it
+    # would carry 2.4.1 along.)
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.3\n" + "a" * 200, heading="2.3", level=2, tokenizer=tokenizer),
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock(
+            "### 2.4.1\n" + "b" * 40, heading="2.4.1", level=3, tokenizer=tokenizer
+        ),
+    ]
+
+    glued = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+    final = _merge_small_blocks(
+        glued,
+        tokenizer=tokenizer,
+        target_max=2000,
+        target_ideal=150,
+        small_tail_threshold=12,
+    )
+
+    assert len(final) == 2
+    assert "## 2.4" not in final[0]["content"]
+    assert "## 2.4" in final[1]["content"] and "### 2.4.1" in final[1]["content"]
+
+
+@pytest.mark.offline
+def test_heading_only_group_backfills_into_unsaturated_prev():
+    # Rule 2 end-to-end: when `## 2.3` is still below target_ideal and the
+    # join fits target_max, LevelMerge packs 2.3 + 2.4 + 2.4.1 into one chunk.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.3\n" + "a" * 40, heading="2.3", level=2, tokenizer=tokenizer),
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock(
+            "### 2.4.1\n" + "b" * 40, heading="2.4.1", level=3, tokenizer=tokenizer
+        ),
+    ]
+
+    glued = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+    final = _merge_small_blocks(
+        glued,
+        tokenizer=tokenizer,
+        target_max=200,
+        target_ideal=150,
+        small_tail_threshold=12,
+    )
+
+    assert len(final) == 1
+    content = final[0]["content"]
+    assert "## 2.3" in content
+    assert "## 2.4" in content
+    assert "### 2.4.1" in content
+
+
+@pytest.mark.offline
+def test_heading_only_not_glued_into_deeper_prev():
+    # `## 2.4` (L2) has no deeper child after it; its previous block is the
+    # DEEPER `### 2.3.9` (L3). It must NOT be pulled backward into that deeper
+    # block — absorbing a shallower heading into a deeper chunk would invert the
+    # hierarchy. It stays separate, left for LevelMerge.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock(
+            "### 2.3.9\n" + "a" * 40, heading="2.3.9", level=3, tokenizer=tokenizer
+        ),
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+    ]
+
+    out = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+
+    assert len(out) == 2
+    assert out[0]["heading"] == "2.3.9"
+    assert "## 2.4" not in out[0]["content"]
+    assert out[1]["heading"] == "2.4"
+
+
+@pytest.mark.offline
+def test_heading_only_not_glued_into_same_level_prev():
+    # The previous block `## 2.3` is the SAME level (a sibling), not deeper, so
+    # the body-less `## 2.4` is not glued backward into it — that is the original
+    # mis-merge. It stays standalone for LevelMerge to handle.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.3\n" + "a" * 40, heading="2.3", level=2, tokenizer=tokenizer),
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+    ]
+
+    out = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+
+    assert len(out) == 2
+    assert "## 2.4" not in out[0]["content"]
+    assert out[1]["heading"] == "2.4"
+
+
+_TABLE_FIRST = '<table id="t" format="json">[["a","b"],["c","d"]]</table>'
+
+
+@pytest.mark.offline
+def test_heading_only_glues_into_first_table_slice():
+    # The deeper child's first emitted block is the "first" slice of a split
+    # table (its body is an oversized table). The pre-pass must glue the
+    # body-less `## 2.4` into it AND keep the "first" role so LevelMerge still
+    # cannot absorb it backward.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock(
+            "### 2.4.1\n" + _TABLE_FIRST,
+            heading="2.4.1",
+            level=3,
+            tokenizer=tokenizer,
+            table_chunk_role="first",
+        ),
+    ]
+
+    out = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+
+    assert len(out) == 1
+    assert "## 2.4" in out[0]["content"]
+    assert "### 2.4.1" in out[0]["content"]
+    assert _TABLE_FIRST in out[0]["content"]
+    # "first" role preserved so LevelMerge keeps the table boundary protected.
+    assert out[0]["table_chunk_role"] == "first"
+
+
+@pytest.mark.offline
+def test_heading_only_with_first_table_child_not_separated_by_stage_d():
+    # End-to-end: `## 2.4` whose child starts with a "first" table slice must
+    # NOT be left on `## 2.3`. After the glue keeps the merged block "first",
+    # LevelMerge cannot pull it backward into the previous sibling.
+    tokenizer = _make_tokenizer()
+    blocks = [
+        _hblock("## 2.3\n" + "a" * 40, heading="2.3", level=2, tokenizer=tokenizer),
+        _hblock("## 2.4", heading="2.4", level=2, tokenizer=tokenizer),
+        _hblock(
+            "### 2.4.1\n" + _TABLE_FIRST,
+            heading="2.4.1",
+            level=3,
+            tokenizer=tokenizer,
+            table_chunk_role="first",
+        ),
+    ]
+
+    glued = _glue_heading_only_blocks(
+        blocks, tokenizer=tokenizer, target_max=10000, target_ideal=7500
+    )
+    final = _merge_small_blocks(
+        glued,
+        tokenizer=tokenizer,
+        target_max=10000,
+        target_ideal=150,
+        small_tail_threshold=12,
+    )
+
+    # `## 2.4` rides with its table child, never glued onto `## 2.3`.
+    chunk_23 = next(b for b in final if b["content"].startswith("## 2.3"))
+    assert "## 2.4" not in chunk_23["content"]
+    chunk_with_table = next(b for b in final if _TABLE_FIRST in b["content"])
+    assert "## 2.4" in chunk_with_table["content"]
+    assert "### 2.4.1" in chunk_with_table["content"]

tests/chunker/test_paragraph_semantic_split_long_block.py

@@ -0,0 +1,323 @@
+"""Regression tests for paragraph-semantic AnchorSplit anchor selection."""
+
+import json
+
+import pytest
+
+from lightrag.chunker.paragraph_semantic import (
+    _split_long_block,
+    chunking_by_paragraph_semantic,
+)
+from lightrag.utils import Tokenizer, TokenizerInterface
+
+
+class _CharTokenizer(TokenizerInterface):
+    """1:1 character-to-token mapping — keeps math obvious in assertions."""
+
+    def encode(self, content: str):
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens):
+        return "".join(chr(t) for t in tokens)
+
+
+def _make_tokenizer() -> Tokenizer:
+    return Tokenizer(model_name="char", tokenizer=_CharTokenizer())
+
+
+def _write_blocks_jsonl(tmp_path, rows: list[dict]) -> str:
+    path = tmp_path / "doc.blocks.jsonl"
+    path.write_text(
+        "\n".join(json.dumps(row, ensure_ascii=False) for row in rows),
+        encoding="utf-8",
+    )
+    return str(path)
+
+
+@pytest.mark.offline
+def test_split_long_block_short_lead_then_huge_does_not_recurse():
+    # Reproduces the case where the only ≤100-char paragraph is at index 0:
+    # before the fix, the anchor at idx=0 was selected, slice_paras was
+    # empty, the tail was the original input, and the recursive guard
+    # re-entered _split_long_block with the same arguments forever.
+    tokenizer = _make_tokenizer()
+    paragraphs = [
+        {"text": "Short lead anchor."},  # idx 0 — short, but unusable as a divider
+        {"text": "x" * 4000},  # idx 1 — huge, no anchor inside
+    ]
+
+    blocks = _split_long_block(
+        paragraphs,
+        heading="Heading",
+        parent_headings=[],
+        level=2,
+        table_chunk_role="none",
+        tokenizer=tokenizer,
+        target_max=1000,
+        target_ideal=750,
+    )
+
+    # Falls through to the "no eligible anchor" branch and now defers to
+    # recursive-character splitting so ``target_max`` is honored without
+    # relying on the embedding-time hard fallback (which uses a different
+    # threshold).  The original recursion-guard contract still holds: the
+    # function returns a finite list rather than recursing forever.
+    assert len(blocks) > 1
+    assert all(b["tokens"] <= 1000 for b in blocks)
+    # Heading hierarchy is preserved on every R-derived sub-block.
+    assert all(b["heading"] == "Heading" for b in blocks)
+
+
+@pytest.mark.offline
+def test_split_long_block_no_anchor_pack_accounts_for_separator():
+    # The no-anchor greedy pack joins pieces with ``"\n"``, which costs
+    # tokens on its own. Without debiting that separator from the buffer
+    # budget, two pieces summing to exactly target_max produced a final
+    # chunk of ``target_max + 1`` tokens — silently violating the cap.
+    tokenizer = _make_tokenizer()
+    # Two paragraphs both > _MAX_ANCHOR_CANDIDATE_LENGTH (100 chars), so
+    # neither qualifies as an anchor and the no-anchor branch fires.
+    # Their lengths sum exactly to ``target_max`` (101 + 101 = 202),
+    # so before the fix the joined output overflowed by the "\n" token.
+    paragraphs = [
+        {"text": "a" * 101},
+        {"text": "b" * 101},
+    ]
+
+    blocks = _split_long_block(
+        paragraphs,
+        heading="Heading",
+        parent_headings=[],
+        level=2,
+        table_chunk_role="none",
+        tokenizer=tokenizer,
+        target_max=202,
+        target_ideal=150,
+    )
+
+    assert blocks, "expected at least one sub-block"
+    assert all(b["tokens"] <= 202 for b in blocks), [b["tokens"] for b in blocks]
+
+
+@pytest.mark.offline
+def test_split_long_block_single_paragraph_oversized_is_character_split():
+    # A single oversized paragraph used to trigger the early-return at
+    # ``len(paragraphs) <= 1`` and the recursive-guard's ``> 1`` clause,
+    # so the function emitted one ~total-token block that silently
+    # blew past target_max. With both gates relaxed, the no-anchor
+    # branch's character fallback honors the cap on this case too.
+    tokenizer = _make_tokenizer()
+    paragraphs = [{"text": "x" * 4000}]
+
+    blocks = _split_long_block(
+        paragraphs,
+        heading="Heading",
+        parent_headings=[],
+        level=2,
+        table_chunk_role="none",
+        tokenizer=tokenizer,
+        target_max=1000,
+        target_ideal=750,
+    )
+
+    assert len(blocks) > 1, "single oversized paragraph must be split, not kept whole"
+    assert all(b["tokens"] <= 1000 for b in blocks), [b["tokens"] for b in blocks]
+    # Heading hierarchy is preserved on every R-derived sub-block.
+    assert all(b["heading"] == "Heading" for b in blocks)
+
+
+@pytest.mark.offline
+def test_split_long_block_character_fallback_keeps_configured_overlap(monkeypatch):
+    tokenizer = _make_tokenizer()
+    captured: dict[str, int] = {}
+
+    def fake_chunker(
+        tokenizer,
+        content,
+        chunk_token_size: int = 1200,
+        *,
+        chunk_overlap_token_size: int = 100,
+        separators=None,
+    ):
+        captured["chunk_overlap_token_size"] = chunk_overlap_token_size
+        step = max(chunk_token_size - chunk_overlap_token_size, 1)
+        tokens = tokenizer.encode(content)
+        chunks = []
+        for start in range(0, len(tokens), step):
+            piece = tokenizer.decode(tokens[start : start + chunk_token_size])
+            chunks.append(
+                {
+                    "tokens": len(tokenizer.encode(piece)),
+                    "content": piece,
+                    "chunk_order_index": len(chunks),
+                }
+            )
+        return chunks
+
+    import lightrag.chunker.recursive_character as rc_mod
+
+    monkeypatch.setattr(rc_mod, "chunking_by_recursive_character", fake_chunker)
+
+    blocks = _split_long_block(
+        [{"text": "x" * 260}],
+        heading="Heading",
+        parent_headings=[],
+        level=2,
+        table_chunk_role="none",
+        tokenizer=tokenizer,
+        target_max=100,
+        target_ideal=75,
+        chunk_overlap_token_size=25,
+    )
+
+    assert captured["chunk_overlap_token_size"] == 25
+    assert len(blocks) > 1
+    assert blocks[0]["content"][-25:] == blocks[1]["content"][:25]
+
+
+@pytest.mark.offline
+def test_split_long_block_uses_later_short_anchor():
+    # Sanity check: a short paragraph at idx>0 IS still a valid divider.
+    tokenizer = _make_tokenizer()
+    paragraphs = [
+        {"text": "x" * 1500},  # idx 0 — huge
+        {"text": "Mid anchor."},  # idx 1 — short, eligible
+        {"text": "y" * 1500},  # idx 2 — huge
+    ]
+
+    blocks = _split_long_block(
+        paragraphs,
+        heading="Heading",
+        parent_headings=[],
+        level=2,
+        table_chunk_role="none",
+        tokenizer=tokenizer,
+        target_max=1000,
+        target_ideal=750,
+    )
+
+    assert len(blocks) >= 2
+    # Anchor paragraph becomes the heading of the post-split sub-block.
+    assert any(b["heading"] == "Mid anchor." for b in blocks)
+
+
+@pytest.mark.offline
+def test_public_chunking_keeps_unsplit_heading_without_part_suffix(tmp_path):
+    tokenizer = _make_tokenizer()
+    blocks_path = _write_blocks_jsonl(
+        tmp_path,
+        [
+            {
+                "type": "content",
+                "heading": "Heading",
+                "parent_headings": [],
+                "level": 2,
+                "content": "short body",
+            }
+        ],
+    )
+
+    chunks = chunking_by_paragraph_semantic(
+        tokenizer,
+        "short body",
+        chunk_token_size=100,
+        blocks_path=blocks_path,
+    )
+
+    assert len(chunks) == 1
+    assert chunks[0]["heading"]["heading"] == "Heading"
+
+
+@pytest.mark.offline
+def test_public_chunking_adds_part_suffixes_for_anchor_split(tmp_path):
+    tokenizer = _make_tokenizer()
+    body = "\n".join(["x" * 800, "Mid anchor.", "y" * 800])
+    blocks_path = _write_blocks_jsonl(
+        tmp_path,
+        [
+            {
+                "type": "content",
+                "heading": "Heading",
+                "parent_headings": [],
+                "level": 2,
+                "content": body,
+            }
+        ],
+    )
+
+    chunks = chunking_by_paragraph_semantic(
+        tokenizer,
+        body,
+        chunk_token_size=1000,
+        blocks_path=blocks_path,
+        chunk_overlap_token_size=0,
+    )
+
+    assert [chunk["heading"]["heading"] for chunk in chunks] == [
+        "Heading [part 1]",
+        "Mid anchor. [part 2]",
+    ]
+    assert all(
+        all("[part " not in parent for parent in chunk["heading"]["parent_headings"])
+        for chunk in chunks
+    )
+
+
+@pytest.mark.offline
+def test_public_chunking_adds_part_suffixes_for_long_text_fallback(
+    tmp_path, monkeypatch
+):
+    tokenizer = _make_tokenizer()
+
+    def fake_chunker(
+        tokenizer,
+        content,
+        chunk_token_size: int = 1200,
+        *,
+        chunk_overlap_token_size: int = 100,
+        separators=None,
+    ):
+        tokens = tokenizer.encode(content)
+        chunks = []
+        for start in range(0, len(tokens), chunk_token_size):
+            piece = tokenizer.decode(tokens[start : start + chunk_token_size])
+            chunks.append(
+                {
+                    "tokens": len(tokenizer.encode(piece)),
+                    "content": piece,
+                    "chunk_order_index": len(chunks),
+                }
+            )
+        return chunks
+
+    import lightrag.chunker.recursive_character as rc_mod
+
+    monkeypatch.setattr(rc_mod, "chunking_by_recursive_character", fake_chunker)
+
+    body = "z" * 260
+    blocks_path = _write_blocks_jsonl(
+        tmp_path,
+        [
+            {
+                "type": "content",
+                "heading": "Heading",
+                "parent_headings": [],
+                "level": 2,
+                "content": body,
+            }
+        ],
+    )
+
+    chunks = chunking_by_paragraph_semantic(
+        tokenizer,
+        body,
+        chunk_token_size=100,
+        blocks_path=blocks_path,
+        chunk_overlap_token_size=0,
+    )
+
+    assert [chunk["heading"]["heading"] for chunk in chunks] == [
+        "Heading [part 1]",
+        "Heading [part 2]",
+        "Heading [part 3]",
+    ]

tests/chunker/test_rerank_chunking.py

@@ -0,0 +1,564 @@
+"""
+Unit tests for rerank document chunking functionality.
+
+Tests the chunk_documents_for_rerank and aggregate_chunk_scores functions
+in lightrag/rerank.py to ensure proper document splitting and score aggregation.
+"""
+
+import pytest
+from unittest.mock import Mock, patch, AsyncMock
+from lightrag.rerank import (
+    chunk_documents_for_rerank,
+    aggregate_chunk_scores,
+    cohere_rerank,
+)
+
+
+class TestChunkDocumentsForRerank:
+    """Test suite for chunk_documents_for_rerank function"""
+
+    def test_no_chunking_needed_for_short_docs(self):
+        """Documents shorter than max_tokens should not be chunked"""
+        documents = [
+            "Short doc 1",
+            "Short doc 2",
+            "Short doc 3",
+        ]
+
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=100, overlap_tokens=10
+        )
+
+        # No chunking should occur
+        assert len(chunked_docs) == 3
+        assert chunked_docs == documents
+        assert doc_indices == [0, 1, 2]
+
+    def test_chunking_with_character_fallback(self):
+        """Test chunking falls back to character-based when tokenizer unavailable"""
+        # Create a very long document that exceeds character limit
+        long_doc = "a" * 2000  # 2000 characters
+        documents = [long_doc, "short doc"]
+
+        with patch("lightrag.utils.TiktokenTokenizer", side_effect=ImportError):
+            chunked_docs, doc_indices = chunk_documents_for_rerank(
+                documents,
+                max_tokens=100,  # 100 tokens = ~400 chars
+                overlap_tokens=10,  # 10 tokens = ~40 chars
+            )
+
+        # First doc should be split into chunks, second doc stays whole
+        assert len(chunked_docs) > 2  # At least one chunk from first doc + second doc
+        assert chunked_docs[-1] == "short doc"  # Last chunk is the short doc
+        # Verify doc_indices maps chunks to correct original document
+        assert doc_indices[-1] == 1  # Last chunk maps to document 1
+
+    def test_chunking_with_tiktoken_tokenizer(self):
+        """Test chunking with actual tokenizer"""
+        # Create document with known token count
+        # Approximate: "word " = ~1 token, so 200 words ~ 200 tokens
+        long_doc = " ".join([f"word{i}" for i in range(200)])
+        documents = [long_doc, "short"]
+
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=50, overlap_tokens=10
+        )
+
+        # Long doc should be split, short doc should remain
+        assert len(chunked_docs) > 2
+        assert doc_indices[-1] == 1  # Last chunk is from second document
+
+        # Verify overlapping chunks contain overlapping content
+        if len(chunked_docs) > 2:
+            # Check that consecutive chunks from same doc have some overlap
+            for i in range(len(doc_indices) - 1):
+                if doc_indices[i] == doc_indices[i + 1] == 0:
+                    # Both chunks from first doc, should have overlap
+                    chunk1_words = chunked_docs[i].split()
+                    chunk2_words = chunked_docs[i + 1].split()
+                    # At least one word should be common due to overlap
+                    assert any(word in chunk2_words for word in chunk1_words[-5:])
+
+    def test_empty_documents(self):
+        """Test handling of empty document list"""
+        documents = []
+        chunked_docs, doc_indices = chunk_documents_for_rerank(documents)
+
+        assert chunked_docs == []
+        assert doc_indices == []
+
+    def test_single_document_chunking(self):
+        """Test chunking of a single long document"""
+        # Create document with ~100 tokens
+        long_doc = " ".join([f"token{i}" for i in range(100)])
+        documents = [long_doc]
+
+        chunked_docs, doc_indices = chunk_documents_for_rerank(
+            documents, max_tokens=30, overlap_tokens=5
+        )
+
+        # Should create multiple chunks
+        assert len(chunked_docs) > 1
+        # All chunks should map to document 0
+        assert all(idx == 0 for idx in doc_indices)
+
+
+class TestAggregateChunkScores:
+    """Test suite for aggregate_chunk_scores function"""
+
+    def test_no_chunking_simple_aggregation(self):
+        """Test aggregation when no chunking occurred (1:1 mapping)"""
+        chunk_results = [
+            {"index": 0, "relevance_score": 0.9},
+            {"index": 1, "relevance_score": 0.7},
+            {"index": 2, "relevance_score": 0.5},
+        ]
+        doc_indices = [0, 1, 2]  # 1:1 mapping
+        num_original_docs = 3
+
+        aggregated = aggregate_chunk_scores(
+            chunk_results, doc_indices, num_original_docs, aggregation="max"
+        )
+
+        # Results should be sorted by score
+        assert len(aggregated) == 3
+        assert aggregated[0]["index"] == 0
+        assert aggregated[0]["relevance_score"] == 0.9
+        assert aggregated[1]["index"] == 1
+        assert aggregated[1]["relevance_score"] == 0.7
+        assert aggregated[2]["index"] == 2
+        assert aggregated[2]["relevance_score"] == 0.5
+
+    def test_max_aggregation_with_chunks(self):
+        """Test max aggregation strategy with multiple chunks per document"""
+        # 5 chunks: first 3 from doc 0, last 2 from doc 1
+        chunk_results = [
+            {"index": 0, "relevance_score": 0.5},
+            {"index": 1, "relevance_score": 0.8},
+            {"index": 2, "relevance_score": 0.6},
+            {"index": 3, "relevance_score": 0.7},
+            {"index": 4, "relevance_score": 0.4},
+        ]
+        doc_indices = [0, 0, 0, 1, 1]
+        num_original_docs = 2
+
+        aggregated = aggregate_chunk_scores(
+            chunk_results, doc_indices, num_original_docs, aggregation="max"
+        )
+
+        # Should take max score for each document
+        assert len(aggregated) == 2
+        assert aggregated[0]["index"] == 0
+        assert aggregated[0]["relevance_score"] == 0.8  # max of 0.5, 0.8, 0.6
+        assert aggregated[1]["index"] == 1
+        assert aggregated[1]["relevance_score"] == 0.7  # max of 0.7, 0.4
+
+    def test_mean_aggregation_with_chunks(self):
+        """Test mean aggregation strategy"""
+        chunk_results = [
+            {"index": 0, "relevance_score": 0.6},
+            {"index": 1, "relevance_score": 0.8},
+            {"index": 2, "relevance_score": 0.4},
+        ]
+        doc_indices = [0, 0, 1]  # First two chunks from doc 0, last from doc 1
+        num_original_docs = 2
+
+        aggregated = aggregate_chunk_scores(
+            chunk_results, doc_indices, num_original_docs, aggregation="mean"
+        )
+
+        assert len(aggregated) == 2
+        assert aggregated[0]["index"] == 0
+        assert aggregated[0]["relevance_score"] == pytest.approx(0.7)  # (0.6 + 0.8) / 2
+        assert aggregated[1]["index"] == 1
+        assert aggregated[1]["relevance_score"] == 0.4
+
+    def test_first_aggregation_with_chunks(self):
+        """Test first aggregation strategy"""
+        chunk_results = [
+            {"index": 0, "relevance_score": 0.6},
+            {"index": 1, "relevance_score": 0.8},
+            {"index": 2, "relevance_score": 0.4},
+        ]
+        doc_indices = [0, 0, 1]
+        num_original_docs = 2
+
+        aggregated = aggregate_chunk_scores(
+            chunk_results, doc_indices, num_original_docs, aggregation="first"
+        )
+
+        assert len(aggregated) == 2
+        # First should use first score seen for each doc
+        assert aggregated[0]["index"] == 0
+        assert aggregated[0]["relevance_score"] == 0.6  # First score for doc 0
+        assert aggregated[1]["index"] == 1
+        assert aggregated[1]["relevance_score"] == 0.4
+
+    def test_empty_chunk_results(self):
+        """Test handling of empty results"""
+        aggregated = aggregate_chunk_scores([], [], 3, aggregation="max")
+        assert aggregated == []
+
+    def test_documents_with_no_scores(self):
+        """Test when some documents have no chunks/scores"""
+        chunk_results = [
+            {"index": 0, "relevance_score": 0.9},
+            {"index": 1, "relevance_score": 0.7},
+        ]
+        doc_indices = [0, 0]  # Both chunks from document 0
+        num_original_docs = 3  # But we have 3 documents total
+
+        aggregated = aggregate_chunk_scores(
+            chunk_results, doc_indices, num_original_docs, aggregation="max"
+        )
+
+        # Only doc 0 should appear in results
+        assert len(aggregated) == 1
+        assert aggregated[0]["index"] == 0
+
+    def test_unknown_aggregation_strategy(self):
+        """Test that unknown strategy falls back to max"""
+        chunk_results = [
+            {"index": 0, "relevance_score": 0.6},
+            {"index": 1, "relevance_score": 0.8},
+        ]
+        doc_indices = [0, 0]
+        num_original_docs = 1
+
+        # Use invalid strategy
+        aggregated = aggregate_chunk_scores(
+            chunk_results, doc_indices, num_original_docs, aggregation="invalid"
+        )
+
+        # Should fall back to max
+        assert aggregated[0]["relevance_score"] == 0.8
+
+
+@pytest.mark.offline
+class TestTopNWithChunking:
+    """Tests for top_n behavior when chunking is enabled (Bug fix verification)"""
+
+    @pytest.mark.asyncio
+    async def test_top_n_limits_documents_not_chunks(self):
+        """
+        Test that top_n correctly limits documents (not chunks) when chunking is enabled.
+
+        Bug scenario: 10 docs expand to 50 chunks. With old behavior, top_n=5 would
+        return scores for only 5 chunks (possibly all from 1-2 docs). After aggregation,
+        fewer than 5 documents would be returned.
+
+        Fixed behavior: top_n=5 should return exactly 5 documents after aggregation.
+        """
+        # Setup: 5 documents, each producing multiple chunks when chunked
+        # Using small max_tokens to force chunking
+        long_docs = [" ".join([f"doc{i}_word{j}" for j in range(50)]) for i in range(5)]
+        query = "test query"
+
+        # First, determine how many chunks will be created by actual chunking
+        _, doc_indices = chunk_documents_for_rerank(
+            long_docs, max_tokens=50, overlap_tokens=10
+        )
+        num_chunks = len(doc_indices)
+
+        # Mock API returns scores for ALL chunks (simulating disabled API-level top_n)
+        # Give different scores to ensure doc 0 gets highest, doc 1 second, etc.
+        # Assign scores based on original document index (lower doc index = higher score)
+        mock_chunk_scores = []
+        for i in range(num_chunks):
+            original_doc = doc_indices[i]
+            # Higher score for lower doc index, with small variation per chunk
+            base_score = 0.9 - (original_doc * 0.1)
+            mock_chunk_scores.append({"index": i, "relevance_score": base_score})
+
+        mock_response = Mock()
+        mock_response.status = 200
+        mock_response.json = AsyncMock(return_value={"results": mock_chunk_scores})
+        mock_response.request_info = None
+        mock_response.history = None
+        mock_response.headers = {}
+        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
+        mock_response.__aexit__ = AsyncMock(return_value=None)
+
+        mock_session = Mock()
+        mock_session.post = Mock(return_value=mock_response)
+        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_session.__aexit__ = AsyncMock(return_value=None)
+
+        with patch("lightrag.rerank.aiohttp.ClientSession", return_value=mock_session):
+            result = await cohere_rerank(
+                query=query,
+                documents=long_docs,
+                api_key="test-key",
+                base_url="http://test.com/rerank",
+                enable_chunking=True,
+                max_tokens_per_doc=50,  # Match chunking above
+                top_n=3,  # Request top 3 documents
+            )
+
+            # Verify: should get exactly 3 documents (not unlimited chunks)
+            assert len(result) == 3
+            # All results should have valid document indices (0-4)
+            assert all(0 <= r["index"] < 5 for r in result)
+            # Results should be sorted by score (descending)
+            assert all(
+                result[i]["relevance_score"] >= result[i + 1]["relevance_score"]
+                for i in range(len(result) - 1)
+            )
+            # The top 3 docs should be 0, 1, 2 (highest scores)
+            result_indices = [r["index"] for r in result]
+            assert set(result_indices) == {0, 1, 2}
+
+    @pytest.mark.asyncio
+    async def test_api_receives_no_top_n_when_chunking_enabled(self):
+        """
+        Test that the API request does NOT include top_n when chunking is enabled.
+
+        This ensures all chunk scores are retrieved for proper aggregation.
+        """
+        documents = [" ".join([f"word{i}" for i in range(100)]), "short doc"]
+        query = "test query"
+
+        captured_payload = {}
+
+        mock_response = Mock()
+        mock_response.status = 200
+        mock_response.json = AsyncMock(
+            return_value={
+                "results": [
+                    {"index": 0, "relevance_score": 0.9},
+                    {"index": 1, "relevance_score": 0.8},
+                    {"index": 2, "relevance_score": 0.7},
+                ]
+            }
+        )
+        mock_response.request_info = None
+        mock_response.history = None
+        mock_response.headers = {}
+        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
+        mock_response.__aexit__ = AsyncMock(return_value=None)
+
+        def capture_post(*args, **kwargs):
+            captured_payload.update(kwargs.get("json", {}))
+            return mock_response
+
+        mock_session = Mock()
+        mock_session.post = Mock(side_effect=capture_post)
+        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_session.__aexit__ = AsyncMock(return_value=None)
+
+        with patch("lightrag.rerank.aiohttp.ClientSession", return_value=mock_session):
+            await cohere_rerank(
+                query=query,
+                documents=documents,
+                api_key="test-key",
+                base_url="http://test.com/rerank",
+                enable_chunking=True,
+                max_tokens_per_doc=30,
+                top_n=1,  # User wants top 1 document
+            )
+
+            # Verify: API payload should NOT have top_n (disabled for chunking)
+            assert "top_n" not in captured_payload
+
+    @pytest.mark.asyncio
+    async def test_top_n_not_modified_when_chunking_disabled(self):
+        """
+        Test that top_n is passed through to API when chunking is disabled.
+        """
+        documents = ["doc1", "doc2"]
+        query = "test query"
+
+        captured_payload = {}
+
+        mock_response = Mock()
+        mock_response.status = 200
+        mock_response.json = AsyncMock(
+            return_value={
+                "results": [
+                    {"index": 0, "relevance_score": 0.9},
+                ]
+            }
+        )
+        mock_response.request_info = None
+        mock_response.history = None
+        mock_response.headers = {}
+        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
+        mock_response.__aexit__ = AsyncMock(return_value=None)
+
+        def capture_post(*args, **kwargs):
+            captured_payload.update(kwargs.get("json", {}))
+            return mock_response
+
+        mock_session = Mock()
+        mock_session.post = Mock(side_effect=capture_post)
+        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_session.__aexit__ = AsyncMock(return_value=None)
+
+        with patch("lightrag.rerank.aiohttp.ClientSession", return_value=mock_session):
+            await cohere_rerank(
+                query=query,
+                documents=documents,
+                api_key="test-key",
+                base_url="http://test.com/rerank",
+                enable_chunking=False,  # Chunking disabled
+                top_n=1,
+            )
+
+            # Verify: API payload should have top_n when chunking is disabled
+            assert captured_payload.get("top_n") == 1
+
+
+@pytest.mark.offline
+class TestCohereRerankChunking:
+    """Integration tests for cohere_rerank with chunking enabled"""
+
+    @pytest.mark.asyncio
+    async def test_cohere_rerank_with_chunking_disabled(self):
+        """Test that chunking can be disabled"""
+        documents = ["doc1", "doc2"]
+        query = "test query"
+
+        # Mock the generic_rerank_api
+        with patch(
+            "lightrag.rerank.generic_rerank_api", new_callable=AsyncMock
+        ) as mock_api:
+            mock_api.return_value = [
+                {"index": 0, "relevance_score": 0.9},
+                {"index": 1, "relevance_score": 0.7},
+            ]
+
+            result = await cohere_rerank(
+                query=query,
+                documents=documents,
+                api_key="test-key",
+                enable_chunking=False,
+                max_tokens_per_doc=100,
+            )
+
+            # Verify generic_rerank_api was called with correct parameters
+            mock_api.assert_called_once()
+            call_kwargs = mock_api.call_args[1]
+            assert call_kwargs["enable_chunking"] is False
+            assert call_kwargs["max_tokens_per_doc"] == 100
+            # Result should mirror mocked scores
+            assert len(result) == 2
+            assert result[0]["index"] == 0
+            assert result[0]["relevance_score"] == 0.9
+            assert result[1]["index"] == 1
+            assert result[1]["relevance_score"] == 0.7
+
+    @pytest.mark.asyncio
+    async def test_cohere_rerank_with_chunking_enabled(self):
+        """Test that chunking parameters are passed through"""
+        documents = ["doc1", "doc2"]
+        query = "test query"
+
+        with patch(
+            "lightrag.rerank.generic_rerank_api", new_callable=AsyncMock
+        ) as mock_api:
+            mock_api.return_value = [
+                {"index": 0, "relevance_score": 0.9},
+                {"index": 1, "relevance_score": 0.7},
+            ]
+
+            result = await cohere_rerank(
+                query=query,
+                documents=documents,
+                api_key="test-key",
+                enable_chunking=True,
+                max_tokens_per_doc=480,
+            )
+
+            # Verify parameters were passed
+            call_kwargs = mock_api.call_args[1]
+            assert call_kwargs["enable_chunking"] is True
+            assert call_kwargs["max_tokens_per_doc"] == 480
+            # Result should mirror mocked scores
+            assert len(result) == 2
+            assert result[0]["index"] == 0
+            assert result[0]["relevance_score"] == 0.9
+            assert result[1]["index"] == 1
+            assert result[1]["relevance_score"] == 0.7
+
+    @pytest.mark.asyncio
+    async def test_cohere_rerank_default_parameters(self):
+        """Test default parameter values for cohere_rerank"""
+        documents = ["doc1"]
+        query = "test"
+
+        with patch(
+            "lightrag.rerank.generic_rerank_api", new_callable=AsyncMock
+        ) as mock_api:
+            mock_api.return_value = [{"index": 0, "relevance_score": 0.9}]
+
+            result = await cohere_rerank(
+                query=query, documents=documents, api_key="test-key"
+            )
+
+            # Verify default values
+            call_kwargs = mock_api.call_args[1]
+            assert call_kwargs["enable_chunking"] is False
+            assert call_kwargs["max_tokens_per_doc"] == 4096
+            assert call_kwargs["model"] == "rerank-v3.5"
+            # Result should mirror mocked scores
+            assert len(result) == 1
+            assert result[0]["index"] == 0
+            assert result[0]["relevance_score"] == 0.9
+
+
+@pytest.mark.offline
+class TestEndToEndChunking:
+    """End-to-end tests for chunking workflow"""
+
+    @pytest.mark.asyncio
+    async def test_end_to_end_chunking_workflow(self):
+        """Test complete chunking workflow from documents to aggregated results"""
+        # Create documents where first one needs chunking
+        long_doc = " ".join([f"word{i}" for i in range(100)])
+        documents = [long_doc, "short doc"]
+        query = "test query"
+
+        # Mock the HTTP call inside generic_rerank_api
+        mock_response = Mock()
+        mock_response.status = 200
+        mock_response.json = AsyncMock(
+            return_value={
+                "results": [
+                    {"index": 0, "relevance_score": 0.5},  # chunk 0 from doc 0
+                    {"index": 1, "relevance_score": 0.8},  # chunk 1 from doc 0
+                    {"index": 2, "relevance_score": 0.6},  # chunk 2 from doc 0
+                    {"index": 3, "relevance_score": 0.7},  # doc 1 (short)
+                ]
+            }
+        )
+        mock_response.request_info = None
+        mock_response.history = None
+        mock_response.headers = {}
+        # Make mock_response an async context manager (for `async with session.post() as response`)
+        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
+        mock_response.__aexit__ = AsyncMock(return_value=None)
+
+        mock_session = Mock()
+        # session.post() returns an async context manager, so return mock_response which is now one
+        mock_session.post = Mock(return_value=mock_response)
+        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
+        mock_session.__aexit__ = AsyncMock(return_value=None)
+
+        with patch("lightrag.rerank.aiohttp.ClientSession", return_value=mock_session):
+            result = await cohere_rerank(
+                query=query,
+                documents=documents,
+                api_key="test-key",
+                base_url="http://test.com/rerank",
+                enable_chunking=True,
+                max_tokens_per_doc=30,  # Force chunking of long doc
+            )
+
+            # Should get 2 results (one per original document)
+            # The long doc's chunks should be aggregated
+            assert len(result) <= len(documents)
+            # Results should be sorted by score
+            assert all(
+                result[i]["relevance_score"] >= result[i + 1]["relevance_score"]
+                for i in range(len(result) - 1)
+            )

tests/chunker/test_sidecar_backfill_integration.py

@@ -0,0 +1,391 @@
+"""Integration tests: real F chunker + hard-split + sidecar backfill + chunks dict.
+
+Unlike ``tests/sidecar/test_backfill.py`` (which hand-builds chunk lists), these
+drive the actual ``chunking_by_fixed_token`` chunker over the exact merged text a
+parsed document would yield, then apply the real
+``enforce_chunk_token_limit_before_embedding`` hard-split and
+``build_chunks_dict_from_chunking_result`` persistence step — verifying that
+backfilled sidecars are precise per final slice and survive into the chunks dict.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+from lightrag.chunker import chunking_by_fixed_token
+from lightrag.sidecar import backfill_chunk_sidecars
+from lightrag.utils import (
+    Tokenizer,
+    enforce_chunk_token_limit_before_embedding,
+)
+from lightrag.utils_pipeline import build_chunks_dict_from_chunking_result
+
+_BLOCK_SEPARATOR = "\n\n"
+
+
+class _CharTokenizerImpl:
+    """Deterministic char-per-token tokenizer; decode(encode(x)) == x so F
+    chunks are verbatim substrings and token sizes are character counts."""
+
+    def encode(self, content: str) -> list[int]:
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens: list[int]) -> str:
+        return "".join(chr(t) for t in tokens)
+
+
+def _tokenizer() -> Tokenizer:
+    return Tokenizer("char", _CharTokenizerImpl())
+
+
+def _write_blocks(tmp_path: Path, blocks: list[tuple[str, str]]) -> tuple[str, str]:
+    """Write blocks.jsonl; return (path, merged_text)."""
+    path = tmp_path / "doc.blocks.jsonl"
+    lines = [json.dumps({"type": "meta", "format": "lightrag", "version": "1.0"})]
+    parts: list[str] = []
+    for blockid, content in blocks:
+        lines.append(
+            json.dumps(
+                {
+                    "type": "content",
+                    "blockid": blockid,
+                    "content": content,
+                    "heading": "",
+                    "parent_headings": [],
+                    "level": 1,
+                }
+            )
+        )
+        if content.strip():
+            parts.append(content)
+    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    return str(path), _BLOCK_SEPARATOR.join(parts)
+
+
+@pytest.mark.offline
+def test_real_overlap_tail_chunk_maps_to_next_block(tmp_path: Path) -> None:
+    # End-to-end reproduction of the overlap-tail ambiguity with the real chunker.
+    # Simple case: b1="aa", b2="a", chunk_size=3, overlap=1 -> ["aa", "a", "a"]. The
+    # middle window [2:5] = "\n\na" strips to b2's "a" (the overlap landed on the
+    # separator), so the tail chunks belong to b2 — not the earlier "a" inside b1.
+    blocks_path, merged = _write_blocks(tmp_path, [("b1", "aa"), ("b2", "a")])
+    tok = _tokenizer()
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=3,
+        chunk_overlap_token_size=1,
+        _emit_source_span=True,
+    )
+    assert [c["content"] for c in chunks] == ["aa", "a", "a"]
+
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    assert [r["id"] for r in chunks[0]["sidecar"]["refs"]] == ["b1"]
+    assert [r["id"] for r in chunks[1]["sidecar"]["refs"]] == ["b2"]
+    assert [r["id"] for r in chunks[2]["sidecar"]["refs"]] == ["b2"]
+
+
+@pytest.mark.offline
+def test_real_overlap_on_stripped_separator_does_not_strand_chunks(
+    tmp_path: Path,
+) -> None:
+    # Harder end-to-end case: b1="a", b2="abab", chunk_size=2, overlap=1 ->
+    # ["a", "", "a", "ab", "ba", "ab", "b"]. The token overlap repeatedly lands on the
+    # stripped separator, so consecutive non-empty chunks can share a normalized start
+    # and only the end advances. The empty chunk is skipped; every other tail chunk
+    # must resolve into b2 without raising ChunkBlockMatchError.
+    blocks_path, merged = _write_blocks(tmp_path, [("b1", "a"), ("b2", "abab")])
+    tok = _tokenizer()
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=2,
+        chunk_overlap_token_size=1,
+        _emit_source_span=True,
+    )
+    assert [c["content"] for c in chunks] == ["a", "", "a", "ab", "ba", "ab", "b"]
+
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    assert [r["id"] for r in chunks[0]["sidecar"]["refs"]] == ["b1"]
+    # The empty chunk is skipped (no sidecar); all remaining chunks belong to b2.
+    assert "sidecar" not in chunks[1]
+    for ch in chunks[2:]:
+        assert [r["id"] for r in ch["sidecar"]["refs"]] == ["b2"]
+
+
+@pytest.mark.offline
+def test_real_identical_adjacent_blocks_no_cross_block_artifact(
+    tmp_path: Path,
+) -> None:
+    # End-to-end: identical adjacent blocks b1="aa", b2="aa", chunk_size=2, overlap=0
+    # -> ["aa", "", "aa"]. Stripping glues them to "aaaa", where "aa" also matches at
+    # offset 1 across the (removed) separator. The final chunk must reference only b2,
+    # not spuriously span both blocks.
+    blocks_path, merged = _write_blocks(tmp_path, [("b1", "aa"), ("b2", "aa")])
+    tok = _tokenizer()
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=2,
+        chunk_overlap_token_size=0,
+        _emit_source_span=True,
+    )
+    assert [c["content"] for c in chunks] == ["aa", "", "aa"]
+
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    assert chunks[0]["_source_span"] == {"start": 0, "end": 2}
+    assert chunks[2]["_source_span"] == {"start": 4, "end": 6}
+
+    assert [r["id"] for r in chunks[0]["sidecar"]["refs"]] == ["b1"]
+    assert "sidecar" not in chunks[1]
+    assert [r["id"] for r in chunks[2]["sidecar"]["refs"]] == ["b2"]
+
+
+@pytest.mark.offline
+def test_real_fixed_token_chunks_get_full_coverage(tmp_path: Path) -> None:
+    blocks = [(f"b{i}", f"Block number {i} body content here.") for i in range(6)]
+    blocks_path, merged = _write_blocks(tmp_path, blocks)
+    tok = _tokenizer()
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=40,
+        chunk_overlap_token_size=5,
+        _emit_source_span=True,
+    )
+    assert len(chunks) >= 2  # multiple blocks per chunk at this size
+
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    # Every chunk is located and carries block provenance.
+    for ch in chunks:
+        assert ch["sidecar"]["type"] == "block"
+        assert ch["sidecar"]["refs"]
+    # Union of all refs covers every block.
+    seen = {r["id"] for ch in chunks for r in ch["sidecar"]["refs"]}
+    assert seen == {b for b, _ in blocks}
+
+
+@pytest.mark.offline
+def test_hard_split_slices_get_precise_refs(tmp_path: Path) -> None:
+    # One large block followed by a small one. A single fixed-token chunk covers
+    # both; the hard-split then breaks the big-content chunk into slices that
+    # each lie inside one block, so refs must NOT all inherit the full set.
+    big = "A" * 120
+    blocks_path, merged = _write_blocks(tmp_path, [("big", big), ("small", "tail.")])
+    tok = _tokenizer()
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=200,
+        chunk_overlap_token_size=0,
+        _emit_source_span=True,
+    )
+    assert len(chunks) == 1  # everything in one chunk pre-split
+    assert chunks[0]["_source_span"] == {"start": 0, "end": len(merged)}
+
+    chunks = enforce_chunk_token_limit_before_embedding(chunks, tok, max_tokens=30)
+    assert len(chunks) > 1  # hard-split fired
+
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    # The early slices live entirely inside "big" -> single ref, not both blocks.
+    assert chunks[0]["sidecar"]["refs"] == [{"type": "block", "id": "big"}]
+    # "small" is referenced only by the slice that actually reaches its content.
+    small_refs = [
+        ch for ch in chunks if any(r["id"] == "small" for r in ch["sidecar"]["refs"])
+    ]
+    assert len(small_refs) == 1
+
+
+@pytest.mark.offline
+def test_hard_split_multi_sentence_rejoin_keeps_provenance(tmp_path: Path) -> None:
+    # A single block whose sentences are separated by single spaces. The hard
+    # split regroups whole sentence units and rejoins them with "\n\n", so the
+    # resulting slice content is NOT byte-verbatim in the source. Span propagation
+    # must fall back to whitespace-normalized matching instead of dropping the span
+    # — otherwise span-first backfill would wrongly FAIL the document.
+    block = "ab. cd. ef. gh. ij. kl."
+    blocks_path, merged = _write_blocks(tmp_path, [("b1", block)])
+    tok = _tokenizer()
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=200,
+        chunk_overlap_token_size=0,
+        _emit_source_span=True,
+    )
+    assert len(chunks) == 1
+    assert chunks[0]["_source_span"] == {"start": 0, "end": len(merged)}
+
+    chunks = enforce_chunk_token_limit_before_embedding(chunks, tok, max_tokens=7)
+    assert len(chunks) > 1  # hard split fired into multiple slices
+    # At least one slice rejoined sentence units with "\n\n" (not byte-verbatim),
+    # which is exactly the case the normalized span fallback must cover.
+    assert any("\n\n" in ch["content"] for ch in chunks)
+
+    # Must NOT raise: every slice keeps a span via the normalized fallback, and
+    # each maps back to the single source block it came from.
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    for ch in chunks:
+        assert [r["id"] for r in ch["sidecar"]["refs"]] == ["b1"]
+
+
+@pytest.mark.offline
+def test_real_tiktoken_token_windows_match_verbatim(tmp_path: Path) -> None:
+    # The char tokenizer guarantees decode(encode(x)) == x; a real BPE tokenizer
+    # does not split on character boundaries, so this exercises that tiktoken's
+    # decoded token windows (after the chunker's .strip()) are still locatable in
+    # the reconstructed merged text — including across the token-overlap fallback.
+    tiktoken = pytest.importorskip("tiktoken")
+    del tiktoken  # only needed to gate the test
+    from lightrag.utils import TiktokenTokenizer
+
+    tok = TiktokenTokenizer()
+    blocks = [
+        ("b1", "The quick brown fox jumps over the lazy dog repeatedly."),
+        ("b2", "Pack my box with five dozen liquor jugs, said the printer."),
+        ("b3", "How vexingly quick daft zebras jump across the meadow."),
+    ]
+    blocks_path, merged = _write_blocks(tmp_path, blocks)
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=12,
+        chunk_overlap_token_size=3,
+        _emit_source_span=True,
+    )
+    assert len(chunks) >= 2  # small window + overlap -> multiple chunks
+
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    for ch in chunks:
+        assert ch["sidecar"]["type"] == "block"
+        assert ch["sidecar"]["refs"]
+    # Union of all refs covers every block.
+    seen = {r["id"] for ch in chunks for r in ch["sidecar"]["refs"]}
+    assert seen == {b for b, _ in blocks}
+
+
+@pytest.mark.offline
+def test_real_tiktoken_multibyte_boundary_degrades_not_fails(tmp_path: Path) -> None:
+    # Regression: tiktoken is byte-level, so a 4-byte UTF-8 char (emoji / rare CJK
+    # extension) can have its bytes split across a token-window boundary. Decoding the
+    # partial window yields U+FFFD in BOTH the chunk content and its span probe, so the
+    # chunk is unlocatable by span or by text. Span-first backfill must skip provenance
+    # for the corrupt chunks while still attributing the clean ones, not FAIL the whole
+    # document.
+    pytest.importorskip("tiktoken")
+    from lightrag.utils import TiktokenTokenizer
+
+    tok = TiktokenTokenizer()
+    # Emoji are supplementary-plane (4-byte) chars that force byte-fallback tokens.
+    block = "Status update 🎉🚀 progress 😀😁😂 and more text 🔥💡✅ keep going. " * 20
+    blocks_path, merged = _write_blocks(tmp_path, [("b1", block)])
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=18,
+        chunk_overlap_token_size=4,
+        _emit_source_span=True,
+    )
+    # The window splits at least one emoji -> some chunks carry U+FFFD and lack a span.
+    assert any("�" in c["content"] for c in chunks)
+    assert any("_source_span" not in c for c in chunks)
+
+    # Must NOT raise: corrupt chunks are skipped, the rest are attributed.
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    for ch in chunks:
+        if "�" in ch["content"]:
+            assert "sidecar" not in ch  # provenance degraded, document not failed
+        elif ch["content"].strip():  # empty tail chunks are skipped entirely
+            assert ch["sidecar"]["refs"] == [{"type": "block", "id": "b1"}]
+    # At least the clean chunks resolved into the single source block.
+    assert any("sidecar" in ch for ch in chunks)
+
+
+@pytest.mark.offline
+def test_backfilled_sidecar_persists_into_chunks_dict(tmp_path: Path) -> None:
+    blocks_path, merged = _write_blocks(
+        tmp_path, [("b1", "Alpha body."), ("b2", "Beta body.")]
+    )
+    tok = _tokenizer()
+
+    chunks = chunking_by_fixed_token(
+        tok, merged, chunk_token_size=200, _emit_source_span=True
+    )
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    chunks_dict = build_chunks_dict_from_chunking_result(
+        chunks, doc_id="doc-xyz", file_path="doc.docx"
+    )
+    assert chunks_dict  # non-empty
+    for record in chunks_dict.values():
+        assert "sidecar" in record
+        assert "_source_span" not in record
+        assert record["sidecar"]["type"] == "block"
+        assert record["sidecar"]["refs"]
+
+
+@pytest.mark.offline
+def test_span_first_disambiguates_repeated_cross_block_text(tmp_path: Path) -> None:
+    # Latest pathological case: the real first chunk is b1 + separator + b2
+    # ("ab\n\ncd"), but stripping whitespace makes it textually equal to b3
+    # ("abcd"). Span-first backfill must map by source coverage, not by the
+    # ambiguous stripped string.
+    blocks_path, merged = _write_blocks(
+        tmp_path, [("b1", "ab"), ("b2", "cd"), ("b3", "abcd")]
+    )
+    tok = _tokenizer()
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        merged,
+        chunk_token_size=6,
+        chunk_overlap_token_size=0,
+        _emit_source_span=True,
+    )
+    assert [c["content"] for c in chunks] == ["ab\n\ncd", "abcd"]
+
+    backfill_chunk_sidecars(chunks, blocks_path)
+
+    assert [r["id"] for r in chunks[0]["sidecar"]["refs"]] == ["b1", "b2"]
+    assert [r["id"] for r in chunks[1]["sidecar"]["refs"]] == ["b3"]
+
+
+@pytest.mark.offline
+def test_split_by_character_only_emits_exact_source_spans() -> None:
+    tok = _tokenizer()
+    content = "  alpha|beta|gamma  "
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        content,
+        chunk_token_size=20,
+        split_by_character="|",
+        split_by_character_only=True,
+        _emit_source_span=True,
+    )
+
+    assert [c["content"] for c in chunks] == ["alpha", "beta", "gamma"]
+    assert [c["_source_span"] for c in chunks] == [
+        {"start": 2, "end": 7},
+        {"start": 8, "end": 12},
+        {"start": 13, "end": 18},
+    ]

tests/chunker/test_token_window_source_span.py

@@ -0,0 +1,201 @@
+"""Regression tests for the O(N) delta-decode in ``_token_window_source_span``.
+
+The fixed-token chunker maps each decoded token window back to its exact source
+span. The offset of a window's start used to be recomputed with a full
+``decode(tokens[:start_token])`` prefix decode on every window — O(N) per window,
+O(N²) over a document. It now decodes only the delta since the previous *verified*
+anchor (O(N) total). These tests pin three properties:
+
+1. **Equivalence** — the delta path yields byte-identical spans to a reference
+   full-prefix implementation, window for window (incl. a non-1:1 tokenizer and a
+   real tiktoken BPE tokenizer).
+2. **Exactness** — every emitted span is a verbatim slice of the source.
+3. **Linear decode budget** — total tokens handed to ``decode`` scales ~O(N), not
+   O(N²); the pre-optimization prefix decode would blow well past the bound.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from lightrag.chunker import chunking_by_fixed_token
+from lightrag.chunker.token_size import _source_span, _token_window_source_span
+from lightrag.utils import Tokenizer, TokenizerInterface
+
+
+class _CharTokenizer(TokenizerInterface):
+    """1:1 char-per-token; ``decode(encode(x)) == x`` so windows are verbatim."""
+
+    def encode(self, content: str) -> list[int]:
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens: list[int]) -> str:
+        return "".join(chr(t) for t in tokens)
+
+
+class _MultiTokenTokenizer(TokenizerInterface):
+    """Non-uniform char→token ratio: uppercase = 2 tokens, ``.?!`` = 3, else 1.
+
+    Exercises the offset arithmetic where token count != char count, so a bug that
+    confuses the two surfaces immediately.
+    """
+
+    def encode(self, content: str) -> list[int]:
+        tokens: list[int] = []
+        for ch in content:
+            if ch.isupper():
+                tokens.extend([ord(ch), ord(ch) + 1000])
+            elif ch in (".", "?", "!"):
+                tokens.extend([ord(ch), ord(ch) + 2000, ord(ch) + 3000])
+            else:
+                tokens.append(ord(ch))
+        return tokens
+
+    def decode(self, tokens: list[int]) -> str:
+        result: list[str] = []
+        i = 0
+        while i < len(tokens):
+            base = tokens[i]
+            if (
+                i + 2 < len(tokens)
+                and tokens[i + 1] == base + 2000
+                and tokens[i + 2] == base + 3000
+            ):
+                result.append(chr(base))
+                i += 3
+            elif i + 1 < len(tokens) and tokens[i + 1] == base + 1000:
+                result.append(chr(base))
+                i += 2
+            else:
+                result.append(chr(base))
+                i += 1
+        return "".join(result)
+
+
+class _CountingTokenizer(TokenizerInterface):
+    """Wraps a char tokenizer and tallies every token handed to ``decode``."""
+
+    def __init__(self) -> None:
+        self.decoded_tokens = 0
+
+    def encode(self, content: str) -> list[int]:
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens: list[int]) -> str:
+        self.decoded_tokens += len(tokens)
+        return "".join(chr(t) for t in tokens)
+
+
+def _ref_full_prefix_span(
+    tokenizer: Tokenizer,
+    content: str,
+    tokens: list[int],
+    start_token: int,
+    end_token: int,
+) -> dict[str, int] | None:
+    """The pre-optimization implementation: full ``decode(tokens[:start_token])``."""
+    window = tokenizer.decode(tokens[start_token:end_token])
+    start = len(tokenizer.decode(tokens[:start_token]))
+    end = start + len(window)
+    if content[start:end] != window:
+        found = content.find(
+            window, max(0, start - 32), min(len(content), end + 32 + len(window))
+        )
+        if found < 0:
+            return None
+        start, end = found, found + len(window)
+    return _source_span(content, start, end)
+
+
+def _windows(n_tokens: int, chunk_size: int, overlap: int) -> list[tuple[int, int]]:
+    step = chunk_size - overlap
+    return [
+        (start, min(start + chunk_size, n_tokens)) for start in range(0, n_tokens, step)
+    ]
+
+
+@pytest.mark.offline
+@pytest.mark.parametrize(
+    "content",
+    [
+        " ".join(f"word{i:03d}" for i in range(400)),
+        "Alpha. BETA gamma? Delta! " * 120,
+        "Repeated phrase. " * 200,  # heavily repeated — exact offsets must not drift
+    ],
+)
+def test_delta_decode_matches_full_prefix_reference(content: str) -> None:
+    for impl in (_CharTokenizer(), _MultiTokenTokenizer()):
+        tok = Tokenizer(model_name="t", tokenizer=impl)
+        tokens = tok.encode(content)
+        chunk_size, overlap = 60, 12
+        anchor = (0, 0)
+        for start_token, end_token in _windows(len(tokens), chunk_size, overlap):
+            got, anchor = _token_window_source_span(
+                tok, content, tokens, start_token, end_token, anchor=anchor
+            )
+            ref = _ref_full_prefix_span(tok, content, tokens, start_token, end_token)
+            assert got == ref
+
+
+@pytest.mark.offline
+def test_delta_decode_matches_full_prefix_with_tiktoken() -> None:
+    pytest.importorskip("tiktoken")
+    from lightrag.utils import TiktokenTokenizer
+
+    tok = TiktokenTokenizer()
+    content = (
+        "The quick brown fox jumps over the lazy dog. "
+        "Pack my box with five dozen liquor jugs. "
+        "How vexingly quick daft zebras jump. "
+    ) * 40
+    tokens = tok.encode(content)
+    anchor = (0, 0)
+    for start_token, end_token in _windows(len(tokens), 24, 6):
+        got, anchor = _token_window_source_span(
+            tok, content, tokens, start_token, end_token, anchor=anchor
+        )
+        ref = _ref_full_prefix_span(tok, content, tokens, start_token, end_token)
+        assert got == ref
+
+
+@pytest.mark.offline
+def test_fixed_token_spans_are_exact_on_long_doc() -> None:
+    content = " ".join(f"token{i:04d}" for i in range(1500))
+    tok = Tokenizer(model_name="char", tokenizer=_CharTokenizer())
+
+    chunks = chunking_by_fixed_token(
+        tok,
+        content,
+        chunk_token_size=120,
+        chunk_overlap_token_size=20,
+        _emit_source_span=True,
+    )
+
+    assert len(chunks) > 5
+    for chunk in chunks:
+        span = chunk["_source_span"]
+        assert content[span["start"] : span["end"]] == chunk["content"]
+
+
+@pytest.mark.offline
+def test_decode_budget_is_linear_not_quadratic() -> None:
+    # char tokenizer => len(tokens) == len(content). With the old full-prefix
+    # decode the helper alone would decode ~sum(starts) == N^2/(2*step) tokens;
+    # the delta path decodes ~one step per window plus each window once. Assert
+    # the total stays under a small linear multiple of N, a bound the quadratic
+    # implementation cannot meet.
+    content = "x" * 8000
+    counting = _CountingTokenizer()
+    tok = Tokenizer(model_name="counting", tokenizer=counting)
+    n = len(tok.encode(content))  # 8000
+
+    chunking_by_fixed_token(
+        tok,
+        content,
+        chunk_token_size=200,
+        chunk_overlap_token_size=20,
+        _emit_source_span=True,
+    )
+
+    # Empirically ~3.2*N for the delta path; the old prefix decode is ~24*N here.
+    assert counting.decoded_tokens <= 6 * n

tests/conftest.py

@@ -0,0 +1,191 @@
+"""
+Pytest configuration for LightRAG tests.
+
+This file provides command-line options and fixtures for test configuration.
+"""
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def _hermetic_mineru_env(monkeypatch):
+    """Make every test start with parser-routing env vars in their unset state.
+
+    ``lightrag/api/{auth,config}.py`` call ``load_dotenv(override=False)``
+    at import time, leaking the developer's local ``.env`` into the test
+    process. The MinerU test fixtures assume ``MINERU_API_MODE`` is unset
+    (so it defaults to ``"local"`` per ``MinerURawClient.__init__`` /
+    ``parser_engine_endpoint_requirement``):
+
+    - A leaked ``MINERU_API_MODE=offical`` typo (or any invalid value)
+      makes ``MinerURawClient()`` raise at construction.
+    - A leaked ``MINERU_API_MODE=official`` flips
+      ``parser_engine_endpoint_requirement`` to return
+      ``"MINERU_API_TOKEN"`` instead of ``"MINERU_LOCAL_ENDPOINT"``,
+      breaking the validation-error string match.
+
+    ``LIGHTRAG_PARSER`` is cleared for the same reason: a routing rule
+    like ``docx:mineru-iet`` in the developer's ``.env`` forces
+    ``parser_routing.validate_parser_routing_config`` to require the
+    corresponding endpoint (``MINERU_LOCAL_ENDPOINT`` /
+    ``DOCLING_ENDPOINT``) at every ``create_app`` call, which then trips
+    unrelated API/FastAPI tests (``test_bedrock_llm.py``,
+    ``test_path_prefixes.py``).
+
+    The ``MINERU_LOCAL_*`` parser options are stripped for the same reason:
+    a developer ``.env`` that pins e.g. ``MINERU_LOCAL_PARSE_METHOD=ocr``
+    leaks a non-default into tests that assume the built-in defaults
+    (``test_client_local_mode_round_trip`` expects ``parse_method=auto``;
+    ``test_invalid_when_local_parser_options_change`` toggles each option
+    and expects the change to invalidate a bundle recorded with defaults).
+
+    Strip these variables globally; tests that need a specific mode can
+    still ``monkeypatch.setenv(...)`` themselves and monkeypatch will
+    restore the inherited value at teardown.
+    """
+    monkeypatch.delenv("MINERU_API_MODE", raising=False)
+    monkeypatch.delenv("MINERU_API_TOKEN", raising=False)
+    monkeypatch.delenv("MINERU_LOCAL_ENDPOINT", raising=False)
+    monkeypatch.delenv("MINERU_OFFICIAL_ENDPOINT", raising=False)
+    monkeypatch.delenv("MINERU_LOCAL_BACKEND", raising=False)
+    monkeypatch.delenv("MINERU_LOCAL_PARSE_METHOD", raising=False)
+    monkeypatch.delenv("MINERU_LOCAL_IMAGE_ANALYSIS", raising=False)
+    monkeypatch.delenv("MINERU_LOCAL_START_PAGE_ID", raising=False)
+    monkeypatch.delenv("LIGHTRAG_PARSER", raising=False)
+    monkeypatch.delenv("DOCLING_ENDPOINT", raising=False)
+
+
+def pytest_configure(config):
+    """Register custom markers for LightRAG tests."""
+    config.addinivalue_line(
+        "markers", "offline: marks tests as offline (no external dependencies)"
+    )
+    config.addinivalue_line(
+        "markers",
+        "integration: marks tests requiring external services (skipped by default)",
+    )
+    config.addinivalue_line("markers", "requires_db: marks tests requiring database")
+    config.addinivalue_line(
+        "markers", "requires_api: marks tests requiring LightRAG API server"
+    )
+
+
+def pytest_addoption(parser):
+    """Add custom command-line options for LightRAG tests."""
+
+    parser.addoption(
+        "--keep-artifacts",
+        action="store_true",
+        default=False,
+        help="Keep test artifacts (temporary directories and files) after test completion for inspection",
+    )
+
+    parser.addoption(
+        "--stress-test",
+        action="store_true",
+        default=False,
+        help="Enable stress test mode with more intensive workloads",
+    )
+
+    parser.addoption(
+        "--test-workers",
+        action="store",
+        default=3,
+        type=int,
+        help="Number of parallel workers for stress tests (default: 3)",
+    )
+
+    parser.addoption(
+        "--run-integration",
+        action="store_true",
+        default=False,
+        help="Run integration tests that require external services (database, API server, etc.)",
+    )
+
+
+def pytest_collection_modifyitems(config, items):
+    """Modify test collection to skip integration tests by default.
+
+    Integration tests are skipped unless --run-integration flag is provided.
+    This allows running offline tests quickly without needing external services.
+    """
+    if config.getoption("--run-integration"):
+        # If --run-integration is specified, run all tests
+        return
+
+    skip_integration = pytest.mark.skip(
+        reason="Requires external services(DB/API), use --run-integration to run"
+    )
+
+    for item in items:
+        if "integration" in item.keywords:
+            item.add_marker(skip_integration)
+
+
+@pytest.fixture(scope="session")
+def keep_test_artifacts(request):
+    """
+    Fixture to determine whether to keep test artifacts.
+
+    Priority: CLI option > Environment variable > Default (False)
+    """
+    import os
+
+    # Check CLI option first
+    if request.config.getoption("--keep-artifacts"):
+        return True
+
+    # Fall back to environment variable
+    return os.getenv("LIGHTRAG_KEEP_ARTIFACTS", "false").lower() == "true"
+
+
+@pytest.fixture(scope="session")
+def stress_test_mode(request):
+    """
+    Fixture to determine whether stress test mode is enabled.
+
+    Priority: CLI option > Environment variable > Default (False)
+    """
+    import os
+
+    # Check CLI option first
+    if request.config.getoption("--stress-test"):
+        return True
+
+    # Fall back to environment variable
+    return os.getenv("LIGHTRAG_STRESS_TEST", "false").lower() == "true"
+
+
+@pytest.fixture(scope="session")
+def parallel_workers(request):
+    """
+    Fixture to determine the number of parallel workers for stress tests.
+
+    Priority: CLI option > Environment variable > Default (3)
+    """
+    import os
+
+    # Check CLI option first
+    cli_workers = request.config.getoption("--test-workers")
+    if cli_workers != 3:  # Non-default value provided
+        return cli_workers
+
+    # Fall back to environment variable
+    return int(os.getenv("LIGHTRAG_TEST_WORKERS", "3"))
+
+
+@pytest.fixture(scope="session")
+def run_integration_tests(request):
+    """
+    Fixture to determine whether to run integration tests.
+
+    Priority: CLI option > Environment variable > Default (False)
+    """
+    import os
+
+    # Check CLI option first
+    if request.config.getoption("--run-integration"):
+        return True
+
+    # Fall back to environment variable
+    return os.getenv("LIGHTRAG_RUN_INTEGRATION", "false").lower() == "true"

tests/evaluation/test_evaluation_offline_retrieval_check.py

@@ -0,0 +1,101 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+from lightrag.evaluation.offline_retrieval_check import (
+    audit_samples,
+    load_cases,
+    load_documents,
+    load_oracle,
+    summarize,
+)
+
+
+class OfflineRetrievalCheckTests(unittest.TestCase):
+    def test_expected_document_ranks_first(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            root = Path(temp_dir)
+            docs_dir = root / "docs"
+            docs_dir.mkdir()
+            (docs_dir / "alpha.md").write_text(
+                "Alpha covers vector search and filtering.",
+                encoding="utf-8",
+            )
+            (docs_dir / "beta.md").write_text(
+                "Beta covers deployment and monitoring.",
+                encoding="utf-8",
+            )
+            dataset = root / "dataset.json"
+            dataset.write_text(
+                '{"test_cases":[{"question":"Which file explains vector search?"}]}',
+                encoding="utf-8",
+            )
+            oracle = root / "oracle.json"
+            oracle.write_text(
+                '{"oracle":[{"question":"Which file explains vector search?",'
+                '"expected_documents":["alpha.md"]}]}',
+                encoding="utf-8",
+            )
+
+            results = audit_samples(
+                load_cases(dataset),
+                load_oracle(oracle),
+                load_documents(docs_dir),
+            )
+            summary = summarize(results, top_k=1)
+
+        self.assertEqual(results[0].ranked[0], "alpha.md")
+        self.assertEqual(summary["queries"], 1)
+        self.assertEqual(summary["average_recall_at_k"], 1.0)
+
+    def test_zero_score_documents_do_not_count_as_hits(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            root = Path(temp_dir)
+            docs_dir = root / "docs"
+            docs_dir.mkdir()
+            (docs_dir / "alpha.md").write_text(
+                "Alpha covers deployment pipelines.",
+                encoding="utf-8",
+            )
+            (docs_dir / "beta.md").write_text(
+                "Beta covers monitoring dashboards.",
+                encoding="utf-8",
+            )
+            dataset = root / "dataset.json"
+            dataset.write_text(
+                '{"test_cases":[{"question":"Which file explains vector search?"}]}',
+                encoding="utf-8",
+            )
+            oracle = root / "oracle.json"
+            oracle.write_text(
+                '{"oracle":[{"question":"Which file explains vector search?",'
+                '"expected_documents":["alpha.md"]}]}',
+                encoding="utf-8",
+            )
+
+            results = audit_samples(
+                load_cases(dataset),
+                load_oracle(oracle),
+                load_documents(docs_dir),
+            )
+            summary = summarize(results, top_k=1)
+
+        self.assertEqual(results[0].ranked, [])
+        self.assertEqual(summary["average_recall_at_k"], 0.0)
+        self.assertEqual(summary["no_hit_queries"], 1)
+
+    def test_sample_oracle_has_full_recall_at_two(self):
+        results = audit_samples(
+            load_cases(Path("lightrag/evaluation/sample_dataset.json")),
+            load_oracle(Path("lightrag/evaluation/sample_retrieval_oracle.json")),
+            load_documents(Path("lightrag/evaluation/sample_documents")),
+        )
+        summary = summarize(results, top_k=2)
+
+        self.assertEqual(summary["queries"], 6)
+        self.assertEqual(summary["full_recall_queries"], 6)
+        self.assertEqual(summary["no_hit_queries"], 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/extraction/test_extract_entities.py

@@ -0,0 +1,178 @@
+"""Tests for entity extraction gleaning token limit guard."""
+
+import logging
+from unittest.mock import AsyncMock
+
+import pytest
+
+from lightrag.utils import Tokenizer, TokenizerInterface
+
+
+@pytest.fixture
+def _propagate_lightrag_logger(monkeypatch):
+    """``lightrag.utils.logger`` sets ``propagate = False`` to avoid noisy
+    test output; restore propagation locally so ``caplog`` can capture
+    WARNING records emitted from inside ``lightrag.operate``."""
+    monkeypatch.setattr(logging.getLogger("lightrag"), "propagate", True)
+
+
+class DummyTokenizer(TokenizerInterface):
+    """Simple 1:1 character-to-token mapping for testing."""
+
+    def encode(self, content: str):
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens):
+        return "".join(chr(token) for token in tokens)
+
+
+def _make_global_config(
+    entity_extract_max_gleaning: int = 1,
+) -> dict:
+    """Build a minimal global_config dict for extract_entities."""
+    tokenizer = Tokenizer("dummy", DummyTokenizer())
+    extract_func = AsyncMock(return_value="")
+    return {
+        "llm_model_func": extract_func,
+        "role_llm_funcs": {
+            "extract": extract_func,
+            "keyword": extract_func,
+            "query": extract_func,
+            "vlm": extract_func,
+        },
+        "entity_extract_max_gleaning": entity_extract_max_gleaning,
+        "entity_extract_max_records": 100,
+        "entity_extract_max_entities": 40,
+        "addon_params": {},
+        "tokenizer": tokenizer,
+        "llm_model_max_async": 1,
+    }
+
+
+# Minimal valid extraction result that _process_extraction_result can parse
+_EXTRACTION_RESULT = (
+    "(entity<|#|>TEST_ENTITY<|#|>CONCEPT<|#|>A test entity)<|COMPLETE|>"
+)
+
+
+def _make_chunks(content: str = "Test content.") -> dict[str, dict]:
+    return {
+        "chunk-001": {
+            "tokens": len(content),
+            "content": content,
+            "full_doc_id": "doc-001",
+            "chunk_order_index": 0,
+        }
+    }
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_gleaning_skipped_when_tokens_exceed_limit(
+    monkeypatch, caplog, _propagate_lightrag_logger
+):
+    """Gleaning must be skipped (with a WARNING) when the projected
+    gleaning input — system + history(user+assistant) + continue prompt —
+    exceeds ``MAX_EXTRACT_INPUT_TOKENS``.  This prevents
+    ``context_length_exceeded`` errors from the LLM provider on the second
+    round when the initial response was long.
+    """
+    from lightrag.operate import extract_entities
+
+    # 10 tokens cannot fit any realistic prompt — guard must trip.
+    monkeypatch.setenv("MAX_EXTRACT_INPUT_TOKENS", "10")
+
+    global_config = _make_global_config(entity_extract_max_gleaning=1)
+    llm_func = global_config["llm_model_func"]
+    llm_func.return_value = _EXTRACTION_RESULT
+
+    with caplog.at_level("WARNING", logger="lightrag"):
+        await extract_entities(
+            chunks=_make_chunks(),
+            global_config=global_config,
+        )
+
+    # Only the initial extraction round ran; gleaning was skipped.
+    assert llm_func.await_count == 1
+
+    warnings_emitted = [
+        rec.getMessage()
+        for rec in caplog.records
+        if rec.levelname == "WARNING"
+        and rec.getMessage().startswith("Gleaning stopped for chunk chunk-001:")
+    ]
+    assert warnings_emitted, (
+        "expected a WARNING log explaining gleaning was skipped due to "
+        "token limit; got: "
+        f"{[r.getMessage() for r in caplog.records]}"
+    )
+    # Message must surface both the measured token count and the limit so
+    # operators can size MAX_EXTRACT_INPUT_TOKENS appropriately.
+    msg = warnings_emitted[0]
+    assert "exceeded limit (10)" in msg
+    assert "Input tokens (" in msg
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_gleaning_proceeds_when_tokens_within_limit(monkeypatch):
+    """Gleaning runs normally when the projected input fits the cap."""
+    from lightrag.operate import extract_entities
+
+    monkeypatch.setenv("MAX_EXTRACT_INPUT_TOKENS", "999999")
+
+    global_config = _make_global_config(entity_extract_max_gleaning=1)
+    llm_func = global_config["llm_model_func"]
+    llm_func.return_value = _EXTRACTION_RESULT
+
+    await extract_entities(
+        chunks=_make_chunks(),
+        global_config=global_config,
+    )
+
+    # Both rounds run: initial extraction + one gleaning pass.
+    assert llm_func.await_count == 2
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_no_gleaning_when_max_gleaning_zero(monkeypatch):
+    """``entity_extract_max_gleaning=0`` disables gleaning regardless of
+    token budget — the guard is downstream of the feature flag."""
+    from lightrag.operate import extract_entities
+
+    monkeypatch.setenv("MAX_EXTRACT_INPUT_TOKENS", "999999")
+
+    global_config = _make_global_config(entity_extract_max_gleaning=0)
+    llm_func = global_config["llm_model_func"]
+    llm_func.return_value = _EXTRACTION_RESULT
+
+    await extract_entities(
+        chunks=_make_chunks(),
+        global_config=global_config,
+    )
+
+    assert llm_func.await_count == 1
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_gleaning_guard_disabled_when_max_tokens_zero(monkeypatch):
+    """Setting ``MAX_EXTRACT_INPUT_TOKENS=0`` opts out of the guard so
+    gleaning always runs regardless of input size — useful for callers
+    whose provider has no hard input ceiling."""
+    from lightrag.operate import extract_entities
+
+    monkeypatch.setenv("MAX_EXTRACT_INPUT_TOKENS", "0")
+
+    global_config = _make_global_config(entity_extract_max_gleaning=1)
+    llm_func = global_config["llm_model_func"]
+    llm_func.return_value = _EXTRACTION_RESULT
+
+    await extract_entities(
+        chunks=_make_chunks(),
+        global_config=global_config,
+    )
+
+    # Guard disabled → gleaning still runs even with tight projected input.
+    assert llm_func.await_count == 2

tests/extraction/test_keyword_extraction_drivers.py

@@ -0,0 +1,270 @@
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from lightrag.llm.lmdeploy import lmdeploy_model_if_cache
+from lightrag.llm.lollms import lollms_model_complete, lollms_model_if_cache
+from lightrag.llm.ollama import _ollama_model_if_cache, ollama_model_complete
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_ollama_response_format_forwards_to_inner():
+    hashing_kv = SimpleNamespace(global_config={"llm_model_name": "ollama-model"})
+
+    with patch(
+        "lightrag.llm.ollama._ollama_model_if_cache",
+        AsyncMock(return_value="{}"),
+    ) as mocked_complete:
+        await ollama_model_complete(
+            prompt="hello",
+            hashing_kv=hashing_kv,
+            response_format={"type": "json_object"},
+        )
+
+    assert mocked_complete.await_args.kwargs["response_format"] == {
+        "type": "json_object"
+    }
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_ollama_legacy_keyword_extraction_emits_deprecation_warning():
+    """_ollama_model_if_cache is the canonical emission site for the shim."""
+    captured_kwargs = {}
+
+    class FakeAsyncClient:
+        def __init__(self, *args, **kwargs):
+            self._client = SimpleNamespace(aclose=AsyncMock())
+
+        async def chat(self, **kwargs):
+            captured_kwargs.update(kwargs)
+            return {"message": {"content": "{}"}}
+
+    with patch("lightrag.llm.ollama.ollama.AsyncClient", FakeAsyncClient):
+        with pytest.warns(DeprecationWarning):
+            await _ollama_model_if_cache(
+                model="ollama-model",
+                prompt="hello",
+                keyword_extraction=True,
+            )
+
+    assert captured_kwargs["format"] == "json"
+    assert "keyword_extraction" not in captured_kwargs
+    assert "response_format" not in captured_kwargs
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_ollama_complete_forwards_legacy_flag_downstream():
+    """ollama_model_complete is a pure forwarder; the shim fires inside _if_cache."""
+    hashing_kv = SimpleNamespace(global_config={"llm_model_name": "ollama-model"})
+
+    with patch(
+        "lightrag.llm.ollama._ollama_model_if_cache",
+        AsyncMock(return_value="{}"),
+    ) as mocked_complete:
+        await ollama_model_complete(
+            prompt="hello",
+            hashing_kv=hashing_kv,
+            keyword_extraction=True,
+        )
+
+    assert mocked_complete.await_args.kwargs.get("keyword_extraction") is True
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_ollama_translates_json_object_response_format_to_native_format():
+    captured_kwargs = {}
+
+    class FakeAsyncClient:
+        def __init__(self, *args, **kwargs):
+            self._client = SimpleNamespace(aclose=AsyncMock())
+
+        async def chat(self, **kwargs):
+            captured_kwargs.update(kwargs)
+            return {"message": {"content": "{}"}}
+
+    with patch("lightrag.llm.ollama.ollama.AsyncClient", FakeAsyncClient):
+        result = await _ollama_model_if_cache(
+            model="ollama-model",
+            prompt="hello",
+            response_format={"type": "json_object"},
+        )
+
+    assert result == "{}"
+    assert captured_kwargs["format"] == "json"
+    assert "response_format" not in captured_kwargs
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_ollama_unwraps_openai_json_schema_response_format():
+    captured_kwargs = {}
+    schema = {
+        "type": "object",
+        "properties": {"answer": {"type": "string"}},
+        "required": ["answer"],
+    }
+
+    class FakeAsyncClient:
+        def __init__(self, *args, **kwargs):
+            self._client = SimpleNamespace(aclose=AsyncMock())
+
+        async def chat(self, **kwargs):
+            captured_kwargs.update(kwargs)
+            return {"message": {"content": "{}"}}
+
+    with patch("lightrag.llm.ollama.ollama.AsyncClient", FakeAsyncClient):
+        result = await _ollama_model_if_cache(
+            model="ollama-model",
+            prompt="hello",
+            response_format={
+                "type": "json_schema",
+                "json_schema": {"name": "answer_payload", "schema": schema},
+            },
+        )
+
+    assert result == "{}"
+    assert captured_kwargs["format"] == schema
+    assert "response_format" not in captured_kwargs
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_lollms_if_cache_strips_response_format_before_request():
+    """lollms_model_if_cache drops response_format; lollms has no JSON mode."""
+    captured_requests = []
+
+    class FakeResponse:
+        def __init__(self):
+            pass
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *exc_info):
+            return False
+
+        async def text(self):
+            return "{}"
+
+    class FakeSession:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *exc_info):
+            return False
+
+        def post(self, url, json):
+            captured_requests.append(json)
+            return FakeResponse()
+
+    with patch("lightrag.llm.lollms.aiohttp.ClientSession", FakeSession):
+        result = await lollms_model_if_cache(
+            model="lollms-model",
+            prompt="hello",
+            response_format={"type": "json_object"},
+        )
+
+    assert result == "{}"
+    assert captured_requests
+    assert "response_format" not in captured_requests[0]
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_lollms_if_cache_emits_deprecation_warning():
+    class FakeResponse:
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *exc_info):
+            return False
+
+        async def text(self):
+            return "{}"
+
+    class FakeSession:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *exc_info):
+            return False
+
+        def post(self, url, json):
+            return FakeResponse()
+
+    with patch("lightrag.llm.lollms.aiohttp.ClientSession", FakeSession):
+        with pytest.warns(DeprecationWarning):
+            await lollms_model_if_cache(
+                model="lollms-model",
+                prompt="hello",
+                keyword_extraction=True,
+            )
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_lollms_complete_forwards_legacy_flag_downstream():
+    hashing_kv = SimpleNamespace(global_config={"llm_model_name": "lollms-model"})
+
+    with patch(
+        "lightrag.llm.lollms.lollms_model_if_cache",
+        AsyncMock(return_value="{}"),
+    ) as mocked_complete:
+        await lollms_model_complete(
+            prompt="hello",
+            hashing_kv=hashing_kv,
+            keyword_extraction=True,
+        )
+
+    assert mocked_complete.await_args.kwargs.get("keyword_extraction") is True
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_lmdeploy_strips_response_format_before_generation_config(monkeypatch):
+    captured_gen_config_kwargs = {}
+
+    class FakeGenerationConfig:
+        def __init__(self, **kwargs):
+            captured_gen_config_kwargs.update(kwargs)
+
+    class FakeVersion:
+        def __lt__(self, other):
+            return False
+
+    async def fake_generate(*_args, **_kwargs):
+        yield SimpleNamespace(response="{}")
+
+    monkeypatch.setattr(
+        "lightrag.llm.lmdeploy.initialize_lmdeploy_pipeline",
+        lambda **_kwargs: SimpleNamespace(generate=fake_generate),
+    )
+
+    import sys
+
+    sys.modules["lmdeploy"] = SimpleNamespace(
+        __version__="0.6.0",
+        version_info=FakeVersion(),
+        GenerationConfig=FakeGenerationConfig,
+    )
+
+    result = await lmdeploy_model_if_cache(
+        model="lmdeploy-model",
+        prompt="hello",
+        response_format={"type": "json_object"},
+    )
+
+    assert result == "{}"
+    assert "response_format" not in captured_gen_config_kwargs
+    assert "keyword_extraction" not in captured_gen_config_kwargs

tests/extraction/test_keyword_parsing.py

@@ -0,0 +1,155 @@
+import pytest
+from unittest.mock import patch
+
+from lightrag.base import QueryParam
+from lightrag.operate import _parse_keywords_payload, extract_keywords_only
+
+
+class _FakeKeywordModel:
+    def model_dump(self):
+        return {
+            "high_level_keywords": ["AI"],
+            "low_level_keywords": ["RAG", "Graph"],
+        }
+
+
+class _FakeTokenizer:
+    def encode(self, content: str) -> list[int]:
+        return [ord(ch) for ch in content]
+
+
+class _FakeKVStorage:
+    def __init__(self):
+        self.global_config = {"enable_llm_cache": True}
+        self._store = {}
+
+    async def get_by_id(self, key):
+        return self._store.get(key)
+
+    async def upsert(self, entries):
+        self._store.update(entries)
+
+
+def _keyword_global_config(
+    model: str, binding: str = "openai", keyword_func=None
+) -> dict:
+    return {
+        "addon_params": {"language": "en"},
+        "tokenizer": _FakeTokenizer(),
+        "role_llm_funcs": {"keyword": keyword_func} if keyword_func else {},
+        "llm_cache_identities": {
+            "keyword": {
+                "role": "keyword",
+                "binding": binding,
+                "model": model,
+                "host": "https://api.example.com/v1",
+            }
+        },
+    }
+
+
+@pytest.mark.offline
+def test_parse_keywords_payload_accepts_model_like_objects():
+    is_valid, hl_keywords, ll_keywords = _parse_keywords_payload(_FakeKeywordModel())
+
+    assert is_valid is True
+    assert hl_keywords == ["AI"]
+    assert ll_keywords == ["RAG", "Graph"]
+
+
+@pytest.mark.offline
+def test_parse_keywords_payload_extracts_json_from_wrapped_text():
+    result = """
+    analysis first
+    {"high_level_keywords":"AI, Agents","low_level_keywords":["RAG","LightRAG"]}
+    trailing note
+    """
+
+    is_valid, hl_keywords, ll_keywords = _parse_keywords_payload(result)
+
+    assert is_valid is True
+    assert hl_keywords == ["AI", "Agents"]
+    assert ll_keywords == ["RAG", "LightRAG"]
+
+
+@pytest.mark.offline
+def test_parse_keywords_payload_warns_when_json_repair_is_used():
+    broken_result = (
+        '{"high_level_keywords":"AI, Agents","low_level_keywords":["RAG","LightRAG"]'
+    )
+
+    with patch("lightrag.operate.logger.warning") as mocked_warning:
+        is_valid, hl_keywords, ll_keywords = _parse_keywords_payload(broken_result)
+
+    assert is_valid is True
+    assert hl_keywords == ["AI", "Agents"]
+    assert ll_keywords == ["RAG", "LightRAG"]
+    mocked_warning.assert_called_once()
+    assert (
+        "Keyword extraction response required JSON repair"
+        in mocked_warning.call_args[0][0]
+    )
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_extract_keywords_only_accepts_empty_keyword_cache_without_requery():
+    async def should_not_run(*_args, **_kwargs):
+        raise AssertionError(
+            "keyword LLM should not be called on a valid empty cache hit"
+        )
+
+    param = QueryParam()
+    global_config = _keyword_global_config("model-a", keyword_func=should_not_run)
+
+    with patch(
+        "lightrag.operate.handle_cache",
+        return_value=('{"high_level_keywords":[],"low_level_keywords":[]}', None),
+    ):
+        hl_keywords, ll_keywords = await extract_keywords_only(
+            "hello",
+            param,
+            global_config,
+            hashing_kv=None,
+        )
+
+    assert hl_keywords == []
+    assert ll_keywords == []
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_extract_keywords_only_partitions_cache_by_keyword_llm_identity():
+    cache = _FakeKVStorage()
+    calls = 0
+
+    async def keyword_model(*_args, **_kwargs):
+        nonlocal calls
+        calls += 1
+        return (
+            '{"high_level_keywords":["model-'
+            + str(calls)
+            + '"],"low_level_keywords":["rag"]}'
+        )
+
+    param = QueryParam()
+
+    first_hl, first_ll = await extract_keywords_only(
+        "same query",
+        param,
+        _keyword_global_config("model-a", keyword_func=keyword_model),
+        hashing_kv=cache,
+    )
+    second_hl, second_ll = await extract_keywords_only(
+        "same query",
+        param,
+        _keyword_global_config("model-b", keyword_func=keyword_model),
+        hashing_kv=cache,
+    )
+
+    assert first_hl == ["model-1"]
+    assert first_ll == ["rag"]
+    assert second_hl == ["model-2"]
+    assert second_ll == ["rag"]
+    assert calls == 2
+    assert len(cache._store) == 2

tests/extraction/test_keyword_prompt_template.py

@@ -0,0 +1,69 @@
+import json
+import re
+
+import pytest
+
+from lightrag.prompt import PROMPTS
+
+
+@pytest.mark.offline
+def test_keywords_extraction_prompt_template_formats_with_literal_json_braces():
+    rendered = PROMPTS["keywords_extraction"].format(
+        query="hello",
+        examples="example",
+        language="en",
+    )
+
+    assert "first character of your response must be `{`" in rendered
+    assert '{"high_level_keywords": [], "low_level_keywords": []}' in rendered
+
+
+@pytest.mark.offline
+def test_keywords_extraction_examples_are_format_only():
+    """Keyword examples must be placeholder-only JSON templates, not sample demos.
+
+    Rather than denylisting specific sample queries (brittle: generic words like
+    "education" would both false-match unrelated content and let new samples slip
+    through), assert the structural shape: no ``Query:``/``Output:`` demo framing,
+    and every keyword is an angle-bracket placeholder.
+    """
+    placeholder = re.compile(r"<[^<>]+>")
+
+    for example in PROMPTS["keywords_extraction_examples"]:
+        assert "Query:" not in example
+        assert "Output:" not in example
+
+        parsed = json.loads(example)
+        assert set(parsed) == {"high_level_keywords", "low_level_keywords"}
+        keywords = parsed["high_level_keywords"] + parsed["low_level_keywords"]
+        assert keywords
+        for keyword in keywords:
+            assert placeholder.fullmatch(keyword), keyword
+
+
+@pytest.mark.offline
+def test_keywords_extraction_prompt_labels_template_as_not_source_text():
+    prompt = PROMPTS["keywords_extraction"]
+
+    assert "---Output Format Template---" in prompt
+    assert "---Examples---" not in prompt
+    assert "Apple Inc." not in prompt
+    assert "output JSON format template only" in prompt
+    assert "not source text" in prompt
+    assert "must never be used as keyword extraction content" in prompt
+    assert (
+        "derived only from the `User Query` in the `---Real Data---` section" in prompt
+    )
+
+
+@pytest.mark.offline
+def test_keywords_extraction_prompt_keeps_single_real_user_query_section():
+    rendered = PROMPTS["keywords_extraction"].format(
+        query="How did LightRAG improve retrieval?",
+        examples="\n".join(PROMPTS["keywords_extraction_examples"]),
+        language="English",
+    )
+
+    assert rendered.count("User Query:") == 1
+    assert "User Query: How did LightRAG improve retrieval?" in rendered
+    assert "User Query:" not in "\n".join(PROMPTS["keywords_extraction_examples"])

tests/extraction/test_summary_map_reduce_perf.py

@@ -0,0 +1,115 @@
+"""Regression tests for the map-reduce chunking path of
+``_handle_entity_relation_summary``.
+
+The map phase tokenizes every description to decide (a) whether the list
+needs splitting and (b) where to cut the chunks. Historically the same
+description was encoded twice per map-reduce iteration — once to sum the
+total token count and again while building the chunks — which doubled the
+tiktoken BPE cost on every entity/edge merge. The counts from the first
+pass are now memoized and reused by the second.
+
+These tests guard:
+
+* the chunking output is unchanged (behavioral), and
+* each description is encoded exactly once per map-reduce iteration
+  (the perf property the fix introduced).
+"""
+
+from unittest.mock import AsyncMock
+
+import pytest
+
+from lightrag.utils import Tokenizer, TokenizerInterface
+
+
+class _CountingTokenizer(TokenizerInterface):
+    """1:1 character-to-token mapping that records every ``encode`` input."""
+
+    def __init__(self) -> None:
+        self.encoded: list[str] = []
+
+    def encode(self, content: str):
+        self.encoded.append(content)
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens):
+        return "".join(chr(token) for token in tokens)
+
+
+def _make_chunking_config(mock_extract) -> tuple[dict, _CountingTokenizer]:
+    """``global_config`` whose ``summary_context_size`` forces the map-reduce
+    chunking path (4 five-token descriptions vs. a 10-token window)."""
+    counter = _CountingTokenizer()
+    tokenizer = Tokenizer("dummy", counter)
+    global_config = {
+        "role_llm_funcs": {"extract": mock_extract},
+        "addon_params": {},
+        "_resolved_summary_language": "English",
+        "summary_length_recommended": 100,
+        # 4 descriptions x 5 tokens = 20 tokens > 10-token window → chunking
+        "summary_context_size": 10,
+        "summary_max_tokens": 100_000,
+        "force_llm_summary_on_merge": 10,
+        "tokenizer": tokenizer,
+    }
+    return global_config, counter
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_map_reduce_chunking_is_preserved():
+    """Forcing the chunking path must still split the descriptions into groups,
+    summarize each via the LLM, and join the results — i.e. memoizing the
+    token counts must not change the chunk boundaries or the final output."""
+    from lightrag.operate import _handle_entity_relation_summary
+
+    # Distinct summaries per reduce call so the joined result reflects the
+    # number of chunks produced.
+    mock_extract = AsyncMock(side_effect=["SUM_ONE", "SUM_TWO", "SUM_THREE"])
+    global_config, _ = _make_chunking_config(mock_extract)
+
+    description, llm_was_used = await _handle_entity_relation_summary(
+        "Entity",
+        "TEST_ENTITY",
+        ["alpha", "beta", "gamma", "delta"],
+        "<SEP>",
+        global_config,
+        llm_response_cache=None,
+    )
+
+    # Reduce phase ran (the LLM was used to summarize at least one chunk).
+    assert llm_was_used is True
+    # 4 five-token descriptions under a 10-token window → 2 chunks → 2 summaries
+    # joined in the next iteration (len == 2 takes the no-LLM join exit).
+    assert description == "SUM_ONE<SEP>SUM_TWO"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_map_phase_encodes_each_description_once():
+    """The map phase must tokenize each description exactly once per
+    map-reduce iteration. Pre-fix the chunk-building pass re-encoded every
+    description, so each of the 4 inputs appeared 2x in the first iteration;
+    post-fix it appears exactly once (the second pass reuses the counts)."""
+    from lightrag.operate import _handle_entity_relation_summary
+
+    mock_extract = AsyncMock(side_effect=["SUM_ONE", "SUM_TWO", "SUM_THREE"])
+    global_config, counter = _make_chunking_config(mock_extract)
+    inputs = ["alpha", "beta", "gamma", "delta"]
+
+    await _handle_entity_relation_summary(
+        "Entity",
+        "TEST_ENTITY",
+        inputs,
+        "<SEP>",
+        global_config,
+        llm_response_cache=None,
+    )
+
+    # In the first (chunking) iteration each input description must be encoded
+    # exactly once by the map phase. Pre-fix this was 2 per description.
+    for desc in inputs:
+        assert counter.encoded.count(desc) == 1, (
+            f"description {desc!r} was encoded {counter.encoded.count(desc)} "
+            f"time(s) by the map phase; expected exactly 1 (memoized counts)"
+        )

tests/extraction/test_summary_sanitization.py

@@ -0,0 +1,160 @@
+"""Regression tests: every merge-stage description outcome must be
+sanitized before it lands on graph nodes/edges.
+
+Extraction-time descriptions are cleaned by
+``sanitize_and_normalize_extracted_text``, but merge-stage descriptions
+can bypass that: LLM re-summaries, descriptions read back from existing
+graph nodes, and multimodal entity descriptions injected straight from
+chunk content. Any of them carrying control characters / surrogates
+(e.g. ``\\frac`` decoded as ``\\x0c`` + ``rac`` from unescaped LaTeX in
+VLM JSON) would break GraphML (XML) serialization on the next NetworkX
+flush with::
+
+    ValueError: All strings must be XML compatible: Unicode or ASCII,
+    no NULL bytes or control characters
+
+``_handle_entity_relation_summary`` (single / join outcomes) and
+``_summarize_descriptions`` (LLM outcome) now sanitize at every exit.
+"""
+
+from unittest.mock import AsyncMock
+
+import networkx as nx
+import pytest
+
+from lightrag.utils import Tokenizer, TokenizerInterface
+
+
+class DummyTokenizer(TokenizerInterface):
+    """Simple 1:1 character-to-token mapping for testing."""
+
+    def encode(self, content: str):
+        return [ord(ch) for ch in content]
+
+    def decode(self, tokens):
+        return "".join(chr(token) for token in tokens)
+
+
+def _make_global_config(summary_return: str) -> dict:
+    """Minimal global_config for ``_summarize_descriptions`` whose ``extract``
+    role LLM returns ``summary_return`` verbatim."""
+    tokenizer = Tokenizer("dummy", DummyTokenizer())
+    extract_func = AsyncMock(return_value=summary_return)
+    return {
+        "role_llm_funcs": {"extract": extract_func},
+        "addon_params": {},
+        "_resolved_summary_language": "English",
+        "summary_length_recommended": 100,
+        "summary_context_size": 10_000,
+        "tokenizer": tokenizer,
+    }
+
+
+# A description that the LLM "echoed" from a dirty source: NULL byte, bell,
+# unit-separator and DEL are all illegal in XML 1.0; \t and \n are legal and
+# must be preserved.
+_DIRTY_SUMMARY = "Clean text\x00with\x07control\x1fchars\x7f\tand\nnewlines"
+_EXPECTED_CLEAN = "Clean textwithcontrolchars\tand\nnewlines"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_summarize_descriptions_strips_control_chars():
+    """The LLM summary path must remove XML-incompatible control characters
+    while preserving legal whitespace."""
+    from lightrag.operate import _summarize_descriptions
+
+    global_config = _make_global_config(_DIRTY_SUMMARY)
+
+    summary = await _summarize_descriptions(
+        "Entity",
+        "TEST_ENTITY",
+        ["first description", "second description"],
+        global_config,
+        llm_response_cache=None,
+    )
+
+    assert summary == _EXPECTED_CLEAN
+    assert "\x00" not in summary
+    assert "\x07" not in summary
+    assert "\x1f" not in summary
+    assert "\x7f" not in summary
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_single_description_early_return_is_sanitized():
+    """The ``len(description_list) == 1`` early return skips the LLM
+    entirely — this is exactly the path a dirty multimodal entity
+    description (chunk content reused verbatim) takes on first insert."""
+    from lightrag.operate import _handle_entity_relation_summary
+
+    description, llm_was_used = await _handle_entity_relation_summary(
+        "Entity",
+        "TEST_ENTITY",
+        [_DIRTY_SUMMARY],
+        "<SEP>",
+        # Early return happens before any config key is read.
+        global_config={},
+        llm_response_cache=None,
+    )
+
+    assert llm_was_used is False
+    assert description == _EXPECTED_CLEAN
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_join_without_llm_is_sanitized():
+    """The no-LLM join outcome must also sanitize: descriptions read back
+    from pre-existing (dirty) graph nodes re-enter the merge here."""
+    from lightrag.operate import _handle_entity_relation_summary
+
+    tokenizer = Tokenizer("dummy", DummyTokenizer())
+    global_config = {
+        "tokenizer": tokenizer,
+        "summary_context_size": 100_000,
+        "summary_max_tokens": 100_000,
+        "force_llm_summary_on_merge": 10,
+    }
+
+    description, llm_was_used = await _handle_entity_relation_summary(
+        "Entity",
+        "TEST_ENTITY",
+        [_DIRTY_SUMMARY, "clean second description"],
+        "<SEP>",
+        global_config,
+        llm_response_cache=None,
+    )
+
+    assert llm_was_used is False
+    assert description == f"{_EXPECTED_CLEAN}<SEP>clean second description"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_summarized_description_is_graphml_writable(tmp_path):
+    """End-to-end guard: a node whose description came from the summary path
+    must serialize to GraphML without the original ValueError."""
+    from lightrag.operate import _summarize_descriptions
+
+    global_config = _make_global_config(_DIRTY_SUMMARY)
+
+    description = await _summarize_descriptions(
+        "Entity",
+        "TEST_ENTITY",
+        ["first description", "second description"],
+        global_config,
+        llm_response_cache=None,
+    )
+
+    graph = nx.Graph()
+    graph.add_node("TEST_ENTITY", description=description)
+
+    # Pre-fix, write_graphml raised:
+    # "All strings must be XML compatible: ... no NULL bytes or control characters"
+    out = tmp_path / "graph.graphml"
+    nx.write_graphml(graph, out)
+
+    reloaded = nx.read_graphml(out)
+    assert reloaded.nodes["TEST_ENTITY"]["description"] == _EXPECTED_CLEAN

tests/kg/faiss_impl/test_atomic_write_faiss.py

@@ -0,0 +1,127 @@
+"""Atomicity coverage for ``FaissVectorDBStorage._save_faiss_index``.
+
+The save path produces two files (``.index`` and ``.meta.json``). Each one
+must land via tmp + rename so a crash during either write preserves the
+prior snapshot. Cross-file consistency between the two renames is
+intentionally out of scope (declared in the PR).
+"""
+
+import json
+import os
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+faiss = pytest.importorskip("faiss")  # noqa: F841 — needed before the import below
+
+from lightrag.kg.faiss_impl import FaissVectorDBStorage  # noqa: E402
+from lightrag.utils import EmbeddingFunc  # noqa: E402
+
+
+def _make_storage(tmp_path, namespace: str = "vectors") -> FaissVectorDBStorage:
+    """Construct a FaissVectorDBStorage that does not need real embeddings.
+
+    ``_save_faiss_index`` only reads ``self._index`` and ``self._id_to_meta``,
+    so a dummy ``EmbeddingFunc`` with the right ``embedding_dim`` is enough.
+    """
+
+    def _unused(*_args, **_kwargs):  # pragma: no cover — never called here
+        raise AssertionError("embedding_func must not be invoked by save path")
+
+    embedding_func = EmbeddingFunc(embedding_dim=4, func=_unused)
+    return FaissVectorDBStorage(
+        namespace=namespace,
+        workspace="",
+        global_config={
+            "working_dir": str(tmp_path),
+            "embedding_batch_num": 1,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=embedding_func,
+    )
+
+
+def _seed(storage: FaissVectorDBStorage, marker: str) -> None:
+    """Push a single vector tagged with ``marker`` into the in-memory state."""
+    vec = np.ones((1, 4), dtype=np.float32)
+    storage._index.add(vec)
+    storage._id_to_meta = {
+        storage._index.ntotal - 1: {"__id__": marker, "content": marker}
+    }
+
+
+@pytest.mark.offline
+def test_save_faiss_index_publishes_both_files_cleanly(tmp_path):
+    storage = _make_storage(tmp_path)
+    _seed(storage, "v1")
+    storage._save_faiss_index()
+
+    assert os.path.exists(storage._faiss_index_file)
+    assert os.path.exists(storage._meta_file)
+    meta = json.load(open(storage._meta_file))
+    assert next(iter(meta.values()))["__id__"] == "v1"
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"Unexpected tmp residue: {leftovers}"
+
+
+@pytest.mark.offline
+def test_save_faiss_index_replace_crash_preserves_prior_index(tmp_path):
+    """If ``os.replace`` raises while renaming the ``.index`` tmp, the old
+    ``.index`` must remain loadable by ``faiss.read_index``."""
+    storage = _make_storage(tmp_path)
+    _seed(storage, "v1")
+    storage._save_faiss_index()
+    assert os.path.exists(storage._faiss_index_file)
+
+    # Bump in-memory state to v2 and then crash the .index rename.
+    storage._index.reset()
+    _seed(storage, "v2")
+    with patch(
+        "lightrag.file_atomic.os.replace",
+        side_effect=OSError("simulated crash"),
+    ):
+        with pytest.raises(OSError, match="simulated crash"):
+            storage._save_faiss_index()
+
+    # Reload the destination — must still be the v1 single-vector index.
+    reloaded = faiss.read_index(storage._faiss_index_file)
+    assert reloaded.ntotal == 1
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"Python-exception path must clean tmp, got {leftovers}"
+
+
+@pytest.mark.offline
+def test_save_faiss_meta_write_failure_preserves_prior_meta(tmp_path):
+    """A failure inside the meta ``write_fn`` (after the index has been
+    written) must leave the previous ``.meta.json`` intact."""
+    storage = _make_storage(tmp_path)
+    _seed(storage, "v1")
+    storage._save_faiss_index()
+    assert json.load(open(storage._meta_file))
+
+    real_dump = json.dump
+    seen: list[bool] = []
+
+    def explode_on_second_dump(*args, **kwargs):
+        # The first dump is from the v1 save above — we are past it because
+        # this patch is only installed for the v2 attempt. Raise immediately
+        # to simulate a serialization failure mid-write.
+        seen.append(True)
+        raise RuntimeError("simulated meta failure")
+
+    storage._index.reset()
+    _seed(storage, "v2")
+    with patch("lightrag.kg.faiss_impl.json.dump", side_effect=explode_on_second_dump):
+        with pytest.raises(RuntimeError, match="simulated meta failure"):
+            storage._save_faiss_index()
+    assert seen, "patched json.dump must have been invoked"
+
+    # .meta.json must still parse and still reflect v1.
+    meta = json.load(open(storage._meta_file))
+    assert any(entry["__id__"] == "v1" for entry in meta.values())
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"meta-write failure must clean tmp, got {leftovers}"
+
+    # Restore real json.dump so subsequent tests don't see the patch (defensive).
+    assert json.dump is real_dump

tests/kg/faiss_impl/test_faiss_deferred_embedding.py

@@ -0,0 +1,824 @@
+"""Deferred-embedding coverage for ``FaissVectorDBStorage``.
+
+The storage no longer embeds eagerly in ``upsert``: it buffers a pending doc
+and embeds once per id at flush time (``index_done_callback`` / ``finalize``).
+These tests pin that contract using a counting mock embedding function — no
+live model or network. They mirror the protocol proven for
+``NanoVectorDBStorage`` (issue #2785) plus three Faiss-specific cases:
+
+- ``test_reupsert_after_flush_replaces_single_fid`` — Faiss has no in-place
+  upsert; verify the rebuild keeps a single fid per custom id.
+- ``test_index_done_callback_save_failure_raises`` — flush succeeds, save IO
+  fails: pending is empty, ``_index_dirty`` stays True, the materialized index
+  is preserved for a finalize retry.
+- ``test_reload_warns_on_index_meta_skew`` — ``index > meta`` on-disk skew
+  (from a crash between the two atomic_writes) is logged on reload but **not**
+  auto-repaired.
+"""
+
+import json
+import os
+
+import numpy as np
+import pytest
+
+faiss = pytest.importorskip("faiss")
+
+from lightrag.kg.faiss_impl import FaissVectorDBStorage  # noqa: E402
+from lightrag.kg.shared_storage import (  # noqa: E402
+    initialize_share_data,
+    finalize_share_data,
+)
+from lightrag.utils import EmbeddingFunc  # noqa: E402
+
+DIM = 8
+
+
+@pytest.fixture(autouse=True)
+def _shared_data():
+    finalize_share_data()
+    initialize_share_data()
+    yield
+    finalize_share_data()
+
+
+class _CountingEmbed:
+    """Async embedding callable that records how many texts it embedded and how
+    many times it was invoked (one invocation == one batch)."""
+
+    def __init__(self, dim: int = DIM):
+        self.dim = dim
+        self.call_count = 0
+        self.embedded_texts: list[str] = []
+
+    async def __call__(self, texts, **kwargs):
+        self.call_count += 1
+        self.embedded_texts.extend(texts)
+        # Deterministic per-text vector so duplicates are still 1-1.
+        return np.array(
+            [
+                np.full(self.dim, (abs(hash(t)) % 97) + 1, dtype=np.float32)
+                for t in texts
+            ]
+        )
+
+
+def _make_storage(tmp_path, embed: _CountingEmbed) -> FaissVectorDBStorage:
+    return FaissVectorDBStorage(
+        namespace="test_vectors",
+        workspace="ws",
+        global_config={
+            "working_dir": str(tmp_path),
+            "embedding_batch_num": 32,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=EmbeddingFunc(embedding_dim=DIM, max_token_size=512, func=embed),
+        meta_fields={"content"},
+    )
+
+
+def _assert_consistent(storage: FaissVectorDBStorage) -> None:
+    """Faiss has two structures (index + meta dict); the root failure mode is
+    them diverging. Every test that mutates state asserts they match."""
+    assert storage._index.ntotal == len(storage._id_to_meta), (
+        f"index ntotal ({storage._index.ntotal}) != meta length "
+        f"({len(storage._id_to_meta)})"
+    )
+
+
+# ---------------------------------------------------------------------------
+# (A) Nano-ported tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_upsert_defers_embedding_to_index_done_callback(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert(
+        {
+            "id1": {"content": "alpha"},
+            "id2": {"content": "beta"},
+        }
+    )
+    assert embed.call_count == 0, "upsert must not embed"
+    assert storage._index.ntotal == 0, "nothing should be materialized yet"
+    _assert_consistent(storage)
+
+    await storage.index_done_callback()
+    assert embed.call_count == 1, "flush should embed in a single batch"
+    assert sorted(embed.embedded_texts) == ["alpha", "beta"]
+    assert storage._index.ntotal == 2
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_repeated_upserts_same_id_embed_once_per_flush(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "v1"}})
+    await storage.upsert({"id1": {"content": "v2"}})
+    await storage.upsert({"id1": {"content": "v3"}})
+
+    await storage.index_done_callback()
+
+    assert embed.call_count == 1
+    assert embed.embedded_texts == ["v3"], "only the latest content is embedded"
+    assert storage._index.ntotal == 1
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_get_vectors_caches_and_flush_reuses(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    vecs = await storage.get_vectors_by_ids(["id1"])
+    assert "id1" in vecs and len(vecs["id1"]) == DIM
+    assert embed.call_count == 1, "get_vectors_by_ids embeds pending lazily"
+
+    # Flush must reuse the cached vector, not re-embed.
+    await storage.index_done_callback()
+    assert embed.call_count == 1, "flush should reuse the cached temp vector"
+    assert storage._index.ntotal == 1
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_reupsert_after_get_vectors_clears_cached_vector(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "old"}})
+    await storage.get_vectors_by_ids(["id1"])  # caches a temp vector for "old"
+    assert embed.call_count == 1
+
+    # New content version must clear the cached vector and re-embed at flush.
+    await storage.upsert({"id1": {"content": "new"}})
+    await storage.index_done_callback()
+
+    assert embed.call_count == 2
+    assert embed.embedded_texts == ["old", "new"]
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_delete_cancels_pending_and_removes_materialized(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    # Materialize id1; leave id2 only as a pending (unflushed) upsert.
+    await storage.upsert({"id1": {"content": "alpha"}})
+    await storage.index_done_callback()
+    await storage.upsert({"id2": {"content": "beta"}})
+
+    await storage.delete(["id1", "id2"])
+
+    assert "id2" not in storage._pending_upserts, "delete cancels pending upsert"
+    assert storage._index.ntotal == 0, "delete removes the materialized row"
+    assert await storage.get_by_id("id1") is None
+    assert await storage.get_by_id("id2") is None
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_stale_client_reload_still_flushes_pending_upsert(tmp_path):
+    embed = _CountingEmbed()
+    writer = _make_storage(tmp_path, embed)
+    stale_writer = _make_storage(tmp_path, embed)
+    await writer.initialize()
+    await stale_writer.initialize()
+
+    await writer.upsert({"id1": {"content": "alpha"}})
+    assert await writer.index_done_callback() is True
+    assert stale_writer.storage_updated.value is True
+
+    await stale_writer.upsert({"id2": {"content": "beta"}})
+    assert await stale_writer.index_done_callback() is True
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    rows = await reader.get_by_ids(["id1", "id2"])
+    assert [row["id"] for row in rows] == ["id1", "id2"]
+    assert stale_writer._pending_upserts == {}
+    _assert_consistent(reader)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_delete_reloads_stale_client_before_mutating(tmp_path):
+    embed = _CountingEmbed()
+    writer = _make_storage(tmp_path, embed)
+    stale_deleter = _make_storage(tmp_path, embed)
+    await writer.initialize()
+    await stale_deleter.initialize()
+
+    await writer.upsert({"id1": {"content": "alpha"}})
+    assert await writer.index_done_callback() is True
+    assert stale_deleter.storage_updated.value is True
+
+    await stale_deleter.delete(["id1"])
+    assert stale_deleter.storage_updated.value is False
+    assert await stale_deleter.index_done_callback() is True
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    assert await reader.get_by_id("id1") is None
+    _assert_consistent(reader)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_finalize_reloads_stale_client_before_flushing(tmp_path):
+    embed = _CountingEmbed()
+    writer = _make_storage(tmp_path, embed)
+    stale_finalizer = _make_storage(tmp_path, embed)
+    await writer.initialize()
+    await stale_finalizer.initialize()
+
+    await writer.upsert({"id1": {"content": "alpha"}})
+    assert await writer.index_done_callback() is True
+    assert stale_finalizer.storage_updated.value is True
+
+    await stale_finalizer.upsert({"id2": {"content": "beta"}})
+    await stale_finalizer.finalize()
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    rows = await reader.get_by_ids(["id1", "id2"])
+    assert [row["id"] for row in rows] == ["id1", "id2"]
+    assert stale_finalizer._pending_upserts == {}
+    _assert_consistent(reader)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_read_your_writes_and_query_after_flush(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    # Before flush: read paths see the pending row, query does not.
+    hit = await storage.get_by_id("id1")
+    assert hit is not None and hit["id"] == "id1" and hit["content"] == "alpha"
+    by_ids = await storage.get_by_ids(["id1", "missing"])
+    assert by_ids[0]["id"] == "id1" and by_ids[1] is None
+    assert await storage.query("alpha", top_k=5) == [], "query ignores unflushed data"
+
+    # After flush: query returns the row.
+    await storage.index_done_callback()
+    results = await storage.query("alpha", top_k=5)
+    assert any(r["id"] == "id1" for r in results)
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_finalize_flushes_pending(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+    await storage.finalize()
+
+    assert embed.call_count == 1
+    assert storage._pending_upserts == {}
+    assert storage._index.ntotal == 1
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_delete_entity_relation_cancels_pending(tmp_path):
+    embed = _CountingEmbed()
+    storage = FaissVectorDBStorage(
+        namespace="test_relations",
+        workspace="ws",
+        global_config={
+            "working_dir": str(tmp_path),
+            "embedding_batch_num": 32,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=EmbeddingFunc(embedding_dim=DIM, max_token_size=512, func=embed),
+        meta_fields={"content", "src_id", "tgt_id"},
+    )
+    await storage.initialize()
+
+    # Materialize r1 (A->B), leave r2 (A->C) and r3 (X->Y) as pending.
+    await storage.upsert({"r1": {"content": "rel1", "src_id": "A", "tgt_id": "B"}})
+    await storage.index_done_callback()
+    await storage.upsert(
+        {
+            "r2": {"content": "rel2", "src_id": "A", "tgt_id": "C"},
+            "r3": {"content": "rel3", "src_id": "X", "tgt_id": "Y"},
+        }
+    )
+
+    await storage.delete_entity_relation("A")
+
+    assert "r2" not in storage._pending_upserts, "incident pending entry cancelled"
+    assert "r3" in storage._pending_upserts, "unrelated pending entry preserved"
+    assert storage._index.ntotal == 0, "materialized A->B removed"
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_flush_embedding_failure_raises_and_keeps_pending(tmp_path):
+    class _FailingEmbed:
+        def __init__(self):
+            self.call_count = 0
+
+        async def __call__(self, texts, **kwargs):
+            self.call_count += 1
+            raise RuntimeError("embed boom")
+
+    embed = _FailingEmbed()
+    storage = FaissVectorDBStorage(
+        namespace="test_vectors",
+        workspace="ws",
+        global_config={
+            "working_dir": str(tmp_path),
+            "embedding_batch_num": 32,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=EmbeddingFunc(embedding_dim=DIM, max_token_size=512, func=embed),
+        meta_fields={"content"},
+    )
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    with pytest.raises(RuntimeError, match="embed boom"):
+        await storage.index_done_callback()
+
+    assert "id1" in storage._pending_upserts, "pending preserved for retry"
+    assert storage._index.ntotal == 0, "nothing materialized on embed failure"
+    # Embed failure happens before self._index.add in _flush_pending_locked,
+    # so _index_dirty must NOT be set. (A save-stage failure would leave it True
+    # — see test_index_done_callback_save_failure_raises.)
+    assert storage._index_dirty is False
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_drop_discards_pending_without_embedding(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+    assert "id1" in storage._pending_upserts
+
+    result = await storage.drop()
+
+    assert result["status"] == "success"
+    assert storage._pending_upserts == {}, "drop discards buffered upserts"
+    assert embed.call_count == 0, "drop must not embed"
+    assert storage._index_dirty is False
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_finalize_retries_save_after_flush_failure(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    original_save = storage._save_faiss_index
+    save_calls = 0
+
+    def fail_once():
+        nonlocal save_calls
+        save_calls += 1
+        if save_calls == 1:
+            raise OSError("boom")
+        original_save()
+
+    storage._save_faiss_index = fail_once
+
+    with pytest.raises(OSError, match="boom"):
+        await storage.finalize()
+
+    assert storage._pending_upserts == {}
+    assert storage._index_dirty is True
+
+    await storage.finalize()
+
+    assert save_calls == 2
+    assert storage._index_dirty is False
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    hit = await reader.get_by_id("id1")
+    assert hit is not None and hit["id"] == "id1"
+    _assert_consistent(reader)
+
+
+# ---------------------------------------------------------------------------
+# (B) Faiss-specific tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_reupsert_after_flush_replaces_single_fid(tmp_path):
+    """Faiss has no in-place upsert: re-upserting an already-materialized id
+    must rebuild the index without the old fid, so we still end up with
+    exactly one row per custom id."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "old"}})
+    await storage.index_done_callback()
+    assert storage._index.ntotal == 1
+    _assert_consistent(storage)
+
+    await storage.upsert({"id1": {"content": "new"}})
+    await storage.index_done_callback()
+
+    assert storage._index.ntotal == 1, "rebuild must remove old fid before adding new"
+    assert len(storage._id_to_meta) == 1
+    _assert_consistent(storage)
+
+    hit = await storage.get_by_id("id1")
+    assert hit is not None and hit["content"] == "new"
+    assert embed.call_count == 2, "each flush embeds the latest content once"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_index_done_callback_save_failure_raises(tmp_path):
+    """Save failure in index_done_callback must propagate, leave pending empty
+    (flush already succeeded), and keep _index_dirty=True so finalize retries."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    original_save = storage._save_faiss_index
+
+    def fail_save():
+        raise OSError("save boom")
+
+    storage._save_faiss_index = fail_save
+
+    with pytest.raises(OSError, match="save boom"):
+        await storage.index_done_callback()
+
+    assert storage._pending_upserts == {}, "flush succeeded so pending is empty"
+    assert storage._index_dirty is True, "save failure preserves dirty for retry"
+    assert storage._index.ntotal == 1, "materialized state is preserved"
+    _assert_consistent(storage)
+
+    # Restore real save; finalize must retry only the save (no re-embed).
+    storage._save_faiss_index = original_save
+    embed_before = embed.call_count
+    await storage.finalize()
+    assert embed.call_count == embed_before, "save retry must not re-embed"
+    assert storage._index_dirty is False
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    hit = await reader.get_by_id("id1")
+    assert hit is not None and hit["id"] == "id1"
+    _assert_consistent(reader)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_reload_warns_on_index_meta_skew(tmp_path, caplog):
+    """A crash between the .index write and the .meta.json write leaves
+    ``ntotal(.index) > rows(.meta)``. ``_load_faiss_index`` must log a warning
+    on reload; auto-repair is intentionally not in scope here."""
+    import logging
+
+    from lightrag.utils import logger as lightrag_logger
+
+    embed = _CountingEmbed()
+    writer = _make_storage(tmp_path, embed)
+    await writer.initialize()
+
+    await writer.upsert({"id1": {"content": "alpha"}, "id2": {"content": "beta"}})
+    await writer.index_done_callback()
+
+    # Corrupt the meta file: drop one entry so disk has index > meta.
+    with open(writer._meta_file, "r", encoding="utf-8") as f:
+        meta = json.load(f)
+    assert len(meta) == 2
+    dropped_key = next(iter(meta))
+    del meta[dropped_key]
+    with open(writer._meta_file, "w", encoding="utf-8") as f:
+        json.dump(meta, f)
+
+    # The lightrag logger sets propagate=False (lightrag/utils.py), so caplog —
+    # which attaches to root by default — never sees its records. Flip propagate
+    # for the duration of the reload, then restore.
+    caplog.clear()
+    old_propagate = lightrag_logger.propagate
+    lightrag_logger.propagate = True
+    try:
+        with caplog.at_level(logging.WARNING, logger="lightrag"):
+            reader = _make_storage(tmp_path, embed)
+            await reader.initialize()
+    finally:
+        lightrag_logger.propagate = old_propagate
+
+    # The reader's index still has 2 vectors but only 1 reachable via meta —
+    # this is the "known risk, not auto-repaired" state.
+    assert reader._index.ntotal == 2
+    assert len(reader._id_to_meta) == 1
+    skew_messages = [
+        rec.message
+        for rec in caplog.records
+        if "skew" in rec.message or "index > meta" in rec.message
+    ]
+    assert skew_messages, (
+        f"expected an index>meta skew warning; got: "
+        f"{[r.message for r in caplog.records]}"
+    )
+
+    # Sanity: state files exist where we left them.
+    assert os.path.exists(writer._faiss_index_file)
+    assert os.path.exists(writer._meta_file)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_query_skips_orphan_faiss_hits(tmp_path):
+    """After an ``index > meta`` skew the orphan vector is still searchable by
+    similarity, but ``query`` must skip it instead of leaking a ghost
+    ``{"id": None, ...}`` row to the caller."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    # Materialize two rows.
+    await storage.upsert({"id1": {"content": "alpha"}, "id2": {"content": "beta"}})
+    await storage.index_done_callback()
+    assert storage._index.ntotal == 2
+
+    # Synthesize the skew: drop one meta row in memory, keeping the faiss
+    # index untouched. This mirrors what _load_faiss_index would surface on
+    # reload after a crash between the two atomic_writes.
+    orphan_fid = next(iter(storage._id_to_meta))
+    del storage._id_to_meta[orphan_fid]
+    assert storage._index.ntotal == 2
+    assert len(storage._id_to_meta) == 1
+
+    # The orphan vector still scores high in similarity search; query must
+    # filter it out instead of returning {"id": None, ...}.
+    results = await storage.query("anything", top_k=5)
+    for row in results:
+        assert row["id"] is not None, f"orphan hit leaked: {row}"
+    # And the surviving row is still returned.
+    surviving_id = next(iter(storage._id_to_meta.values()))["__id__"]
+    assert any(r["id"] == surviving_id for r in results)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_reupsert_cleans_duplicate_custom_id_rows(tmp_path):
+    """Defends against legacy / externally corrupted stores where multiple
+    fids in ``_id_to_meta`` share the same ``__id__``. A re-upsert + flush
+    must collapse them to a single row; a ``delete`` must remove all of them."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    # Hand-craft a corrupt state: two fids carry the same custom id "dup".
+    matrix = np.array([[1.0] * DIM, [2.0] * DIM], dtype=np.float32)
+    faiss.normalize_L2(matrix)
+    storage._index.add(matrix)
+    storage._id_to_meta[0] = {
+        "__id__": "dup",
+        "__created_at__": 1,
+        "content": "v1",
+        "__vector__": matrix[0].tolist(),
+    }
+    storage._id_to_meta[1] = {
+        "__id__": "dup",
+        "__created_at__": 1,
+        "content": "v2",
+        "__vector__": matrix[1].tolist(),
+    }
+    _assert_consistent(storage)
+    assert storage._find_faiss_ids_by_custom_id("dup") == [0, 1]
+
+    # Re-upsert + flush: both duplicates must be removed in the rebuild
+    # before the new vector is added; final state is a single row.
+    await storage.upsert({"dup": {"content": "v3"}})
+    await storage.index_done_callback()
+
+    assert storage._index.ntotal == 1, "flush rebuild must drop both duplicates"
+    assert len(storage._id_to_meta) == 1
+    assert storage._find_faiss_ids_by_custom_id("dup") == list(
+        storage._id_to_meta.keys()
+    )
+    hit = await storage.get_by_id("dup")
+    assert hit is not None and hit["content"] == "v3"
+    _assert_consistent(storage)
+
+    # Re-seed two more duplicates and verify delete also removes them all.
+    matrix2 = np.array([[3.0] * DIM, [4.0] * DIM], dtype=np.float32)
+    faiss.normalize_L2(matrix2)
+    storage._index.add(matrix2)
+    next_fid = max(storage._id_to_meta) + 1
+    storage._id_to_meta[next_fid] = {
+        "__id__": "dup",
+        "__created_at__": 2,
+        "content": "dup-a",
+        "__vector__": matrix2[0].tolist(),
+    }
+    storage._id_to_meta[next_fid + 1] = {
+        "__id__": "dup",
+        "__created_at__": 2,
+        "content": "dup-b",
+        "__vector__": matrix2[1].tolist(),
+    }
+    assert len(storage._find_faiss_ids_by_custom_id("dup")) == 3
+
+    await storage.delete(["dup"])
+    assert storage._find_faiss_ids_by_custom_id("dup") == []
+    assert storage._index.ntotal == 0
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_delete_propagates_errors(tmp_path, monkeypatch):
+    """Faiss ``delete`` must NOT swallow errors — the caller (document
+    deletion / status update path) needs to abort if vectors weren't
+    actually removed. This intentionally diverges from Nano."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+    await storage.index_done_callback()
+
+    def boom(_self, _fids):
+        raise RuntimeError("rebuild boom")
+
+    # _remove_faiss_ids_locked is what delete calls under the hood.
+    monkeypatch.setattr(
+        FaissVectorDBStorage, "_remove_faiss_ids_locked", boom, raising=True
+    )
+
+    with pytest.raises(RuntimeError, match="rebuild boom"):
+        await storage.delete(["id1"])
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_flush_recovers_from_index_add_failure_without_re_embedding(tmp_path):
+    """Self-heal contract: if ``index.add`` raises mid-flush (after embedding
+    already succeeded), the pending buffer keeps the cached vectors and a
+    subsequent ``finalize`` retries the flush **without re-embedding**. Pins
+    the "pending is the source of truth on mid-write failure" invariant
+    documented on ``_flush_pending_locked``."""
+
+    class _AddFailsOnce:
+        """Wraps a real faiss index, raising on the first ``.add`` call. After
+        the second add succeeds it swaps the storage's ``_index`` attribute
+        back to the real instance, so ``faiss.write_index`` (which requires a
+        real SWIG-wrapped object) can run during the retry's save step. This
+        is a test-only shim — in production ``self._index`` is always a real
+        faiss index throughout the retry.
+        """
+
+        def __init__(self, storage, real):
+            self._storage = storage
+            self._real = real
+            self._calls = 0
+
+        def __getattr__(self, name):
+            return getattr(self._real, name)
+
+        def add(self, arr):
+            self._calls += 1
+            if self._calls == 1:
+                raise RuntimeError("add boom")
+            result = self._real.add(arr)
+            self._storage._index = self._real
+            return result
+
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    real_index = storage._index
+    storage._index = _AddFailsOnce(storage, real_index)
+
+    with pytest.raises(RuntimeError, match="add boom"):
+        await storage.index_done_callback()
+
+    # Embedding completed once (failure happened after embed, in index.add).
+    assert embed.call_count == 1
+    # Pending preserved with cached vectors — that's the self-healing key.
+    assert "id1" in storage._pending_upserts
+    assert storage._pending_upserts["id1"].vector is not None
+    # _index_dirty stays False: docstring says we deliberately don't flip it
+    # on mid-write failure (pending is the source of truth).
+    assert storage._index_dirty is False
+    assert storage._index.ntotal == 0
+
+    # Retry through the same public entry point. The wrapper's second add
+    # succeeds, unwraps itself, and the rest of finalize (save + notify)
+    # runs against the real index.
+    await storage.finalize()
+
+    assert embed.call_count == 1, "retry must reuse cached vectors, not re-embed"
+    assert storage._index is real_index, "wrapper unwrapped itself on the second add"
+    assert storage._index.ntotal == 1
+    assert storage._pending_upserts == {}
+    assert storage._index_dirty is False
+    _assert_consistent(storage)
+
+    # And the row was persisted to disk by the retry's save.
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    hit = await reader.get_by_id("id1")
+    assert hit is not None and hit["content"] == "alpha"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_drop_pending_index_ops_clears_buffer(tmp_path):
+    """An internal-error abort calls drop_pending_index_ops to discard the
+    not-yet-flushed buffer without materializing anything into the index."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}, "id2": {"content": "beta"}})
+    assert storage._pending_upserts, "upsert buffers, does not flush"
+
+    await storage.drop_pending_index_ops()
+
+    assert storage._pending_upserts == {}
+    assert embed.call_count == 0, "drop must not embed"
+    assert storage._index.ntotal == 0, "nothing was materialized"
+    _assert_consistent(storage)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_drop_pending_does_not_rollback_materialized(tmp_path):
+    """drop_pending_index_ops discards ONLY the pending buffer; rows already
+    materialized into the index by a flush whose save then failed
+    (``_index_dirty=True``) are intentionally NOT rolled back."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    # Flush id1 into the index, then fail the save so it stays
+    # materialized-but-unsaved (dirty) and the pending buffer is emptied.
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    def fail_save():
+        raise OSError("save boom")
+
+    storage._save_faiss_index = fail_save
+    with pytest.raises(OSError, match="save boom"):
+        await storage.index_done_callback()
+    assert storage._pending_upserts == {}, "flush succeeded so pending is empty"
+    assert storage._index_dirty is True
+    assert storage._index.ntotal == 1, "id1 materialized, not saved"
+
+    # A new pending op arrives, then the batch aborts and drops pending.
+    await storage.upsert({"id2": {"content": "beta"}})
+    assert "id2" in storage._pending_upserts
+
+    await storage.drop_pending_index_ops()
+
+    assert storage._pending_upserts == {}, "pending id2 dropped"
+    assert storage._index.ntotal == 1, "materialized id1 NOT rolled back"
+    assert storage._index_dirty is True, "still dirty for a later save retry"
+    _assert_consistent(storage)

tests/kg/faiss_impl/test_faiss_meta_inconsistency.py

@@ -0,0 +1,163 @@
+"""
+Regression tests for Faiss meta/index inconsistency handling.
+
+Verifies that FaissVectorDBStorage gracefully handles cases where
+meta.json has more rows than the Faiss index (e.g., after a crash
+during save), and that delete/upsert operations don't crash.
+"""
+
+import json
+import os
+import tempfile
+
+import numpy as np
+import pytest
+
+faiss = pytest.importorskip("faiss")
+
+
+@pytest.mark.offline
+class TestFaissMetaInconsistency:
+    """Test that stale metadata rows are handled gracefully."""
+
+    def _create_index_and_meta(self, tmp_dir, dim=4, n_vectors=3, n_extra_meta=2):
+        """
+        Helper: create a Faiss index with `n_vectors` vectors and a meta.json
+        that has `n_vectors + n_extra_meta` entries (simulating a crash where
+        meta was written but index wasn't fully updated).
+        """
+        index_file = os.path.join(tmp_dir, "faiss_index_test.index")
+        meta_file = index_file + ".meta.json"
+
+        # Build real index with n_vectors
+        index = faiss.IndexFlatIP(dim)
+        vectors = np.random.rand(n_vectors, dim).astype(np.float32)
+        # Normalize for cosine similarity
+        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+        vectors = vectors / norms
+        index.add(vectors)
+        faiss.write_index(index, index_file)
+
+        # Build meta with extra rows beyond index.ntotal
+        meta = {}
+        for i in range(n_vectors):
+            meta[str(i)] = {"__id__": f"id_{i}", "content": f"text_{i}"}
+        for i in range(n_vectors, n_vectors + n_extra_meta):
+            meta[str(i)] = {"__id__": f"stale_{i}", "content": f"stale_{i}"}
+
+        with open(meta_file, "w", encoding="utf-8") as f:
+            json.dump(meta, f)
+
+        return index_file, meta_file, vectors
+
+    def test_load_skips_invalid_metadata_rows(self):
+        """
+        Loading an index where meta.json has fids beyond index.ntotal
+        should skip those rows with a warning, not crash.
+        """
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            dim = 4
+            n_vectors = 3
+            n_extra = 2
+            index_file, meta_file, vectors = self._create_index_and_meta(
+                tmp_dir, dim=dim, n_vectors=n_vectors, n_extra_meta=n_extra
+            )
+
+            # Manually load and verify behavior
+            index = faiss.read_index(index_file)
+            with open(meta_file, "r", encoding="utf-8") as f:
+                stored_dict = json.load(f)
+
+            assert len(stored_dict) == n_vectors + n_extra
+
+            # Simulate the load logic from _load_faiss_index
+            id_to_meta = {}
+            skipped = 0
+            for fid_str, meta in stored_dict.items():
+                fid = int(fid_str)
+                if fid >= index.ntotal:
+                    skipped += 1
+                    continue
+                if "__vector__" not in meta:
+                    meta["__vector__"] = index.reconstruct(fid).tolist()
+                id_to_meta[fid] = meta
+
+            assert len(id_to_meta) == n_vectors
+            assert skipped == n_extra
+
+            # Verify reconstructed vectors match originals
+            for fid in range(n_vectors):
+                reconstructed = np.array(
+                    id_to_meta[fid]["__vector__"], dtype=np.float32
+                )
+                np.testing.assert_allclose(reconstructed, vectors[fid], atol=1e-6)
+
+    def test_remove_with_missing_vector_uses_reconstruct(self):
+        """
+        _remove_faiss_ids should reconstruct vectors from the index
+        when __vector__ is not present in metadata.
+        """
+        dim = 4
+        n_vectors = 3
+
+        index = faiss.IndexFlatIP(dim)
+        vectors = np.random.rand(n_vectors, dim).astype(np.float32)
+        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+        vectors = vectors / norms
+        index.add(vectors)
+
+        # Metadata WITHOUT __vector__ (as stored on disk after our PR)
+        id_to_meta = {}
+        for i in range(n_vectors):
+            id_to_meta[i] = {"__id__": f"id_{i}", "content": f"text_{i}"}
+
+        # Simulate rebuild logic from _remove_faiss_ids (remove fid=1)
+        fid_list = [1]
+        keep_fids = [fid for fid in id_to_meta if fid not in fid_list]
+
+        vectors_to_keep = []
+        new_id_to_meta = {}
+        for new_fid, old_fid in enumerate(keep_fids):
+            vec_meta = id_to_meta[old_fid]
+            if "__vector__" in vec_meta:
+                vec = vec_meta["__vector__"]
+            elif old_fid < index.ntotal:
+                vec = index.reconstruct(old_fid).tolist()
+                vec_meta["__vector__"] = vec
+            else:
+                continue
+            vectors_to_keep.append(vec)
+            new_id_to_meta[new_fid] = vec_meta
+
+        assert len(vectors_to_keep) == 2
+        assert len(new_id_to_meta) == 2
+        # Verify the kept vectors match originals (fid 0 and 2)
+        np.testing.assert_allclose(
+            np.array(vectors_to_keep[0], dtype=np.float32), vectors[0], atol=1e-6
+        )
+        np.testing.assert_allclose(
+            np.array(vectors_to_keep[1], dtype=np.float32), vectors[2], atol=1e-6
+        )
+
+    def test_atomic_save_meta(self):
+        """
+        _save_faiss_index should write meta.json atomically via temp file + os.replace.
+        Verify no .tmp file remains after save.
+        """
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            meta_file = os.path.join(tmp_dir, "test.meta.json")
+            tmp_meta_file = meta_file + ".tmp"
+
+            serializable_dict = {"0": {"__id__": "id_0", "content": "text_0"}}
+
+            # Simulate atomic write
+            with open(tmp_meta_file, "w", encoding="utf-8") as f:
+                json.dump(serializable_dict, f)
+            os.replace(tmp_meta_file, meta_file)
+
+            assert os.path.exists(meta_file)
+            assert not os.path.exists(tmp_meta_file)
+
+            with open(meta_file, "r", encoding="utf-8") as f:
+                loaded = json.load(f)
+            assert loaded == serializable_dict

tests/kg/json_impl/test_atomic_write_write_json.py

@@ -0,0 +1,73 @@
+"""Atomicity coverage for ``lightrag.utils.write_json``.
+
+The two storages that ride on this function (``JsonDocStatusStorage``,
+``JsonKVStorage``) inherit crash safety from it, so the contract lives here:
+
+- A crash during the rename leaves the prior snapshot intact.
+- The sanitize fallback also lands atomically (one tmp, one rename).
+"""
+
+import json
+import os
+import threading
+from unittest.mock import patch
+
+import pytest
+
+from lightrag.utils import write_json
+
+
+@pytest.mark.offline
+def test_write_json_publishes_clean_payload(tmp_path):
+    target = str(tmp_path / "kv.json")
+    needs_reload = write_json({"a": 1, "b": "hello"}, target)
+    assert needs_reload is False
+    assert json.load(open(target)) == {"a": 1, "b": "hello"}
+    assert [p for p in os.listdir(tmp_path) if ".tmp." in p] == []
+
+
+@pytest.mark.offline
+def test_write_json_replace_crash_preserves_prior_snapshot(tmp_path):
+    target = str(tmp_path / "kv.json")
+    write_json({"v": 1}, target)
+
+    with patch(
+        "lightrag.file_atomic.os.replace",
+        side_effect=OSError("simulated crash"),
+    ):
+        with pytest.raises(OSError, match="simulated crash"):
+            write_json({"v": 2}, target)
+
+    assert json.load(open(target)) == {"v": 1}
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"write_json must clean tmp on crash, got {leftovers}"
+
+
+@pytest.mark.offline
+def test_write_json_concurrent_writers_land_intact(tmp_path):
+    """Multiple threads racing on the same destination must each rename
+    cleanly. The final file must be valid JSON (one writer's payload)."""
+    target = str(tmp_path / "kv.json")
+    errors: list[BaseException] = []
+    barrier = threading.Barrier(5)
+
+    def writer(tid: int) -> None:
+        try:
+            barrier.wait()
+            write_json({"writer": tid}, target)
+        except BaseException as exc:
+            errors.append(exc)
+
+    threads = [threading.Thread(target=writer, args=(i,)) for i in range(5)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert errors == [], f"Concurrent writers raised: {errors}"
+
+    payload = json.load(open(target))
+    assert payload.keys() == {"writer"}
+    assert payload["writer"] in range(5)
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"Unexpected tmp residue: {leftovers}"

tests/kg/json_impl/test_json_doc_status_storage.py

@@ -0,0 +1,168 @@
+import pytest
+
+from lightrag.base import DocStatus
+from lightrag.kg.json_doc_status_impl import JsonDocStatusStorage
+from lightrag.kg.json_kv_impl import JsonKVStorage
+from lightrag.kg.shared_storage import finalize_share_data, initialize_share_data
+from lightrag.namespace import NameSpace
+
+pytestmark = pytest.mark.offline
+
+
+class _DummyEmbeddingFunc:
+    embedding_dim = 1
+    max_token_size = 1
+
+    async def __call__(self, texts, **kwargs):
+        return [[0.0] for _ in texts]
+
+
+def _doc(status: str, file_path: str) -> dict:
+    return {
+        "content_summary": f"{status} summary",
+        "content_length": 10,
+        "file_path": file_path,
+        "status": status,
+        "created_at": "2024-01-01T00:00:00+00:00",
+        "updated_at": "2024-01-01T00:00:00+00:00",
+        "metadata": {},
+        "error_msg": None,
+    }
+
+
+@pytest.fixture(autouse=True)
+def setup_shared_data():
+    initialize_share_data()
+    yield
+    finalize_share_data()
+
+
+@pytest.mark.asyncio
+async def test_get_docs_paginated_with_status_filters(tmp_path):
+    storage = JsonDocStatusStorage(
+        namespace="doc_status",
+        global_config={"working_dir": str(tmp_path)},
+        embedding_func=_DummyEmbeddingFunc(),
+        workspace="test",
+    )
+    await storage.initialize()
+
+    async with storage._storage_lock:
+        storage._data.update(
+            {
+                "doc-1": _doc("preprocessed", "a.pdf"),
+                "doc-2": _doc("parsing", "b.pdf"),
+                "doc-3": _doc("analyzing", "c.pdf"),
+                "doc-4": _doc("processed", "d.pdf"),
+            }
+        )
+
+    docs, total = await storage.get_docs_paginated(
+        status_filter=DocStatus.PROCESSED,
+        status_filters=[
+            DocStatus.PREPROCESSED,
+            DocStatus.PARSING,
+            DocStatus.ANALYZING,
+        ],
+        page=1,
+        page_size=10,
+        sort_field="id",
+        sort_direction="asc",
+    )
+
+    assert total == 3
+    assert [doc_id for doc_id, _ in docs] == ["doc-1", "doc-2", "doc-3"]
+    assert [doc.status for _, doc in docs] == [
+        DocStatus.PREPROCESSED,
+        DocStatus.PARSING,
+        DocStatus.ANALYZING,
+    ]
+
+
+@pytest.mark.asyncio
+async def test_doc_status_upsert_preserves_caller_file_path(tmp_path):
+    storage = JsonDocStatusStorage(
+        namespace=NameSpace.DOC_STATUS,
+        global_config={"working_dir": str(tmp_path)},
+        embedding_func=_DummyEmbeddingFunc(),
+        workspace="test",
+    )
+    await storage.initialize()
+
+    await storage.upsert(
+        {
+            "doc-1": _doc(
+                DocStatus.PENDING.value,
+                "/tmp/uploads/report.[native-Fi].pdf",
+            )
+        }
+    )
+
+    assert (await storage.get_by_id("doc-1"))["file_path"] == (
+        "/tmp/uploads/report.[native-Fi].pdf"
+    )
+    assert await storage.get_doc_by_file_basename("report.pdf") is None
+
+
+@pytest.mark.asyncio
+async def test_doc_status_basename_lookup_requires_canonical_stored_path(tmp_path):
+    storage = JsonDocStatusStorage(
+        namespace=NameSpace.DOC_STATUS,
+        global_config={"working_dir": str(tmp_path)},
+        embedding_func=_DummyEmbeddingFunc(),
+        workspace="test",
+    )
+    await storage.initialize()
+
+    async with storage._storage_lock:
+        storage._data["doc-1"] = _doc(
+            DocStatus.PROCESSED.value,
+            "report.[native].pdf",
+        )
+
+    assert await storage.get_doc_by_file_basename("report.pdf") is None
+
+
+@pytest.mark.asyncio
+async def test_json_kv_upsert_preserves_caller_file_paths(tmp_path):
+    full_docs = JsonKVStorage(
+        namespace=NameSpace.KV_STORE_FULL_DOCS,
+        global_config={"working_dir": str(tmp_path)},
+        embedding_func=_DummyEmbeddingFunc(),
+        workspace="test",
+    )
+    text_chunks = JsonKVStorage(
+        namespace=NameSpace.KV_STORE_TEXT_CHUNKS,
+        global_config={"working_dir": str(tmp_path)},
+        embedding_func=_DummyEmbeddingFunc(),
+        workspace="test",
+    )
+    await full_docs.initialize()
+    await text_chunks.initialize()
+
+    await full_docs.upsert(
+        {
+            "doc-1": {
+                "content": "full text",
+                "file_path": "/tmp/uploads/report.[native-Fi].pdf",
+            }
+        }
+    )
+    await text_chunks.upsert(
+        {
+            "chunk-1": {
+                "content": "chunk text",
+                "tokens": 2,
+                "chunk_order_index": 0,
+                "full_doc_id": "doc-1",
+                "file_path": "/tmp/uploads/report.[native-Fi].pdf",
+            }
+        }
+    )
+
+    assert (await full_docs.get_by_id("doc-1"))["file_path"] == (
+        "/tmp/uploads/report.[native-Fi].pdf"
+    )
+    assert (await text_chunks.get_by_id("chunk-1"))["file_path"] == (
+        "/tmp/uploads/report.[native-Fi].pdf"
+    )

tests/kg/json_impl/test_write_json_optimization.py

@@ -0,0 +1,490 @@
+"""
+Test suite for write_json optimization
+
+This test verifies:
+1. Fast path works for clean data (no sanitization)
+2. Slow path applies sanitization for dirty data
+3. Sanitization is done during encoding (memory-efficient)
+4. Reloading updates shared memory with cleaned data
+"""
+
+import os
+import json
+import tempfile
+import pytest
+from lightrag.utils import (
+    write_json,
+    load_json,
+    SanitizingJSONEncoder,
+    sanitize_text_for_encoding,
+)
+
+
+@pytest.mark.offline
+class TestWriteJsonOptimization:
+    """Test write_json optimization with two-stage approach"""
+
+    def test_fast_path_clean_data(self):
+        """Test that clean data takes the fast path without sanitization"""
+        clean_data = {
+            "name": "John Doe",
+            "age": 30,
+            "items": ["apple", "banana", "cherry"],
+            "nested": {"key": "value", "number": 42},
+        }
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            # Write clean data - should return False (no sanitization)
+            needs_reload = write_json(clean_data, temp_file)
+            assert not needs_reload, "Clean data should not require sanitization"
+
+            # Verify data was written correctly
+            loaded_data = load_json(temp_file)
+            assert loaded_data == clean_data, "Loaded data should match original"
+        finally:
+            os.unlink(temp_file)
+
+    def test_slow_path_dirty_data(self):
+        """Test that dirty data triggers sanitization"""
+        # Create data with surrogate characters (U+D800 to U+DFFF)
+        dirty_string = "Hello\ud800World"  # Contains surrogate character
+        dirty_data = {"text": dirty_string, "number": 123}
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            # Write dirty data - should return True (sanitization applied)
+            needs_reload = write_json(dirty_data, temp_file)
+            assert needs_reload, "Dirty data should trigger sanitization"
+
+            # Verify data was written and sanitized
+            loaded_data = load_json(temp_file)
+            assert loaded_data is not None, "Data should be written"
+            assert loaded_data["number"] == 123, "Clean fields should remain unchanged"
+            # Surrogate character should be removed
+            assert "\ud800" not in loaded_data["text"], (
+                "Surrogate character should be removed"
+            )
+        finally:
+            os.unlink(temp_file)
+
+    def test_sanitizing_encoder_removes_surrogates(self):
+        """Test that SanitizingJSONEncoder removes surrogate characters"""
+        data_with_surrogates = {
+            "text": "Hello\ud800\udc00World",  # Contains surrogate pair
+            "clean": "Clean text",
+            "nested": {"dirty_key\ud801": "value", "clean_key": "clean\ud802value"},
+        }
+
+        # Encode using custom encoder
+        encoded = json.dumps(
+            data_with_surrogates, cls=SanitizingJSONEncoder, ensure_ascii=False
+        )
+
+        # Verify no surrogate characters in output
+        assert "\ud800" not in encoded, "Surrogate U+D800 should be removed"
+        assert "\udc00" not in encoded, "Surrogate U+DC00 should be removed"
+        assert "\ud801" not in encoded, "Surrogate U+D801 should be removed"
+        assert "\ud802" not in encoded, "Surrogate U+D802 should be removed"
+
+        # Verify clean parts remain
+        assert "Clean text" in encoded, "Clean text should remain"
+        assert "clean_key" in encoded, "Clean keys should remain"
+
+    def test_nested_structure_sanitization(self):
+        """Test sanitization of deeply nested structures"""
+        nested_data = {
+            "level1": {
+                "level2": {
+                    "level3": {"dirty": "text\ud800here", "clean": "normal text"},
+                    "list": ["item1", "item\ud801dirty", "item3"],
+                }
+            }
+        }
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            needs_reload = write_json(nested_data, temp_file)
+            assert needs_reload, "Nested dirty data should trigger sanitization"
+
+            # Verify nested structure is preserved
+            loaded_data = load_json(temp_file)
+            assert "level1" in loaded_data
+            assert "level2" in loaded_data["level1"]
+            assert "level3" in loaded_data["level1"]["level2"]
+
+            # Verify surrogates are removed
+            dirty_text = loaded_data["level1"]["level2"]["level3"]["dirty"]
+            assert "\ud800" not in dirty_text, "Nested surrogate should be removed"
+
+            # Verify list items are sanitized
+            list_items = loaded_data["level1"]["level2"]["list"]
+            assert "\ud801" not in list_items[1], (
+                "List item surrogates should be removed"
+            )
+        finally:
+            os.unlink(temp_file)
+
+    def test_unicode_non_characters_removed(self):
+        """Test that Unicode non-characters (U+FFFE, U+FFFF) don't cause encoding errors
+
+        Note: U+FFFE and U+FFFF are valid UTF-8 characters (though discouraged),
+        so they don't trigger sanitization. They only get removed when explicitly
+        using the SanitizingJSONEncoder.
+        """
+        data_with_nonchars = {"text1": "Hello\ufffeWorld", "text2": "Test\uffffString"}
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            # These characters are valid UTF-8, so they take the fast path
+            needs_reload = write_json(data_with_nonchars, temp_file)
+            assert not needs_reload, "U+FFFE/U+FFFF are valid UTF-8 characters"
+
+            loaded_data = load_json(temp_file)
+            # They're written as-is in the fast path
+            assert loaded_data == data_with_nonchars
+        finally:
+            os.unlink(temp_file)
+
+    def test_mixed_clean_dirty_data(self):
+        """Test data with both clean and dirty fields"""
+        mixed_data = {
+            "clean_field": "This is perfectly fine",
+            "dirty_field": "This has\ud800issues",
+            "number": 42,
+            "boolean": True,
+            "null_value": None,
+            "clean_list": [1, 2, 3],
+            "dirty_list": ["clean", "dirty\ud801item"],
+        }
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            needs_reload = write_json(mixed_data, temp_file)
+            assert needs_reload, (
+                "Mixed data with dirty fields should trigger sanitization"
+            )
+
+            loaded_data = load_json(temp_file)
+
+            # Clean fields should remain unchanged
+            assert loaded_data["clean_field"] == "This is perfectly fine"
+            assert loaded_data["number"] == 42
+            assert loaded_data["boolean"]
+            assert loaded_data["null_value"] is None
+            assert loaded_data["clean_list"] == [1, 2, 3]
+
+            # Dirty fields should be sanitized
+            assert "\ud800" not in loaded_data["dirty_field"]
+            assert "\ud801" not in loaded_data["dirty_list"][1]
+        finally:
+            os.unlink(temp_file)
+
+    def test_empty_and_none_strings(self):
+        """Test handling of empty and None values"""
+        data = {
+            "empty": "",
+            "none": None,
+            "zero": 0,
+            "false": False,
+            "empty_list": [],
+            "empty_dict": {},
+        }
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            needs_reload = write_json(data, temp_file)
+            assert not needs_reload, (
+                "Clean empty values should not trigger sanitization"
+            )
+
+            loaded_data = load_json(temp_file)
+            assert loaded_data == data, "Empty/None values should be preserved"
+        finally:
+            os.unlink(temp_file)
+
+    def test_specific_surrogate_udc9a(self):
+        """Test specific surrogate character \\udc9a mentioned in the issue"""
+        # Test the exact surrogate character from the error message:
+        # UnicodeEncodeError: 'utf-8' codec can't encode character '\\udc9a'
+        data_with_udc9a = {
+            "text": "Some text with surrogate\udc9acharacter",
+            "position": 201,  # As mentioned in the error
+            "clean_field": "Normal text",
+        }
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            # Write data - should trigger sanitization
+            needs_reload = write_json(data_with_udc9a, temp_file)
+            assert needs_reload, "Data with \\udc9a should trigger sanitization"
+
+            # Verify surrogate was removed
+            loaded_data = load_json(temp_file)
+            assert loaded_data is not None
+            assert "\udc9a" not in loaded_data["text"], "\\udc9a should be removed"
+            assert loaded_data["clean_field"] == "Normal text", (
+                "Clean fields should remain"
+            )
+        finally:
+            os.unlink(temp_file)
+
+    def test_migration_with_surrogate_sanitization(self):
+        """Test that migration process handles surrogate characters correctly
+
+        This test simulates the scenario where legacy cache contains surrogate
+        characters and ensures they are cleaned during migration.
+        """
+        # Simulate legacy cache data with surrogate characters
+        legacy_data_with_surrogates = {
+            "cache_entry_1": {
+                "return": "Result with\ud800surrogate",
+                "cache_type": "extract",
+                "original_prompt": "Some\udc9aprompt",
+            },
+            "cache_entry_2": {
+                "return": "Clean result",
+                "cache_type": "query",
+                "original_prompt": "Clean prompt",
+            },
+        }
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            # First write the dirty data directly (simulating legacy cache file)
+            # Use custom encoder to force write even with surrogates
+            with open(temp_file, "w", encoding="utf-8") as f:
+                json.dump(
+                    legacy_data_with_surrogates,
+                    f,
+                    cls=SanitizingJSONEncoder,
+                    ensure_ascii=False,
+                )
+
+            # Load and verify surrogates were cleaned during initial write
+            loaded_data = load_json(temp_file)
+            assert loaded_data is not None
+
+            # The data should be sanitized
+            assert "\ud800" not in loaded_data["cache_entry_1"]["return"], (
+                "Surrogate in return should be removed"
+            )
+            assert "\udc9a" not in loaded_data["cache_entry_1"]["original_prompt"], (
+                "Surrogate in prompt should be removed"
+            )
+
+            # Clean data should remain unchanged
+            assert loaded_data["cache_entry_2"]["return"] == "Clean result", (
+                "Clean data should remain"
+            )
+
+        finally:
+            os.unlink(temp_file)
+
+    def test_empty_values_after_sanitization(self):
+        """Test that data with empty values after sanitization is properly handled
+
+        Critical edge case: When sanitization results in data with empty string values,
+        we must use 'if cleaned_data is not None' instead of 'if cleaned_data' to ensure
+        proper reload, since truthy check on dict depends on content, not just existence.
+        """
+        # Create data where ALL values are only surrogate characters
+        all_dirty_data = {
+            "key1": "\ud800\udc00\ud801",
+            "key2": "\ud802\ud803",
+        }
+
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
+            temp_file = f.name
+
+        try:
+            # Write dirty data - should trigger sanitization
+            needs_reload = write_json(all_dirty_data, temp_file)
+            assert needs_reload, "All-dirty data should trigger sanitization"
+
+            # Load the sanitized data
+            cleaned_data = load_json(temp_file)
+
+            # Critical assertions for the edge case
+            assert cleaned_data is not None, "Cleaned data should not be None"
+            # Sanitization removes surrogates but preserves keys with empty values
+            assert cleaned_data == {
+                "key1": "",
+                "key2": "",
+            }, "Surrogates should be removed, keys preserved"
+            # This dict is truthy because it has keys (even with empty values)
+            assert cleaned_data, "Dict with keys is truthy"
+
+            # Test the actual edge case: empty dict
+            empty_data = {}
+            needs_reload2 = write_json(empty_data, temp_file)
+            assert not needs_reload2, "Empty dict is clean"
+
+            reloaded_empty = load_json(temp_file)
+            assert reloaded_empty is not None, "Empty dict should not be None"
+            assert reloaded_empty == {}, "Empty dict should remain empty"
+            assert not reloaded_empty, (
+                "Empty dict evaluates to False (the critical check)"
+            )
+
+        finally:
+            os.unlink(temp_file)
+
+
+@pytest.mark.offline
+class TestSanitizeTextForEncoding:
+    """Direct unit tests for sanitize_text_for_encoding function."""
+
+    def test_empty_string_returns_empty(self):
+        assert sanitize_text_for_encoding("") == ""
+
+    def test_none_like_falsy_returns_as_is(self):
+        # The function checks `if not text`, so empty string returns early
+        assert sanitize_text_for_encoding("") == ""
+
+    def test_whitespace_only_returns_empty(self):
+        assert sanitize_text_for_encoding("   ") == ""
+
+    def test_clean_text_unchanged(self):
+        assert sanitize_text_for_encoding("hello world") == "hello world"
+
+    def test_strips_leading_trailing_whitespace(self):
+        assert sanitize_text_for_encoding("  hello  ") == "hello"
+
+    def test_lone_surrogate_removed(self):
+        assert sanitize_text_for_encoding("hello\ud800world") == "helloworld"
+
+    def test_lone_surrogate_with_replacement_char(self):
+        assert (
+            sanitize_text_for_encoding("hello\ud800world", replacement_char="?")
+            == "hello?world"
+        )
+
+    def test_surrogate_range_boundaries(self):
+        # U+D800 and U+DFFF are the surrogate range boundaries
+        assert "\ud800" not in sanitize_text_for_encoding("\ud800")
+        assert "\udfff" not in sanitize_text_for_encoding("\udfff")
+
+    def test_non_characters_fffe_ffff_removed(self):
+        # U+FFFE and U+FFFF are included in _SURROGATE_PATTERN
+        assert sanitize_text_for_encoding("a\ufffeb") == "ab"
+        assert sanitize_text_for_encoding("a\uffffb") == "ab"
+
+    def test_html_entities_unescaped(self):
+        assert sanitize_text_for_encoding("&") == "&"
+        assert sanitize_text_for_encoding("&lt;p&gt;") == "<p>"
+        assert sanitize_text_for_encoding("&quot;hello&quot;") == '"hello"'
+
+    def test_html_entity_that_becomes_surrogate_is_removed(self):
+        # &#xD800; — Python's html.unescape follows HTML5 spec and maps surrogate code
+        # points to U+FFFD (replacement character), so \uD800 never appears in output.
+        # Either way the result must not contain an actual lone surrogate.
+        result = sanitize_text_for_encoding("&#xD800;")
+        assert "\ud800" not in result
+
+    def test_control_chars_removed(self):
+        # C0 control characters (excluding \t \n \r)
+        assert sanitize_text_for_encoding("\x01hello\x1fworld") == "helloworld"
+        assert sanitize_text_for_encoding("\x00null") == "null"
+        assert sanitize_text_for_encoding("del\x7f") == "del"
+
+    def test_control_chars_with_replacement_char(self):
+        # replacement_char must apply to control chars, not just surrogates.
+        # Note: \x1f is treated as Unicode whitespace by Python's str.strip(),
+        # so place control chars in the middle to avoid them being stripped first.
+        result = sanitize_text_for_encoding("a\x01b\x08c", replacement_char="?")
+        assert result == "a?b?c"
+
+    def test_common_whitespace_preserved(self):
+        # \t, \n, \r must NOT be removed (excluded from control char pattern)
+        assert sanitize_text_for_encoding("line1\nline2") == "line1\nline2"
+        assert sanitize_text_for_encoding("col1\tcol2") == "col1\tcol2"
+        assert sanitize_text_for_encoding("line1\rline2") == "line1\rline2"
+
+    def test_c1_control_chars_not_removed(self):
+        # \x80-\x9F range must NOT be removed (restored original behavior).
+        # These are valid in Latin-1 encoded European language text.
+        result = sanitize_text_for_encoding("caf\x85e")
+        assert "\x85" in result
+
+    def test_replacement_char_default_is_deletion(self):
+        # Default replacement_char="" means characters are deleted, not replaced
+        assert sanitize_text_for_encoding("\ud800hello\x01") == "hello"
+
+    def test_mixed_issues_in_one_string(self):
+        # Surrogate + control char + HTML entity + clean text
+        text = "\ud800&amp;\x01clean"
+        result = sanitize_text_for_encoding(text)
+        assert result == "&clean"
+
+    def test_large_text_with_scattered_surrogates(self):
+        # Regression guard: regex must handle large inputs correctly
+        clean_segment = "a" * 10000
+        text = f"prefix\ud800{clean_segment}\udfffsuffix"
+        result = sanitize_text_for_encoding(text)
+        assert "\ud800" not in result
+        assert "\udfff" not in result
+        assert clean_segment in result
+
+
+if __name__ == "__main__":
+    # Run tests
+    test = TestWriteJsonOptimization()
+
+    print("Running test_fast_path_clean_data...")
+    test.test_fast_path_clean_data()
+    print("✓ Passed")
+
+    print("Running test_slow_path_dirty_data...")
+    test.test_slow_path_dirty_data()
+    print("✓ Passed")
+
+    print("Running test_sanitizing_encoder_removes_surrogates...")
+    test.test_sanitizing_encoder_removes_surrogates()
+    print("✓ Passed")
+
+    print("Running test_nested_structure_sanitization...")
+    test.test_nested_structure_sanitization()
+    print("✓ Passed")
+
+    print("Running test_unicode_non_characters_removed...")
+    test.test_unicode_non_characters_removed()
+    print("✓ Passed")
+
+    print("Running test_mixed_clean_dirty_data...")
+    test.test_mixed_clean_dirty_data()
+    print("✓ Passed")
+
+    print("Running test_empty_and_none_strings...")
+    test.test_empty_and_none_strings()
+    print("✓ Passed")
+
+    print("Running test_specific_surrogate_udc9a...")
+    test.test_specific_surrogate_udc9a()
+    print("✓ Passed")
+
+    print("Running test_migration_with_surrogate_sanitization...")
+    test.test_migration_with_surrogate_sanitization()
+    print("✓ Passed")
+
+    print("Running test_empty_values_after_sanitization...")
+    test.test_empty_values_after_sanitization()
+    print("✓ Passed")
+
+    print("\n✅ All tests passed!")

tests/kg/memgraph_impl/test_memgraph_storage.py

@@ -0,0 +1,137 @@
+import pytest
+
+from lightrag.kg.memgraph_impl import MemgraphStorage
+
+
+pytestmark = pytest.mark.offline
+
+
+class _FakeNode(dict):
+    def __init__(self, node_id: int, entity_id: str, **properties):
+        super().__init__(entity_id=entity_id, **properties)
+        self.id = node_id
+
+
+class _FakeResult:
+    def __init__(self, record):
+        self._record = record
+
+    async def single(self):
+        return self._record
+
+    async def consume(self):
+        return None
+
+
+class _FakeSession:
+    def __init__(self, record, calls):
+        self._record = record
+        self._calls = calls
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    async def run(self, query, parameters=None, **kwargs):
+        if parameters is None:
+            parameters = kwargs
+        self._calls.append((query, parameters))
+        return _FakeResult(self._record)
+
+
+class _FakeDriver:
+    def __init__(self, record, calls):
+        self._record = record
+        self._calls = calls
+
+    def session(self, **kwargs):
+        return _FakeSession(self._record, self._calls)
+
+
+def _make_storage(record):
+    calls = []
+    storage = MemgraphStorage(
+        namespace="chunk_entity_relation",
+        global_config={"max_graph_nodes": 1000},
+        embedding_func=None,
+        workspace="test",
+    )
+    storage._driver = _FakeDriver(record, calls)
+    storage._DATABASE = "memgraph"
+    return storage, calls
+
+
+@pytest.mark.asyncio
+async def test_get_knowledge_graph_preserves_isolated_start_node():
+    start_node = _FakeNode(1, "Start", description="isolated")
+    storage, calls = _make_storage(
+        {
+            "node_info": [{"node": start_node}],
+            "relationships": [],
+            "is_truncated": False,
+        }
+    )
+
+    result = await storage.get_knowledge_graph("Start", max_depth=0, max_nodes=1)
+
+    # Verify result data: isolated node must appear with correct labels and properties
+    assert len(result.nodes) == 1
+    assert result.nodes[0].labels == ["Start"]
+    assert result.nodes[0].properties["entity_id"] == "Start"
+    assert result.edges == []
+    assert result.is_truncated is False
+
+    # Verify query parameters: max_other_nodes must reserve a slot for the start node
+    assert len(calls) == 1
+    _, params = calls[0]
+    assert params["entity_id"] == "Start"
+    assert params["max_nodes"] == 1
+    assert (
+        params["max_other_nodes"] == 0
+    )  # max_nodes - 1 = 0, start node occupies the only slot
+
+
+@pytest.mark.asyncio
+async def test_get_knowledge_graph_reserves_capacity_for_start_node_when_truncating():
+    start_node = _FakeNode(1, "Start")
+    storage, calls = _make_storage(
+        {
+            "node_info": [{"node": start_node}],
+            "relationships": [],
+            "is_truncated": True,
+        }
+    )
+
+    result = await storage.get_knowledge_graph("Start", max_depth=2, max_nodes=2)
+
+    # Verify truncation is reflected in result
+    assert result.is_truncated is True
+    assert len(result.nodes) == 1
+    assert result.edges == []
+
+    # Verify max_other_nodes leaves exactly one slot for the start node
+    assert len(calls) == 1
+    _, params = calls[0]
+    assert params["max_nodes"] == 2
+    assert (
+        params["max_other_nodes"] == 1
+    )  # max_nodes - 1 = 1, start node always included
+
+
+@pytest.mark.asyncio
+async def test_get_knowledge_graph_max_nodes_zero_does_not_underflow():
+    """max_other_nodes must not go negative when max_nodes=0."""
+    storage, calls = _make_storage(
+        {
+            "node_info": [],
+            "relationships": [],
+            "is_truncated": False,
+        }
+    )
+
+    await storage.get_knowledge_graph("Start", max_depth=1, max_nodes=0)
+
+    _, params = calls[0]
+    assert params["max_other_nodes"] == 0  # max(0 - 1, 0) = 0, no underflow

tests/kg/milvus_impl/test_milvus_deferred_embedding.py

@@ -0,0 +1,720 @@
+"""Unit tests for MilvusVectorDBStorage's deferred-embedding flush pipeline.
+
+All tests use mocks — no running Milvus instance required.
+Mirrors the structure of tests/kg/opensearch_impl/test_opensearch_storage.py's
+TestVectorStorageBatching to keep behaviour aligned across backends.
+"""
+
+import asyncio
+import os
+
+import numpy as np
+import pytest
+from unittest.mock import MagicMock, patch
+
+from lightrag.kg.milvus_impl import MILVUS_MAX_VARCHAR_BYTES, MilvusVectorDBStorage
+
+pytestmark = pytest.mark.offline
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+class MockEmbeddingFunc:
+    """Mock embedding function that returns random vectors."""
+
+    def __init__(self, dim=8):
+        self.embedding_dim = dim
+        self.max_token_size = 512
+        self.model_name = "mock-embed"
+
+    async def __call__(self, texts, **kwargs):
+        return np.random.rand(len(texts), self.embedding_dim).astype(np.float32)
+
+
+class CountingEmbeddingFunc(MockEmbeddingFunc):
+    """Embedding test double that records calls and can fail a fixed number of times."""
+
+    def __init__(self, dim=8, fail_times=0):
+        super().__init__(dim=dim)
+        self.fail_times = fail_times
+        self.call_count = 0
+        self.batches: list[list[str]] = []
+        self.texts: list[str] = []
+
+    async def __call__(self, texts, **kwargs):
+        self.call_count += 1
+        batch = list(texts)
+        self.batches.append(batch)
+        self.texts.extend(batch)
+        if self.fail_times > 0:
+            self.fail_times -= 1
+            raise RuntimeError("embedding failed")
+        return await super().__call__(texts, **kwargs)
+
+
+@pytest.fixture(autouse=True)
+def patch_namespace_lock():
+    """Cache real asyncio.Locks per (namespace, workspace) for shared semantics.
+
+    Two storage instances whose ``final_namespace`` matches must observe the
+    same Lock instance — this fixture lets us assert that and also exercises
+    real serialization between concurrent flush/upsert coroutines.
+    """
+    cache: dict[tuple[str, str | None], asyncio.Lock] = {}
+
+    def factory(namespace, workspace=None, enable_logging=False):
+        key = (namespace, workspace or "")
+        lock = cache.get(key)
+        if lock is None:
+            lock = asyncio.Lock()
+            cache[key] = lock
+        return lock
+
+    with patch("lightrag.kg.milvus_impl.get_namespace_lock", side_effect=factory):
+        yield cache
+
+
+def _make_storage(
+    embed_func,
+    *,
+    namespace="entities",
+    workspace="test",
+    meta_fields=None,
+):
+    """Build a MilvusVectorDBStorage skipping `initialize()` (no real client)."""
+    if meta_fields is None:
+        meta_fields = {"content", "entity_name", "src_id", "tgt_id"}
+    storage = MilvusVectorDBStorage(
+        namespace=namespace,
+        workspace=workspace,
+        global_config={
+            "embedding_batch_num": 10,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=embed_func,
+        meta_fields=meta_fields,
+    )
+    # Bypass real Milvus client; manually wire the bits initialize() would set.
+    # The flush lock is already constructed in __post_init__ via the patched
+    # get_namespace_lock factory, so no manual lock wiring is needed here.
+    storage._client = MagicMock()
+    storage._client.has_collection.return_value = True
+    storage._client.upsert = MagicMock(return_value={"upsert_count": 0})
+    storage._client.delete = MagicMock(return_value={"delete_count": 0})
+    storage._client.query = MagicMock(return_value=[])
+    storage._client.load_collection = MagicMock()
+    storage._initialized = True
+    return storage
+
+
+# ---------------------------------------------------------------------------
+# Tests: deferred embedding + batched flush
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_upsert_buffers_without_embedding():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}, "v2": {"content": "world"}})
+
+    assert embed.call_count == 0
+    assert set(s._pending_vector_docs.keys()) == {"v1", "v2"}
+    assert s._pending_vector_docs["v1"].vector is None
+    assert s._pending_vector_docs["v2"].vector is None
+    s._client.upsert.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_index_done_callback_triggers_flush():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}, "v2": {"content": "world"}})
+    await s.index_done_callback()
+
+    assert embed.call_count == 1
+    s._client.upsert.assert_called_once()
+    call_kwargs = s._client.upsert.call_args.kwargs
+    assert call_kwargs["collection_name"] == s.final_namespace
+    upserted = call_kwargs["data"]
+    assert {row["id"] for row in upserted} == {"v1", "v2"}
+    assert all("vector" in row for row in upserted)
+    # Buffers cleared after a successful flush.
+    assert s._pending_vector_docs == {}
+    assert s._pending_vector_deletes == set()
+
+
+@pytest.mark.asyncio
+async def test_upsert_truncates_oversized_content_for_payload_only():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed, meta_fields={"content"})
+    content = "x" * (MILVUS_MAX_VARCHAR_BYTES + 10)
+
+    await s.upsert({"v1": {"content": content}})
+    await s.index_done_callback()
+
+    upserted = s._client.upsert.call_args.kwargs["data"][0]
+    assert len(upserted["content"].encode("utf-8")) == MILVUS_MAX_VARCHAR_BYTES
+    assert embed.texts == [content]
+    assert s._pending_vector_docs == {}
+
+
+@pytest.mark.asyncio
+async def test_upsert_truncates_multibyte_content_on_character_boundary():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed, meta_fields={"content"})
+    content = "源" * (MILVUS_MAX_VARCHAR_BYTES // 3 + 10)
+
+    await s.upsert({"v1": {"content": content}})
+    await s.index_done_callback()
+
+    upserted = s._client.upsert.call_args.kwargs["data"][0]
+    assert len(upserted["content"].encode("utf-8")) <= MILVUS_MAX_VARCHAR_BYTES
+    upserted["content"].encode("utf-8").decode("utf-8")
+    assert embed.texts == [content]
+
+
+@pytest.mark.asyncio
+async def test_upsert_rejects_oversized_primary_id():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    long_id = "v" * 65
+
+    with pytest.raises(ValueError, match="primary keys cannot be truncated"):
+        await s.upsert({long_id: {"content": "hello"}})
+
+    assert s._pending_vector_docs == {}
+    assert embed.call_count == 0
+    s._client.upsert.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_upsert_rejects_oversized_entity_name():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed, meta_fields={"content", "entity_name"})
+    long_entity_name = "e" * 513
+
+    with pytest.raises(ValueError, match="identity fields cannot be truncated"):
+        await s.upsert({"ent-1": {"content": "hello", "entity_name": long_entity_name}})
+
+    assert s._pending_vector_docs == {}
+    assert embed.call_count == 0
+    s._client.upsert.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_upsert_rejects_oversized_relation_identity_fields():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed, namespace="relationships")
+    long_entity_id = "e" * 513
+
+    with pytest.raises(ValueError, match="identity fields cannot be truncated"):
+        await s.upsert(
+            {
+                "rel-1": {
+                    "content": "hello",
+                    "src_id": long_entity_id,
+                    "tgt_id": "B",
+                }
+            }
+        )
+
+    assert s._pending_vector_docs == {}
+    assert embed.call_count == 0
+    s._client.upsert.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_upsert_rejects_oversized_full_doc_id():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed, namespace="chunks", meta_fields={"content", "full_doc_id"})
+    long_doc_id = "d" * 65
+
+    with pytest.raises(ValueError, match="identity fields cannot be truncated"):
+        await s.upsert({"chunk-1": {"content": "hello", "full_doc_id": long_doc_id}})
+
+    assert s._pending_vector_docs == {}
+    assert embed.call_count == 0
+    s._client.upsert.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_upsert_truncates_oversized_source_id():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed, meta_fields={"content", "source_id"})
+    source_id = "s" * (MILVUS_MAX_VARCHAR_BYTES + 10)
+
+    await s.upsert({"v1": {"content": "hello", "source_id": source_id}})
+    await s.index_done_callback()
+
+    upserted = s._client.upsert.call_args.kwargs["data"][0]
+    assert len(upserted["source_id"].encode("utf-8")) == MILVUS_MAX_VARCHAR_BYTES
+
+
+@pytest.mark.asyncio
+async def test_upsert_truncates_source_id_on_separator_boundary():
+    # source_id is a <SEP>-joined list of chunk ids; truncation must drop whole
+    # ids at a separator boundary rather than leave a dangling partial id.
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed, meta_fields={"content", "source_id"})
+    # Each chunk id is comfortably sized; enough of them to overflow the limit.
+    chunk_id = "chunk-" + "a" * 50
+    chunk_ids = [chunk_id] * ((MILVUS_MAX_VARCHAR_BYTES // len(chunk_id)) + 5)
+    source_id = "<SEP>".join(chunk_ids)
+    assert len(source_id.encode("utf-8")) > MILVUS_MAX_VARCHAR_BYTES
+
+    await s.upsert({"v1": {"content": "hello", "source_id": source_id}})
+    await s.index_done_callback()
+
+    upserted = s._client.upsert.call_args.kwargs["data"][0]
+    stored = upserted["source_id"]
+    assert len(stored.encode("utf-8")) <= MILVUS_MAX_VARCHAR_BYTES
+    # No trailing partial id: every retained id is intact.
+    assert all(part == chunk_id for part in stored.split("<SEP>"))
+
+
+@pytest.mark.asyncio
+async def test_upsert_truncates_oversized_single_source_id_without_separator():
+    # A single id longer than the limit has no separator to back off to, so it
+    # falls back to the raw byte cut instead of dropping the value entirely.
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed, meta_fields={"content", "source_id"})
+    source_id = "s" * (MILVUS_MAX_VARCHAR_BYTES + 10)
+
+    await s.upsert({"v1": {"content": "hello", "source_id": source_id}})
+    await s.index_done_callback()
+
+    upserted = s._client.upsert.call_args.kwargs["data"][0]
+    assert len(upserted["source_id"].encode("utf-8")) == MILVUS_MAX_VARCHAR_BYTES
+
+
+@pytest.mark.asyncio
+async def test_repeated_upsert_same_id_embeds_once():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "first"}})
+    await s.upsert({"v1": {"content": "second"}})
+    await s.upsert({"v1": {"content": "third"}})
+    await s.index_done_callback()
+
+    assert embed.call_count == 1
+    # Only the latest content survives and was embedded.
+    assert embed.texts == ["third"]
+    s._client.upsert.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_deferred_embeddings_respect_batch_size():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._max_batch_size = 2
+
+    await s.upsert({f"v{i}": {"content": f"doc {i}"} for i in range(5)})
+    await s.index_done_callback()
+
+    # 5 docs / batch 2 → 3 batches → 3 embedding calls
+    assert embed.call_count == 3
+    assert [len(b) for b in embed.batches] == [2, 2, 1]
+
+
+@pytest.mark.asyncio
+async def test_get_vectors_by_ids_lazy_embed_then_reuse_in_flush():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}})
+
+    vectors = await s.get_vectors_by_ids(["v1"])
+    assert "v1" in vectors
+    assert embed.call_count == 1  # lazy embed inside get_vectors_by_ids
+    # The lazy-embedded vector is cached on the pending doc.
+    assert s._pending_vector_docs["v1"].vector is not None
+
+    await s.index_done_callback()
+    # Flush reused the cached vector — no extra embedding call.
+    assert embed.call_count == 1
+    s._client.upsert.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_flush_failure_keeps_buffer_and_no_double_embed_on_retry():
+    embed = CountingEmbeddingFunc(fail_times=1)  # first flush raises
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}})
+
+    with pytest.raises(RuntimeError, match="embedding failed"):
+        await s.index_done_callback()
+
+    # Buffer must remain so the next flush can retry.
+    assert "v1" in s._pending_vector_docs
+    assert s._pending_vector_docs["v1"].vector is None
+    s._client.upsert.assert_not_called()
+
+    # Second attempt succeeds; total embed calls is 2 (one failed + one ok),
+    # not 3 — the same content was retried exactly once.
+    await s.index_done_callback()
+    assert embed.call_count == 2
+    s._client.upsert.assert_called_once()
+    assert s._pending_vector_docs == {}
+
+
+@pytest.mark.asyncio
+async def test_server_upsert_failure_keeps_buffer():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._client.upsert.side_effect = RuntimeError("milvus down")
+
+    await s.upsert({"v1": {"content": "hello"}})
+
+    with pytest.raises(RuntimeError, match="milvus down"):
+        await s.index_done_callback()
+
+    # Embedding ran but server write failed; buffer must remain populated.
+    assert "v1" in s._pending_vector_docs
+    # Vector should be cached so retry doesn't re-embed.
+    assert s._pending_vector_docs["v1"].vector is not None
+
+    # On retry, no further embedding call; only the server write is reattempted.
+    s._client.upsert.side_effect = None
+    s._client.upsert.return_value = {"upsert_count": 1}
+    await s.index_done_callback()
+    assert embed.call_count == 1
+    assert s._pending_vector_docs == {}
+
+
+@pytest.mark.asyncio
+async def test_finalize_raises_when_buffer_unflushed():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._client.upsert.side_effect = RuntimeError("transient milvus error")
+
+    await s.upsert({"v1": {"content": "hello"}})
+
+    with pytest.raises(RuntimeError, match="finalize.*flush raised"):
+        await s.finalize()
+
+    # Buffer still populated — caller knows data was lost.
+    assert "v1" in s._pending_vector_docs
+
+
+@pytest.mark.asyncio
+async def test_delete_then_upsert_same_id_keeps_upsert():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.delete(["v1"])
+    assert "v1" in s._pending_vector_deletes
+
+    await s.upsert({"v1": {"content": "hello"}})
+    assert "v1" in s._pending_vector_docs
+    assert "v1" not in s._pending_vector_deletes
+
+    await s.index_done_callback()
+    s._client.upsert.assert_called_once()
+    s._client.delete.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_upsert_then_delete_same_id_keeps_delete():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}})
+    await s.delete(["v1"])
+    assert "v1" not in s._pending_vector_docs
+    assert "v1" in s._pending_vector_deletes
+
+    await s.index_done_callback()
+    # No upsert payload, only the delete batch.
+    s._client.upsert.assert_not_called()
+    s._client.delete.assert_called_once()
+    assert s._client.delete.call_args.kwargs["pks"] == ["v1"]
+
+
+@pytest.mark.asyncio
+async def test_delete_entity_relation_raises_on_server_failure():
+    """Server-side failure must bubble up — no log-and-swallow."""
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._client.query.return_value = [{"id": "rel1"}, {"id": "rel2"}]
+    s._client.delete.side_effect = RuntimeError("milvus delete failed")
+
+    with pytest.raises(RuntimeError, match="milvus delete failed"):
+        await s.delete_entity_relation("X")
+
+
+@pytest.mark.asyncio
+async def test_delete_entity_relation_prunes_pending_buffer():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert(
+        {
+            "rel-A-B": {"content": "A → B", "src_id": "A", "tgt_id": "B"},
+            "rel-C-D": {"content": "C → D", "src_id": "C", "tgt_id": "D"},
+        }
+    )
+    s._client.query.return_value = []  # no server-side hits
+
+    await s.delete_entity_relation("A")
+    # Pending doc whose src_id == A is pruned, the other survives.
+    assert "rel-A-B" not in s._pending_vector_docs
+    assert "rel-C-D" in s._pending_vector_docs
+
+
+@pytest.mark.asyncio
+async def test_delete_entity_relation_diverges_when_buffer_overwrites_persisted():
+    """Pins the deferred ↔ eager semantic divergence documented on
+    ``delete_entity_relation``.
+
+    Scenario: a persisted row ``rel-X-Y`` has ``src_id="X" / tgt_id="Y"``,
+    and a pending upsert is about to rewrite that same id so it would
+    instead carry ``src_id="A" / tgt_id="B"``. A call to
+    ``delete_entity_relation("A")`` arrives before the buffer is flushed.
+
+    Expected (deferred mode, current implementation):
+      * server-side filter ``src_id == "A" or tgt_id == "A"`` does NOT
+        match the persisted row (its src/tgt are still X/Y), so the
+        server-side delete is a no-op;
+      * the buffered upsert IS pruned (its buffered src/tgt match);
+      * net effect: persisted ``rel-X-Y`` (old values) survives and the
+        pending overwrite is lost.
+
+    Under eager ordering (upsert → flush → delete) the persisted row
+    would have been rewritten first and then matched by the filter, so
+    the final state would have been a deleted ``rel-X-Y``. This test
+    locks in the divergence so a future refactor can't silently change
+    it without touching the docstring.
+    """
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    # Buffered upsert rewriting an (assumed) already-persisted rel-X-Y
+    # so that its new src/tgt would match entity "A".
+    await s.upsert({"rel-X-Y": {"content": "A → B", "src_id": "A", "tgt_id": "B"}})
+    assert "rel-X-Y" in s._pending_vector_docs
+
+    # Server still sees the OLD persisted row (src_id="X" / tgt_id="Y"),
+    # so a filter on entity "A" returns nothing.
+    s._client.query.return_value = []
+
+    await s.delete_entity_relation("A")
+
+    # Buffered overwrite is pruned (matches buffered src/tgt view) …
+    assert "rel-X-Y" not in s._pending_vector_docs
+    # … but the server-side delete is not issued, because the filter
+    # didn't match the persisted row's actual src/tgt.
+    s._client.delete.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_delete_entity_relation_eager_ordering_matches_persisted():
+    """Counterpart to the divergence test: if the caller flushes before
+    invoking ``delete_entity_relation``, the persisted row reflects the
+    buffered overwrite and the server-side filter catches it.
+
+    This documents the recommended workaround called out in the
+    ``delete_entity_relation`` docstring: ``index_done_callback()`` first
+    when eager-equivalent semantics are required.
+    """
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"rel-X-Y": {"content": "A → B", "src_id": "A", "tgt_id": "B"}})
+    await s.index_done_callback()  # buffered upsert is now persisted
+    assert s._pending_vector_docs == {}
+    s._client.upsert.assert_called_once()
+
+    # With the row persisted, the server filter on entity "A" now hits.
+    s._client.query.return_value = [{"id": "rel-X-Y"}]
+
+    await s.delete_entity_relation("A")
+
+    s._client.delete.assert_called_once()
+    assert s._client.delete.call_args.kwargs["pks"] == ["rel-X-Y"]
+
+
+@pytest.mark.asyncio
+async def test_get_by_id_reads_pending_buffer_without_vector():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello", "entity_name": "E1"}})
+    doc = await s.get_by_id("v1")
+    assert doc is not None
+    assert doc["id"] == "v1"
+    assert doc.get("entity_name") == "E1"
+    assert "vector" not in doc
+    # Server was not queried because the buffer answered the read.
+    s._client.query.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_get_by_id_returns_none_for_pending_delete():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.delete(["v1"])
+    assert await s.get_by_id("v1") is None
+    s._client.query.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_env_workspace_override_shares_flush_lock(patch_namespace_lock):
+    """Two instances whose final_namespace collides must share the flush lock."""
+    cache = patch_namespace_lock
+    embed = CountingEmbeddingFunc()
+
+    with patch.dict(os.environ, {"MILVUS_WORKSPACE": "shared_ws"}, clear=False):
+        a = _make_storage(embed, workspace="caller_a")
+        b = _make_storage(embed, workspace="caller_b")
+        assert (
+            a.final_namespace == b.final_namespace == "shared_ws_entities_mock_embed_8d"
+        )
+        assert a._flush_lock is b._flush_lock
+        # Sanity: only one lock object was cached for that final_namespace.
+        assert len([k for k in cache if k[0] == a.final_namespace]) == 1
+
+
+@pytest.mark.asyncio
+async def test_distinct_namespaces_get_independent_locks(patch_namespace_lock):
+    """Different final_namespaces must NOT share a lock."""
+    embed = CountingEmbeddingFunc()
+
+    # Two instances with no env override and different workspaces produce
+    # different final_namespaces ("a_entities" vs "b_entities").
+    a = _make_storage(embed, workspace="a")
+    b = _make_storage(embed, workspace="b")
+    assert a.final_namespace != b.final_namespace
+    assert a._flush_lock is not b._flush_lock
+
+
+@pytest.mark.asyncio
+async def test_mixed_upsert_and_delete_in_single_flush():
+    """A flush carrying both pending upserts and pending deletes (on disjoint
+    ids) must dispatch one server upsert and one server delete in a single
+    pass, then clear both buffers."""
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"a": {"content": "alpha"}})
+    await s.delete(["b"])
+
+    assert set(s._pending_vector_docs.keys()) == {"a"}
+    assert s._pending_vector_deletes == {"b"}
+
+    await s.index_done_callback()
+
+    s._client.upsert.assert_called_once()
+    upsert_kwargs = s._client.upsert.call_args.kwargs
+    assert {row["id"] for row in upsert_kwargs["data"]} == {"a"}
+
+    s._client.delete.assert_called_once()
+    assert s._client.delete.call_args.kwargs["pks"] == ["b"]
+
+    # Both buffers cleared after a successful flush.
+    assert s._pending_vector_docs == {}
+    assert s._pending_vector_deletes == set()
+
+
+@pytest.mark.asyncio
+async def test_finalize_clean_flush_no_raise():
+    """Happy-path counterpart to test_finalize_raises_when_buffer_unflushed:
+    a successful flush during finalize() must leave both buffers empty and
+    must not raise."""
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}})
+    await s.delete(["v2"])
+
+    await s.finalize()  # must not raise
+
+    s._client.upsert.assert_called_once()
+    s._client.delete.assert_called_once()
+    assert s._pending_vector_docs == {}
+    assert s._pending_vector_deletes == set()
+
+
+@pytest.mark.asyncio
+async def test_drop_clears_pending_buffers():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._client.has_collection.return_value = False  # skip drop_collection call
+    # Stub out the recreate path to avoid hitting MilvusIndexConfig logic.
+    with (
+        patch.object(s, "_create_collection_with_schema"),
+        patch.object(s, "_ensure_collection_loaded"),
+    ):
+        await s.upsert({"v1": {"content": "hello"}})
+        await s.delete(["v2"])
+        assert s._pending_vector_docs and s._pending_vector_deletes
+
+        result = await s.drop()
+        assert result["status"] == "success"
+        assert s._pending_vector_docs == {}
+        assert s._pending_vector_deletes == set()
+
+
+@pytest.mark.asyncio
+async def test_drop_recreates_empty_without_legacy_migration():
+    # drop() must leave the collection EMPTY. Recreating via
+    # _create_collection_if_not_exist would re-run the legacy->suffixed
+    # migration (the legacy collection is intentionally kept after migration),
+    # pulling the just-dropped rows back in and forcing a needless full
+    # migration on every rebuild/clear. Regression for that path.
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._client.has_collection.return_value = True
+
+    with (
+        patch.object(s, "_create_collection_if_not_exist") as recreate_via_migration,
+        patch.object(s, "_create_collection_with_schema") as create_empty,
+        patch.object(s, "_ensure_collection_loaded") as load,
+    ):
+        result = await s.drop()
+
+    assert result["status"] == "success"
+    s._client.drop_collection.assert_called_once_with(s.final_namespace)
+    create_empty.assert_called_once_with(s.final_namespace)
+    load.assert_called_once_with()
+    # Never the migration-capable path.
+    recreate_via_migration.assert_not_called()
+    s._client.query_iterator.assert_not_called()
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_drop_pending_index_ops_clears_buffers():
+    """On an internal-error abort the pipeline calls drop_pending_index_ops to
+    discard buffered upserts/deletes without flushing them (PR #3187)."""
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    await s.upsert({"v1": {"content": "x"}, "v2": {"content": "y"}})
+    s._pending_vector_deletes.add("old-id")
+    assert s._pending_vector_docs
+    await s.drop_pending_index_ops()
+    assert not s._pending_vector_docs
+    assert not s._pending_vector_deletes
+
+
+@pytest.mark.asyncio
+async def test_finalize_closes_milvus_client():
+    """finalize() must release the Milvus gRPC channel instead of leaving it
+    for GC — mirroring the close-on-release pattern of the other server-backed
+    storages and the in-file ``_rebuild_milvus_client``."""
+    s = _make_storage(MockEmbeddingFunc())
+    client = s._client
+    assert client is not None
+
+    # No pending writes → finalize must not raise.
+    await s.finalize()
+
+    client.close.assert_called_once()

tests/kg/milvus_impl/test_milvus_index_config.py

@@ -0,0 +1,415 @@
+"""
+Tests for Milvus index configuration
+
+This test suite validates the MilvusIndexConfig class and its integration
+with MilvusVectorDBStorage.
+"""
+
+import pytest
+import os
+from unittest.mock import patch, MagicMock
+from lightrag.kg.milvus_impl import (
+    MilvusIndexConfig,
+    SUPPORTED_INDEX_TYPES,
+    SUPPORTED_METRIC_TYPES,
+    SUPPORTED_SQ_TYPES,
+    SUPPORTED_REFINE_TYPES,
+)
+
+
+class TestMilvusIndexConfig:
+    """MilvusIndexConfig unit tests"""
+
+    def test_default_values(self):
+        """Test default configuration"""
+        config = MilvusIndexConfig()
+        assert config.index_type == "AUTOINDEX"
+        assert config.metric_type == "COSINE"
+        assert config.hnsw_m == 16
+        assert config.hnsw_ef_construction == 360
+        assert config.hnsw_ef == 200
+        assert config.sq_type == "SQ8"
+        assert not config.sq_refine
+        assert config.sq_refine_type == "FP32"
+        assert config.sq_refine_k == 10
+        assert config.ivf_nlist == 1024
+        assert config.ivf_nprobe == 16
+
+    def test_env_override(self):
+        """Test environment variable override"""
+        with patch.dict(
+            os.environ,
+            {
+                "MILVUS_INDEX_TYPE": "HNSW",
+                "MILVUS_METRIC_TYPE": "L2",
+                "MILVUS_HNSW_M": "64",
+            },
+        ):
+            config = MilvusIndexConfig()
+            assert config.index_type == "HNSW"
+            assert config.metric_type == "L2"
+            assert config.hnsw_m == 64
+
+    def test_init_param_priority(self):
+        """Test initialization parameter priority over environment variables"""
+        with patch.dict(os.environ, {"MILVUS_INDEX_TYPE": "IVF_FLAT"}):
+            config = MilvusIndexConfig(index_type="HNSW")
+            assert config.index_type == "HNSW"  # Init param takes precedence
+
+    def test_case_insensitive_index_type(self):
+        """Test that index type is case insensitive"""
+        config = MilvusIndexConfig(index_type="hnsw")
+        assert config.index_type == "HNSW"
+
+    def test_case_insensitive_metric_type(self):
+        """Test that metric type is case insensitive"""
+        config = MilvusIndexConfig(metric_type="cosine")
+        assert config.metric_type == "COSINE"
+
+    def test_invalid_index_type(self):
+        """Test invalid index type raises exception"""
+        with pytest.raises(ValueError, match="Unsupported index type"):
+            MilvusIndexConfig(index_type="INVALID_INDEX")
+
+    def test_invalid_metric_type(self):
+        """Test invalid metric type raises exception"""
+        with pytest.raises(ValueError, match="Unsupported metric type"):
+            MilvusIndexConfig(metric_type="INVALID_METRIC")
+
+    def test_invalid_hnsw_m_range_low(self):
+        """Test HNSW M parameter range validation (too low)"""
+        with pytest.raises(ValueError, match="hnsw_m must be in"):
+            MilvusIndexConfig(hnsw_m=1)  # Less than 2
+
+    def test_invalid_hnsw_m_range_high(self):
+        """Test HNSW M parameter range validation (too high)"""
+        with pytest.raises(ValueError, match="hnsw_m must be in"):
+            MilvusIndexConfig(hnsw_m=3000)  # Greater than 2048
+
+    def test_valid_hnsw_m_boundary(self):
+        """Test HNSW M parameter boundary values"""
+        config_low = MilvusIndexConfig(hnsw_m=2)
+        assert config_low.hnsw_m == 2
+
+        config_high = MilvusIndexConfig(hnsw_m=2048)
+        assert config_high.hnsw_m == 2048
+
+    def test_invalid_hnsw_ef_construction(self):
+        """Test HNSW efConstruction validation"""
+        with pytest.raises(ValueError, match="hnsw_ef_construction must be"):
+            MilvusIndexConfig(hnsw_ef_construction=0)
+
+    def test_invalid_ivf_nlist_low(self):
+        """Test IVF nlist parameter range validation (too low)"""
+        with pytest.raises(ValueError, match="ivf_nlist must be in"):
+            MilvusIndexConfig(ivf_nlist=0)
+
+    def test_invalid_ivf_nlist_high(self):
+        """Test IVF nlist parameter range validation (too high)"""
+        with pytest.raises(ValueError, match="ivf_nlist must be in"):
+            MilvusIndexConfig(ivf_nlist=70000)
+
+    def test_invalid_sq_type(self):
+        """Test invalid sq_type"""
+        with pytest.raises(ValueError, match="Unsupported sq_type"):
+            MilvusIndexConfig(index_type="HNSW_SQ", sq_type="INVALID")
+
+    def test_invalid_refine_type(self):
+        """Test invalid refine_type"""
+        with pytest.raises(ValueError, match="Unsupported refine_type"):
+            MilvusIndexConfig(
+                index_type="HNSW_SQ", sq_refine=True, sq_refine_type="INVALID"
+            )
+
+    def test_version_validation_hnsw_sq_pass(self):
+        """Test HNSW_SQ version validation passes with valid versions"""
+        config = MilvusIndexConfig(index_type="HNSW_SQ")
+
+        # Version meets requirement
+        config.validate_milvus_version("2.6.8")  # Exactly required
+        config.validate_milvus_version("2.6.9")  # Above requirement
+        config.validate_milvus_version("2.7.0")  # Higher version
+
+    def test_version_validation_hnsw_sq_fail(self):
+        """Test HNSW_SQ version validation fails with invalid versions"""
+        config = MilvusIndexConfig(index_type="HNSW_SQ")
+
+        # Version does not meet requirement
+        with pytest.raises(ValueError, match="HNSW_SQ requires Milvus 2.6.8"):
+            config.validate_milvus_version("2.6.7")  # Below 2.6.8
+
+        with pytest.raises(ValueError, match="HNSW_SQ requires Milvus 2.6.8"):
+            config.validate_milvus_version("2.5.0")  # Much lower
+
+    def test_version_validation_hnsw_sq_with_sq4u(self):
+        """Test HNSW_SQ + SQ4U version validation"""
+        config = MilvusIndexConfig(index_type="HNSW_SQ", sq_type="SQ4U")
+
+        # Passes with valid version
+        config.validate_milvus_version("2.6.9")
+
+        # Fails with invalid version
+        with pytest.raises(ValueError, match="HNSW_SQ requires Milvus 2.6.8"):
+            config.validate_milvus_version("2.6.0")
+
+    def test_version_validation_hnsw_no_requirement(self):
+        """Test normal HNSW has no version restriction"""
+        config = MilvusIndexConfig(index_type="HNSW")
+
+        # No version restriction
+        config.validate_milvus_version("2.4.0")  # Lower version OK
+        config.validate_milvus_version("2.6.9")  # Higher version OK
+
+    def test_version_validation_with_dev_suffix(self):
+        """Test version validation handles dev suffixes"""
+        config = MilvusIndexConfig(index_type="HNSW_SQ")
+
+        # Should handle "2.6.9-dev" format
+        config.validate_milvus_version("2.6.9-dev")
+
+    def test_build_index_params_autoindex(self):
+        """Test AUTOINDEX generates explicit index parameters"""
+        config = MilvusIndexConfig(index_type="AUTOINDEX")
+        mock_index_params = MagicMock()
+
+        result = config.build_index_params(mock_index_params)
+        assert result is mock_index_params
+        mock_index_params.add_index.assert_called_once_with(
+            field_name="vector",
+            index_type="AUTOINDEX",
+            metric_type="COSINE",
+            params={},
+        )
+
+    def test_build_index_params_hnsw(self):
+        """Test HNSW index parameters construction"""
+        config = MilvusIndexConfig(
+            index_type="HNSW",
+            metric_type="COSINE",
+            hnsw_m=32,
+            hnsw_ef_construction=256,
+        )
+
+        mock_index_params = MagicMock()
+
+        config.build_index_params(mock_index_params)
+
+        mock_index_params.add_index.assert_called_once()
+        call_kwargs = mock_index_params.add_index.call_args[1]
+        assert call_kwargs["index_type"] == "HNSW"
+        assert call_kwargs["metric_type"] == "COSINE"
+        assert call_kwargs["params"]["M"] == 32
+        assert call_kwargs["params"]["efConstruction"] == 256
+
+    def test_build_index_params_hnsw_sq(self):
+        """Test HNSW_SQ index parameters construction"""
+        config = MilvusIndexConfig(
+            index_type="HNSW_SQ",
+            sq_type="SQ8",
+            sq_refine=True,
+            sq_refine_type="FP32",
+        )
+
+        mock_index_params = MagicMock()
+
+        config.build_index_params(mock_index_params)
+
+        call_kwargs = mock_index_params.add_index.call_args[1]
+        assert call_kwargs["index_type"] == "HNSW_SQ"
+        assert call_kwargs["params"]["sq_type"] == "SQ8"
+        assert call_kwargs["params"]["refine"] is True
+        assert call_kwargs["params"]["refine_type"] == "FP32"
+
+    def test_build_index_params_hnsw_sq_no_refine(self):
+        """Test HNSW_SQ without refinement"""
+        config = MilvusIndexConfig(index_type="HNSW_SQ", sq_type="SQ8", sq_refine=False)
+
+        mock_index_params = MagicMock()
+
+        config.build_index_params(mock_index_params)
+
+        call_kwargs = mock_index_params.add_index.call_args[1]
+        assert call_kwargs["index_type"] == "HNSW_SQ"
+        assert call_kwargs["params"]["sq_type"] == "SQ8"
+        assert "refine" not in call_kwargs["params"]
+        assert "refine_type" not in call_kwargs["params"]
+
+    def test_build_index_params_ivf_flat(self):
+        """Test IVF_FLAT index parameters construction"""
+        config = MilvusIndexConfig(index_type="IVF_FLAT", ivf_nlist=2048)
+
+        mock_index_params = MagicMock()
+
+        config.build_index_params(mock_index_params)
+
+        call_kwargs = mock_index_params.add_index.call_args[1]
+        assert call_kwargs["index_type"] == "IVF_FLAT"
+        assert call_kwargs["params"]["nlist"] == 2048
+
+    def test_build_index_params_with_none(self):
+        """Test that RuntimeError is raised when index_params is None for custom types"""
+        config = MilvusIndexConfig(index_type="HNSW")
+
+        # Pass None to simulate when compatibility helper returns None
+        with pytest.raises(RuntimeError, match="IndexParams not available"):
+            config.build_index_params(None)
+
+    def test_build_search_params_hnsw(self):
+        """Test HNSW search parameters construction"""
+        config = MilvusIndexConfig(index_type="HNSW", hnsw_ef=150)
+        params = config.build_search_params()
+        assert params["params"]["ef"] == 150
+
+    def test_build_search_params_hnsw_sq_with_refine(self):
+        """Test HNSW_SQ with refinement search parameters"""
+        config = MilvusIndexConfig(
+            index_type="HNSW_SQ", hnsw_ef=200, sq_refine=True, sq_refine_k=20
+        )
+        params = config.build_search_params()
+        assert params["params"]["ef"] == 200
+        assert params["params"]["refine_k"] == 20
+
+    def test_build_search_params_hnsw_sq_no_refine(self):
+        """Test HNSW_SQ without refinement search parameters"""
+        config = MilvusIndexConfig(index_type="HNSW_SQ", hnsw_ef=200, sq_refine=False)
+        params = config.build_search_params()
+        assert params["params"]["ef"] == 200
+        assert "refine_k" not in params["params"]
+
+    def test_build_search_params_ivf(self):
+        """Test IVF search parameters construction"""
+        config = MilvusIndexConfig(index_type="IVF_FLAT", ivf_nprobe=32)
+        params = config.build_search_params()
+        assert params["params"]["nprobe"] == 32
+
+    def test_build_search_params_autoindex(self):
+        """Test AUTOINDEX search parameters (empty)"""
+        config = MilvusIndexConfig(index_type="AUTOINDEX")
+        params = config.build_search_params()
+        assert params == {}
+
+    def test_to_dict_hnsw(self):
+        """Test configuration export for HNSW"""
+        config = MilvusIndexConfig(index_type="HNSW")
+        d = config.to_dict()
+        assert d["index_type"] == "HNSW"
+        assert d["hnsw_m"] == 16
+        assert d["sq_type"] is None  # Not HNSW_SQ
+        assert d["ivf_nlist"] is None  # Not IVF
+
+    def test_to_dict_hnsw_sq(self):
+        """Test configuration export for HNSW_SQ"""
+        config = MilvusIndexConfig(index_type="HNSW_SQ", sq_type="SQ8")
+        d = config.to_dict()
+        assert d["index_type"] == "HNSW_SQ"
+        assert d["sq_type"] == "SQ8"
+        assert d["ivf_nlist"] is None
+
+    def test_to_dict_ivf(self):
+        """Test configuration export for IVF"""
+        config = MilvusIndexConfig(index_type="IVF_FLAT")
+        d = config.to_dict()
+        assert d["index_type"] == "IVF_FLAT"
+        assert d["ivf_nlist"] == 1024
+        assert d["sq_type"] is None
+
+    def test_env_bool_parsing(self):
+        """Test boolean environment variable parsing"""
+        with patch.dict(os.environ, {"MILVUS_HNSW_SQ_REFINE": "true"}):
+            config = MilvusIndexConfig(index_type="HNSW_SQ")
+            assert config.sq_refine is True
+
+        with patch.dict(os.environ, {"MILVUS_HNSW_SQ_REFINE": "false"}):
+            config = MilvusIndexConfig(index_type="HNSW_SQ")
+            assert not config.sq_refine
+
+        with patch.dict(os.environ, {"MILVUS_HNSW_SQ_REFINE": "1"}):
+            config = MilvusIndexConfig(index_type="HNSW_SQ")
+            assert config.sq_refine is True
+
+        with patch.dict(os.environ, {"MILVUS_HNSW_SQ_REFINE": "0"}):
+            config = MilvusIndexConfig(index_type="HNSW_SQ")
+            assert not config.sq_refine
+
+    def test_env_int_parsing_invalid(self):
+        """Test integer environment variable parsing with invalid value"""
+        with patch.dict(os.environ, {"MILVUS_HNSW_M": "invalid"}):
+            config = MilvusIndexConfig()
+            assert config.hnsw_m == 16  # Falls back to default (Milvus 2.4+)
+
+    def test_all_index_types_supported(self):
+        """Test all supported index types can be configured"""
+        for index_type in SUPPORTED_INDEX_TYPES:
+            if index_type == "HNSW_SQ":
+                # HNSW_SQ requires special parameters
+                config = MilvusIndexConfig(index_type=index_type, sq_type="SQ8")
+            else:
+                config = MilvusIndexConfig(index_type=index_type)
+            assert config.index_type == index_type
+
+    def test_all_metric_types_supported(self):
+        """Test all supported metric types can be configured"""
+        for metric_type in SUPPORTED_METRIC_TYPES:
+            config = MilvusIndexConfig(metric_type=metric_type)
+            assert config.metric_type == metric_type
+
+    def test_all_sq_types_supported(self):
+        """Test all supported sq_types can be configured"""
+        for sq_type in SUPPORTED_SQ_TYPES:
+            config = MilvusIndexConfig(index_type="HNSW_SQ", sq_type=sq_type)
+            assert config.sq_type == sq_type
+
+    def test_all_refine_types_supported(self):
+        """Test all supported refine_types can be configured"""
+        for refine_type in SUPPORTED_REFINE_TYPES:
+            config = MilvusIndexConfig(
+                index_type="HNSW_SQ", sq_refine=True, sq_refine_type=refine_type
+            )
+            assert config.sq_refine_type == refine_type
+
+    def test_get_config_field_names(self):
+        """Test get_config_field_names() returns all dataclass fields"""
+        field_names = MilvusIndexConfig.get_config_field_names()
+
+        # Check that it's a set
+        assert isinstance(field_names, set)
+
+        # Check that all expected fields are present
+        expected_fields = {
+            "index_type",
+            "metric_type",
+            "hnsw_m",
+            "hnsw_ef_construction",
+            "hnsw_ef",
+            "sq_type",
+            "sq_refine",
+            "sq_refine_type",
+            "sq_refine_k",
+            "ivf_nlist",
+            "ivf_nprobe",
+        }
+        assert field_names == expected_fields
+
+    def test_get_config_field_names_single_source_of_truth(self):
+        """Test that get_config_field_names() provides single source of truth for configuration parameters"""
+        # This test ensures that when we add new fields to MilvusIndexConfig,
+        # they are automatically included in get_config_field_names()
+        # without needing to update hardcoded lists elsewhere
+
+        from dataclasses import fields as dataclass_fields
+
+        # Get fields directly from dataclass
+        direct_fields = {f.name for f in dataclass_fields(MilvusIndexConfig)}
+
+        # Get fields via the method
+        method_fields = MilvusIndexConfig.get_config_field_names()
+
+        # They should be identical
+        assert direct_fields == method_fields, (
+            f"Method should return same fields as dataclass. "
+            f"Difference: {direct_fields.symmetric_difference(method_fields)}"
+        )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/kg/milvus_impl/test_milvus_kwargs_bridge.py

@@ -0,0 +1,296 @@
+"""
+Tests for bridging vector_db_storage_cls_kwargs to MilvusIndexConfig
+
+This test suite validates that MilvusIndexConfig parameters can be passed
+through vector_db_storage_cls_kwargs and that backward compatibility is maintained.
+"""
+
+import pytest
+from unittest.mock import patch, MagicMock
+from lightrag.kg.milvus_impl import MilvusVectorDBStorage
+
+
+@pytest.mark.offline
+class TestMilvusKwargsParameterBridge:
+    """Test parameter bridging from vector_db_storage_cls_kwargs to MilvusIndexConfig"""
+
+    def test_kwargs_to_index_config_basic(self):
+        """Test that basic HNSW parameters are passed from kwargs to MilvusIndexConfig"""
+        # Mock the embedding function
+        mock_embedding_func = MagicMock()
+        mock_embedding_func.embedding_dim = 128
+
+        # Create storage instance with custom index config parameters in kwargs
+        storage = MilvusVectorDBStorage(
+            namespace="test_entities",
+            workspace="test_workspace",
+            global_config={
+                "embedding_batch_num": 100,
+                "vector_db_storage_cls_kwargs": {
+                    "cosine_better_than_threshold": 0.3,
+                    "hnsw_m": 32,
+                    "hnsw_ef": 256,
+                    "hnsw_ef_construction": 300,
+                },
+            },
+            embedding_func=mock_embedding_func,
+            meta_fields=set(),
+        )
+
+        # Verify that parameters were passed to index_config
+        assert storage.index_config.hnsw_m == 32
+        assert storage.index_config.hnsw_ef == 256
+        assert storage.index_config.hnsw_ef_construction == 300
+
+    def test_kwargs_to_index_config_index_and_metric_types(self):
+        """Test that index_type and metric_type are passed from kwargs"""
+        mock_embedding_func = MagicMock()
+        mock_embedding_func.embedding_dim = 128
+
+        storage = MilvusVectorDBStorage(
+            namespace="test_entities",
+            workspace="test_workspace",
+            global_config={
+                "embedding_batch_num": 100,
+                "vector_db_storage_cls_kwargs": {
+                    "cosine_better_than_threshold": 0.3,
+                    "index_type": "IVF_FLAT",
+                    "metric_type": "L2",
+                    "ivf_nlist": 2048,
+                },
+            },
+            embedding_func=mock_embedding_func,
+            meta_fields=set(),
+        )
+
+        # Verify that parameters were passed to index_config
+        assert storage.index_config.index_type == "IVF_FLAT"
+        assert storage.index_config.metric_type == "L2"
+        assert storage.index_config.ivf_nlist == 2048
+
+    def test_kwargs_to_index_config_sq_parameters(self):
+        """Test that HNSW_SQ parameters are passed from kwargs"""
+        mock_embedding_func = MagicMock()
+        mock_embedding_func.embedding_dim = 128
+
+        storage = MilvusVectorDBStorage(
+            namespace="test_entities",
+            workspace="test_workspace",
+            global_config={
+                "embedding_batch_num": 100,
+                "vector_db_storage_cls_kwargs": {
+                    "cosine_better_than_threshold": 0.3,
+                    "index_type": "HNSW_SQ",
+                    "sq_type": "SQ8",
+                    "sq_refine": True,
+                    "sq_refine_type": "FP16",
+                    "sq_refine_k": 20,
+                },
+            },
+            embedding_func=mock_embedding_func,
+            meta_fields=set(),
+        )
+
+        # Verify that parameters were passed to index_config
+        assert storage.index_config.index_type == "HNSW_SQ"
+        assert storage.index_config.sq_type == "SQ8"
+        assert storage.index_config.sq_refine is True
+        assert storage.index_config.sq_refine_type == "FP16"
+        assert storage.index_config.sq_refine_k == 20
+
+    def test_backward_compatibility_no_index_params(self):
+        """Test backward compatibility when no index parameters are provided in kwargs"""
+        mock_embedding_func = MagicMock()
+        mock_embedding_func.embedding_dim = 128
+
+        # Create storage without any index config parameters in kwargs
+        storage = MilvusVectorDBStorage(
+            namespace="test_entities",
+            workspace="test_workspace",
+            global_config={
+                "embedding_batch_num": 100,
+                "vector_db_storage_cls_kwargs": {
+                    "cosine_better_than_threshold": 0.3,
+                },
+            },
+            embedding_func=mock_embedding_func,
+            meta_fields=set(),
+        )
+
+        # Verify that default values are used (from environment variables or defaults)
+        # Defaults aligned with Milvus 2.4+ official documentation
+        assert storage.index_config.index_type == "AUTOINDEX"  # Default
+        assert storage.index_config.metric_type == "COSINE"  # Default
+        assert storage.index_config.hnsw_m == 16  # Default (Milvus 2.4+)
+        assert storage.index_config.hnsw_ef_construction == 360  # Default (Milvus 2.4+)
+
+    def test_kwargs_params_override_environment_variables(self):
+        """Test that kwargs parameters take precedence over environment variables"""
+        mock_embedding_func = MagicMock()
+        mock_embedding_func.embedding_dim = 128
+
+        # Set environment variables
+        with patch.dict(
+            "os.environ",
+            {
+                "MILVUS_INDEX_TYPE": "IVF_FLAT",
+                "MILVUS_HNSW_M": "16",
+            },
+        ):
+            # Create storage with kwargs parameters that should override env vars
+            storage = MilvusVectorDBStorage(
+                namespace="test_entities",
+                workspace="test_workspace",
+                global_config={
+                    "embedding_batch_num": 100,
+                    "vector_db_storage_cls_kwargs": {
+                        "cosine_better_than_threshold": 0.3,
+                        "index_type": "HNSW",
+                        "hnsw_m": 64,
+                    },
+                },
+                embedding_func=mock_embedding_func,
+                meta_fields=set(),
+            )
+
+            # Verify that kwargs parameters override environment variables
+            assert (
+                storage.index_config.index_type == "HNSW"
+            )  # From kwargs, not IVF_FLAT
+            assert storage.index_config.hnsw_m == 64  # From kwargs, not 16
+
+    def test_non_index_params_ignored(self):
+        """Test that non-index-config parameters in kwargs are ignored"""
+        mock_embedding_func = MagicMock()
+        mock_embedding_func.embedding_dim = 128
+
+        storage = MilvusVectorDBStorage(
+            namespace="test_entities",
+            workspace="test_workspace",
+            global_config={
+                "embedding_batch_num": 100,
+                "vector_db_storage_cls_kwargs": {
+                    "cosine_better_than_threshold": 0.3,
+                    "hnsw_m": 32,
+                    "some_other_param": "ignored",  # Should be ignored
+                    "another_param": 123,  # Should be ignored
+                },
+            },
+            embedding_func=mock_embedding_func,
+            meta_fields=set(),
+        )
+
+        # Verify that valid parameter was passed
+        assert storage.index_config.hnsw_m == 32
+        # Verify that invalid parameters were ignored (no AttributeError)
+        assert not hasattr(storage.index_config, "some_other_param")
+        assert not hasattr(storage.index_config, "another_param")
+
+    def test_raganything_framework_integration_scenario(self):
+        """Test configuration passing through frameworks like RAGAnything
+
+        This test validates the use case where a framework (like RAGAnything)
+        sits on top of LightRAG and needs to pass Milvus index configuration
+        through to LightRAG without modifying environment variables.
+
+        The framework can pass all index config parameters via
+        vector_db_storage_cls_kwargs, and they will be properly extracted
+        and applied to MilvusIndexConfig.
+        """
+        mock_embedding_func = MagicMock()
+        mock_embedding_func.embedding_dim = 128
+
+        # Simulate RAGAnything framework passing configuration to LightRAG
+        # All index configuration parameters are passed through kwargs
+        framework_config = {
+            "embedding_batch_num": 100,
+            "vector_db_storage_cls_kwargs": {
+                # Required for vector storage
+                "cosine_better_than_threshold": 0.2,
+                # Milvus index configuration - all parameters supported
+                "index_type": "HNSW",
+                "metric_type": "L2",
+                "hnsw_m": 48,
+                "hnsw_ef_construction": 400,
+                "hnsw_ef": 200,
+                # Framework-specific parameters (should be ignored by Milvus)
+                "framework_version": "1.0.0",
+                "custom_setting": "value",
+            },
+        }
+
+        # Create storage instance with framework configuration
+        storage = MilvusVectorDBStorage(
+            namespace="test_entities",
+            workspace="raganything_workspace",
+            global_config=framework_config,
+            embedding_func=mock_embedding_func,
+            meta_fields=set(),
+        )
+
+        # Verify all Milvus parameters were correctly extracted and applied
+        assert storage.index_config.index_type == "HNSW"
+        assert storage.index_config.metric_type == "L2"
+        assert storage.index_config.hnsw_m == 48
+        assert storage.index_config.hnsw_ef_construction == 400
+        assert storage.index_config.hnsw_ef == 200
+
+        # Verify framework-specific parameters were ignored
+        assert not hasattr(storage.index_config, "framework_version")
+        assert not hasattr(storage.index_config, "custom_setting")
+
+        # Verify workspace isolation is maintained
+        assert storage.workspace == "raganything_workspace"
+
+    def test_all_milvus_parameters_supported_via_kwargs(self):
+        """Test that all 11 MilvusIndexConfig parameters can be configured via kwargs
+
+        This comprehensive test ensures that every single index configuration
+        parameter defined in MilvusIndexConfig can be passed through
+        vector_db_storage_cls_kwargs, which is critical for framework integration.
+        """
+        mock_embedding_func = MagicMock()
+        mock_embedding_func.embedding_dim = 128
+
+        # Pass ALL 11 MilvusIndexConfig parameters via kwargs
+        storage = MilvusVectorDBStorage(
+            namespace="test_entities",
+            workspace="test_workspace",
+            global_config={
+                "embedding_batch_num": 100,
+                "vector_db_storage_cls_kwargs": {
+                    "cosine_better_than_threshold": 0.3,
+                    # All 11 MilvusIndexConfig parameters
+                    "index_type": "HNSW_SQ",
+                    "metric_type": "IP",
+                    "hnsw_m": 64,
+                    "hnsw_ef_construction": 512,
+                    "hnsw_ef": 256,
+                    "sq_type": "SQ8",
+                    "sq_refine": True,
+                    "sq_refine_type": "FP16",
+                    "sq_refine_k": 30,
+                    "ivf_nlist": 4096,
+                    "ivf_nprobe": 64,
+                },
+            },
+            embedding_func=mock_embedding_func,
+            meta_fields=set(),
+        )
+
+        # Verify EVERY parameter was correctly applied
+        assert storage.index_config.index_type == "HNSW_SQ"
+        assert storage.index_config.metric_type == "IP"
+        assert storage.index_config.hnsw_m == 64
+        assert storage.index_config.hnsw_ef_construction == 512
+        assert storage.index_config.hnsw_ef == 256
+        assert storage.index_config.sq_type == "SQ8"
+        assert storage.index_config.sq_refine is True
+        assert storage.index_config.sq_refine_type == "FP16"
+        assert storage.index_config.sq_refine_k == 30
+        assert storage.index_config.ivf_nlist == 4096
+        assert storage.index_config.ivf_nprobe == 64
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/kg/milvus_impl/test_milvus_migration_memory.py

@@ -0,0 +1,217 @@
+"""
+Tests for Milvus schema-migration memory back-pressure.
+
+The bulk copy inserts the whole source collection into a temp collection with
+no client-side throttle, so the Milvus data node accumulates insert-buffer
+segments until its auto-flush catches up — a large migration can exhaust
+server memory. These tests cover the periodic/final flush that bounds that
+buffer and the optional inter-batch throttle.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch
+from pymilvus import MilvusException
+
+from lightrag.kg.milvus_impl import MilvusVectorDBStorage
+
+
+class _EmbeddingFunc:
+    def __init__(self, dim=128, model_name="text-embedding-3-small"):
+        self.embedding_dim = dim
+        self.model_name = model_name
+
+
+def _make_model_storage(namespace="entities", workspace="test_workspace", dim=128):
+    return MilvusVectorDBStorage(
+        namespace=namespace,
+        workspace=workspace,
+        global_config={
+            "embedding_batch_num": 100,
+            "vector_db_storage_cls_kwargs": {
+                "cosine_better_than_threshold": 0.3,
+            },
+        },
+        embedding_func=_EmbeddingFunc(dim=dim),
+        meta_fields=set(),
+    )
+
+
+def _wire_collection_state(storage, collections):
+    storage._client = MagicMock()
+
+    def has_collection(collection_name):
+        return collection_name in collections
+
+    def create_collection(collection_name, schema):
+        collections.add(collection_name)
+
+    def drop_collection(collection_name):
+        collections.discard(collection_name)
+
+    def rename_collection(source, target):
+        collections.discard(source)
+        collections.add(target)
+
+    storage._client.has_collection.side_effect = has_collection
+    storage._client.create_collection.side_effect = create_collection
+    storage._client.drop_collection.side_effect = drop_collection
+    storage._client.rename_collection.side_effect = rename_collection
+    return storage._client
+
+
+def _wire_iterator(client, batches):
+    """query_iterator yields the given list of row-batches then []."""
+
+    def make_iterator(**kwargs):
+        iterator = MagicMock()
+        iterator.next.side_effect = list(batches) + [[]]
+        return iterator
+
+    client.query_iterator.side_effect = make_iterator
+
+
+def _rows(n, start=0):
+    return [
+        {"id": f"ent-{i}", "vector": [0.0] * 128, "content": "body"}
+        for i in range(start, start + n)
+    ]
+
+
+@pytest.mark.offline
+class TestMigrationFlushBackpressure:
+    def test_periodic_and_final_flush_by_row_interval(self):
+        storage = _make_model_storage()
+        storage._migration_flush_interval_rows = 3
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        temp = f"{storage.final_namespace}_temp"
+        # Three iterator batches of 2 rows: cumulative 2, 4, 6.
+        # Periodic flush trips once (at 4 >= 3), then a final flush at the end.
+        _wire_iterator(client, [_rows(2), _rows(2, 2), _rows(2, 4)])
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema(
+                source_collection_name=storage.legacy_namespace,
+                target_collection_name=storage.final_namespace,
+            )
+
+        flush_targets = [call.args[0] for call in client.flush.call_args_list]
+        assert flush_targets == [temp, temp]  # one periodic + one final
+        assert storage.final_namespace in collections
+
+    def test_flush_disabled_when_interval_zero(self):
+        storage = _make_model_storage()
+        storage._migration_flush_interval_rows = 0
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_iterator(client, [_rows(2), _rows(2, 2)])
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema(
+                source_collection_name=storage.legacy_namespace,
+                target_collection_name=storage.final_namespace,
+            )
+
+        client.flush.assert_not_called()
+        assert storage.final_namespace in collections
+
+    def test_no_final_flush_for_empty_source(self):
+        storage = _make_model_storage()
+        storage._migration_flush_interval_rows = 50000
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_iterator(client, [])  # empty source
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema(
+                source_collection_name=storage.legacy_namespace,
+                target_collection_name=storage.final_namespace,
+            )
+
+        client.flush.assert_not_called()
+
+    def test_flush_connection_error_is_retryable(self):
+        # A flush failing with a connection error must be retried like any other
+        # connection-class migration failure (rebuild client, restart attempt).
+        storage = _make_model_storage()
+        storage._migration_flush_interval_rows = 1
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_iterator(client, [_rows(2)])
+
+        flush_calls = {"n": 0}
+
+        def flush(collection_name, **kwargs):
+            flush_calls["n"] += 1
+            if flush_calls["n"] == 1:
+                raise MilvusException(
+                    code=2, message="Fail connecting to server on host:19530"
+                )
+
+        client.flush.side_effect = flush
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch.object(storage, "_rebuild_milvus_client") as rebuild:
+                with patch("lightrag.kg.milvus_impl.time.sleep"):
+                    storage._migrate_collection_schema(
+                        source_collection_name=storage.legacy_namespace,
+                        target_collection_name=storage.final_namespace,
+                    )
+
+        rebuild.assert_called_once()
+        assert storage.final_namespace in collections
+
+
+@pytest.mark.offline
+class TestMigrationBatchThrottle:
+    def test_batch_sleep_throttles_between_batches(self):
+        storage = _make_model_storage()
+        storage._migration_flush_interval_rows = 0
+        storage._migration_batch_sleep = 0.05
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_iterator(client, [_rows(2), _rows(2, 2)])
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch("lightrag.kg.milvus_impl.time.sleep") as sleep:
+                storage._migrate_collection_schema(
+                    source_collection_name=storage.legacy_namespace,
+                    target_collection_name=storage.final_namespace,
+                )
+
+        # One sleep per non-empty iterator batch, none for the terminating [].
+        assert [call.args[0] for call in sleep.call_args_list] == [0.05, 0.05]
+
+    def test_no_sleep_when_disabled(self):
+        storage = _make_model_storage()
+        storage._migration_flush_interval_rows = 0
+        storage._migration_batch_sleep = 0.0
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_iterator(client, [_rows(2)])
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch("lightrag.kg.milvus_impl.time.sleep") as sleep:
+                storage._migrate_collection_schema(
+                    source_collection_name=storage.legacy_namespace,
+                    target_collection_name=storage.final_namespace,
+                )
+
+        sleep.assert_not_called()
+
+
+@pytest.mark.offline
+class TestMigrationBackpressureConfig:
+    def test_env_overrides(self, monkeypatch):
+        monkeypatch.setenv("MILVUS_MIGRATION_FLUSH_INTERVAL_ROWS", "1234")
+        monkeypatch.setenv("MILVUS_MIGRATION_BATCH_SLEEP", "0.25")
+        storage = _make_model_storage()
+        assert storage._migration_flush_interval_rows == 1234
+        assert storage._migration_batch_sleep == 0.25
+
+    def test_negative_values_clamped(self, monkeypatch):
+        monkeypatch.setenv("MILVUS_MIGRATION_FLUSH_INTERVAL_ROWS", "-5")
+        monkeypatch.setenv("MILVUS_MIGRATION_BATCH_SLEEP", "-1")
+        storage = _make_model_storage()
+        assert storage._migration_flush_interval_rows == 0
+        assert storage._migration_batch_sleep == 0.0

tests/kg/milvus_impl/test_milvus_migration_retry.py

@@ -0,0 +1,681 @@
+"""
+Tests for Milvus schema-migration resilience.
+
+This suite validates the failure-handling hardening added after a production
+outage where a transient Milvus connection failure mid-migration permanently
+killed worker startup:
+
+1. Connection-class failures retry the whole migration with a rebuilt client.
+2. Non-connection failures keep failing fast (single attempt).
+3. _is_retryable_connection_error classifies errors through cause chains.
+4. The force-create fallback never fires on a connection error.
+5. The temp collection is not loaded during the bulk copy.
+6. Backup collections are released from memory after a successful migration.
+7. A stale _old backup is dropped so the in-place rename can succeed.
+8. An orphaned temp collection (crash between drop-source and rename-temp)
+   is recovered instead of being shadowed by a fresh empty collection.
+"""
+
+import grpc
+import pytest
+from unittest.mock import MagicMock, patch
+from pymilvus import MilvusException
+
+from lightrag.kg.milvus_impl import MilvusVectorDBStorage
+
+
+class _EmbeddingFunc:
+    def __init__(self, dim=128, model_name="text-embedding-3-small"):
+        self.embedding_dim = dim
+        self.model_name = model_name
+
+
+def _make_model_storage(namespace="entities", workspace="test_workspace", dim=128):
+    return MilvusVectorDBStorage(
+        namespace=namespace,
+        workspace=workspace,
+        global_config={
+            "embedding_batch_num": 100,
+            "vector_db_storage_cls_kwargs": {
+                "cosine_better_than_threshold": 0.3,
+            },
+        },
+        embedding_func=_EmbeddingFunc(dim=dim),
+        meta_fields=set(),
+    )
+
+
+def _wire_collection_state(storage, collections):
+    storage._client = MagicMock()
+
+    def has_collection(collection_name):
+        return collection_name in collections
+
+    def create_collection(collection_name, schema):
+        collections.add(collection_name)
+
+    def drop_collection(collection_name):
+        collections.discard(collection_name)
+
+    def rename_collection(source, target):
+        collections.discard(source)
+        collections.add(target)
+
+    storage._client.has_collection.side_effect = has_collection
+    storage._client.create_collection.side_effect = create_collection
+    storage._client.drop_collection.side_effect = drop_collection
+    storage._client.rename_collection.side_effect = rename_collection
+    return storage._client
+
+
+def _rows():
+    return [{"id": "ent-1", "vector": [0.0] * 128, "content": "body"}]
+
+
+def _wire_fresh_iterator_per_attempt(client, rows=None):
+    """query_iterator returns a new exhausted-after-one-batch iterator per call."""
+
+    def make_iterator(**kwargs):
+        iterator = MagicMock()
+        iterator.next.side_effect = [rows if rows is not None else _rows(), []]
+        return iterator
+
+    client.query_iterator.side_effect = make_iterator
+
+
+class _FakeRpcError(grpc.RpcError):
+    def __init__(self, status_code):
+        self._status_code = status_code
+
+    def code(self):
+        return self._status_code
+
+
+_CONNECT_FAILED = MilvusException(
+    code=2, message="Fail connecting to server on host:19530"
+)
+
+
+@pytest.mark.offline
+class TestMigrationRetry:
+    def test_retryable_error_retries_with_rebuilt_client(self):
+        storage = _make_model_storage()
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+
+        insert_calls = {"n": 0}
+
+        def insert(collection_name, data):
+            insert_calls["n"] += 1
+            if insert_calls["n"] == 1:
+                raise MilvusException(
+                    code=2, message="Fail connecting to server on host:19530"
+                )
+
+        client.insert.side_effect = insert
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch.object(storage, "_rebuild_milvus_client") as rebuild:
+                with patch("lightrag.kg.milvus_impl.time.sleep") as sleep:
+                    storage._migrate_collection_schema(
+                        source_collection_name=storage.legacy_namespace,
+                        target_collection_name=storage.final_namespace,
+                    )
+
+        rebuild.assert_called_once()
+        sleep.assert_called_once()
+        assert sleep.call_args.args[0] == pytest.approx(5.0)
+        assert client.query_iterator.call_count == 2
+        assert storage.final_namespace in collections
+        assert f"{storage.final_namespace}_temp" not in collections
+        # The legacy source stays in place as the backup.
+        assert storage.legacy_namespace in collections
+
+    def test_non_retryable_error_fails_fast(self):
+        storage = _make_model_storage()
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+        client.insert.side_effect = RuntimeError("insert failed")
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch.object(storage, "_rebuild_milvus_client") as rebuild:
+                with patch("lightrag.kg.milvus_impl.time.sleep") as sleep:
+                    with pytest.raises(
+                        RuntimeError, match="Iterator-based migration failed"
+                    ):
+                        storage._migrate_collection_schema(
+                            source_collection_name=storage.legacy_namespace,
+                            target_collection_name=storage.final_namespace,
+                        )
+
+        rebuild.assert_not_called()
+        sleep.assert_not_called()
+        assert client.query_iterator.call_count == 1
+        assert storage.legacy_namespace in collections
+        assert storage.final_namespace not in collections
+
+    def test_retries_exhausted_raises_with_backoff_sequence(self):
+        storage = _make_model_storage()
+        storage._migration_max_retries = 2
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+        client.insert.side_effect = ValueError("Cannot invoke RPC on closed channel!")
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch.object(storage, "_rebuild_milvus_client"):
+                with patch("lightrag.kg.milvus_impl.time.sleep") as sleep:
+                    with pytest.raises(
+                        RuntimeError, match="Iterator-based migration failed"
+                    ):
+                        storage._migrate_collection_schema(
+                            source_collection_name=storage.legacy_namespace,
+                            target_collection_name=storage.final_namespace,
+                        )
+
+        assert client.query_iterator.call_count == 3
+        assert [call.args[0] for call in sleep.call_args_list] == [
+            pytest.approx(5.0),
+            pytest.approx(15.0),
+        ]
+        assert storage.legacy_namespace in collections
+
+    def test_max_backoff_is_capped(self):
+        storage = _make_model_storage()
+        storage._migration_max_retries = 4
+        storage._migration_retry_max_backoff = 20.0
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+        client.insert.side_effect = _CONNECT_FAILED
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch.object(storage, "_rebuild_milvus_client"):
+                with patch("lightrag.kg.milvus_impl.time.sleep") as sleep:
+                    with pytest.raises(RuntimeError):
+                        storage._migrate_collection_schema(
+                            source_collection_name=storage.legacy_namespace,
+                            target_collection_name=storage.final_namespace,
+                        )
+
+        assert [call.args[0] for call in sleep.call_args_list] == [
+            pytest.approx(5.0),
+            pytest.approx(15.0),
+            pytest.approx(20.0),
+            pytest.approx(20.0),
+        ]
+
+
+@pytest.mark.offline
+class TestRetryableErrorClassification:
+    @pytest.mark.parametrize(
+        "error",
+        [
+            _FakeRpcError(grpc.StatusCode.UNAVAILABLE),
+            _FakeRpcError(grpc.StatusCode.DEADLINE_EXCEEDED),
+            MilvusException(code=2, message="Fail connecting to server on host:19530"),
+            MilvusException(code=1, message="server unavailable: ping timeout"),
+            ValueError("Cannot invoke RPC on closed channel!"),
+        ],
+        ids=[
+            "grpc-unavailable",
+            "grpc-deadline-exceeded",
+            "milvus-connect-failed",
+            "milvus-connection-message",
+            "closed-channel-valueerror",
+        ],
+    )
+    def test_retryable_errors(self, error):
+        assert MilvusVectorDBStorage._is_retryable_connection_error(error) is True
+
+    @pytest.mark.parametrize(
+        "error",
+        [
+            ValueError("primary keys cannot be truncated"),
+            MilvusException(code=1100, message="schema mismatch"),
+            RuntimeError("Target collection already exists: foo"),
+            KeyError("vector"),
+        ],
+        ids=[
+            "value-error",
+            "milvus-schema-error",
+            "runtime-error",
+            "key-error",
+        ],
+    )
+    def test_non_retryable_errors(self, error):
+        assert MilvusVectorDBStorage._is_retryable_connection_error(error) is False
+
+    def test_cause_chain_is_walked(self):
+        inner = ValueError("Cannot invoke RPC on closed channel!")
+        outer = RuntimeError("Iterator-based migration failed for collection foo")
+        outer.__cause__ = inner
+        assert MilvusVectorDBStorage._is_retryable_connection_error(outer) is True
+
+    def test_context_chain_is_walked(self):
+        inner = MilvusException(code=2, message="Fail connecting to server")
+        outer = RuntimeError("wrapper")
+        outer.__context__ = inner
+        assert MilvusVectorDBStorage._is_retryable_connection_error(outer) is True
+
+    def test_self_referencing_chain_terminates(self):
+        error = RuntimeError("loop")
+        error.__cause__ = error
+        assert MilvusVectorDBStorage._is_retryable_connection_error(error) is False
+
+
+@pytest.mark.offline
+class TestRebuildMilvusClient:
+    def test_rebuild_replaces_client_and_closes_old(self):
+        storage = _make_model_storage()
+        old_client = MagicMock()
+        new_client = MagicMock()
+        storage._client = old_client
+
+        with patch.object(
+            storage, "_create_milvus_client", return_value=new_client
+        ) as create:
+            storage._rebuild_milvus_client()
+
+        old_client.close.assert_called_once()
+        create.assert_called_once()
+        assert storage._client is new_client
+
+    def test_rebuild_tolerates_close_failure_on_dead_client(self):
+        storage = _make_model_storage()
+        old_client = MagicMock()
+        old_client.close.side_effect = ValueError(
+            "Cannot invoke RPC on closed channel!"
+        )
+        new_client = MagicMock()
+        storage._client = old_client
+
+        with patch.object(storage, "_create_milvus_client", return_value=new_client):
+            storage._rebuild_milvus_client()
+
+        assert storage._client is new_client
+
+
+@pytest.mark.offline
+class TestForceCreateGuard:
+    def test_connection_error_propagates_without_force_create(self):
+        storage = _make_model_storage()
+        storage._client = MagicMock()
+        storage._client.has_collection.side_effect = MilvusException(
+            code=2, message="Fail connecting to server on host:19530"
+        )
+
+        with pytest.raises(MilvusException):
+            storage._create_collection_if_not_exist()
+
+        storage._client.create_collection.assert_not_called()
+        storage._client.drop_collection.assert_not_called()
+
+    def test_non_connection_error_still_force_creates(self):
+        storage = _make_model_storage()
+        storage._client = MagicMock()
+        storage._client.has_collection.side_effect = KeyError("boom")
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch.object(storage, "_ensure_collection_loaded"):
+                storage._create_collection_if_not_exist()
+
+        storage._client.create_collection.assert_called_once()
+
+
+@pytest.mark.offline
+class TestMigrationMemoryFootprint:
+    def test_temp_collection_is_not_loaded_during_migration(self):
+        storage = _make_model_storage()
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema(
+                source_collection_name=storage.legacy_namespace,
+                target_collection_name=storage.final_namespace,
+            )
+
+        # The source must be loaded (query_iterator requires it), but the temp
+        # collection must never be loaded during the bulk copy.
+        loaded = [call.args[0] for call in client.load_collection.call_args_list]
+        assert f"{storage.final_namespace}_temp" not in loaded
+        assert storage.legacy_namespace in loaded
+        assert storage.final_namespace in collections
+
+    def test_source_is_loaded_before_query_iterator(self):
+        # query_iterator runs a server-side query that needs the source loaded;
+        # an unloaded legacy/suffix source otherwise fails with code 101
+        # "collection not loaded". The source must be loaded before the
+        # iterator is created.
+        storage = _make_model_storage()
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+
+        order = []
+        client.load_collection.side_effect = lambda name, **kw: order.append(
+            ("load", name)
+        )
+
+        def make_iterator(**kwargs):
+            order.append(("query_iterator", kwargs["collection_name"]))
+            iterator = MagicMock()
+            iterator.next.side_effect = [_rows(), []]
+            return iterator
+
+        client.query_iterator.side_effect = make_iterator
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema(
+                source_collection_name=storage.legacy_namespace,
+                target_collection_name=storage.final_namespace,
+            )
+
+        assert ("load", storage.legacy_namespace) in order
+        assert order.index(("load", storage.legacy_namespace)) < order.index(
+            ("query_iterator", storage.legacy_namespace)
+        )
+
+    def test_failed_suffix_migration_releases_loaded_source(self):
+        # The source we loaded for a suffix migration must be released again on
+        # failure, so a failed startup/retry does not leave a full backup
+        # resident in query-node memory.
+        storage = _make_model_storage()
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+        client.insert.side_effect = RuntimeError("insert failed")  # non-retryable
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with pytest.raises(RuntimeError, match="Iterator-based migration failed"):
+                storage._migrate_collection_schema(
+                    source_collection_name=storage.legacy_namespace,
+                    target_collection_name=storage.final_namespace,
+                )
+
+        client.release_collection.assert_any_call(storage.legacy_namespace)
+
+    def test_failed_inplace_migration_does_not_release_active_source(self):
+        # An in-place source is the live active collection; it must not be
+        # released on a pre-commit failure.
+        storage = _make_model_storage()
+        collections = {storage.final_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+        client.insert.side_effect = RuntimeError("insert failed")  # non-retryable
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with pytest.raises(RuntimeError, match="Iterator-based migration failed"):
+                storage._migrate_collection_schema()  # in-place
+
+        client.release_collection.assert_not_called()
+
+    def test_suffix_migration_releases_legacy_backup(self):
+        storage = _make_model_storage()
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema(
+                source_collection_name=storage.legacy_namespace,
+                target_collection_name=storage.final_namespace,
+            )
+
+        client.release_collection.assert_called_once_with(storage.legacy_namespace)
+        assert storage.legacy_namespace in collections
+
+    def test_in_place_migration_releases_old_backup(self):
+        storage = _make_model_storage()
+        collections = {storage.final_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema()
+
+        backup_name = f"{storage.final_namespace}_old"
+        client.release_collection.assert_called_once_with(backup_name)
+        assert backup_name in collections
+        assert storage.final_namespace in collections
+
+    def test_release_failure_does_not_fail_migration(self):
+        storage = _make_model_storage()
+        collections = {storage.legacy_namespace}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+        client.release_collection.side_effect = MilvusException(
+            code=1, message="release failed"
+        )
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema(
+                source_collection_name=storage.legacy_namespace,
+                target_collection_name=storage.final_namespace,
+            )
+
+        assert storage.final_namespace in collections
+
+
+@pytest.mark.offline
+class TestInPlaceMigrationSafety:
+    def test_stale_old_backup_is_dropped_before_rename(self):
+        storage = _make_model_storage()
+        backup_name = f"{storage.final_namespace}_old"
+        collections = {storage.final_namespace, backup_name}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            storage._migrate_collection_schema()
+
+        # The stale backup was dropped so the rename could succeed; the source
+        # now lives on as the fresh backup instead of being dropped outright.
+        assert backup_name in collections
+        assert storage.final_namespace in collections
+        drop_calls = [call.args[0] for call in client.drop_collection.call_args_list]
+        assert backup_name in drop_calls
+
+    def test_orphaned_temp_collection_is_recovered_on_startup(self):
+        storage = _make_model_storage()
+        temp_name = f"{storage.final_namespace}_temp"
+        collections = {temp_name}
+        client = _wire_collection_state(storage, collections)
+
+        with patch.object(storage, "_ensure_collection_loaded"):
+            storage._create_collection_if_not_exist()
+
+        assert storage.final_namespace in collections
+        assert temp_name not in collections
+        client.create_collection.assert_not_called()
+
+    def test_startup_restores_old_backup_when_no_temp_survives(self):
+        # Target gone, no temp, only the _old backup left: the source was
+        # vacated but the migrated copy did not survive. Restore _old rather
+        # than creating an empty collection over the last copy.
+        storage = _make_model_storage()
+        old_name = f"{storage.final_namespace}_old"
+        collections = {old_name}
+        client = _wire_collection_state(storage, collections)
+
+        with patch.object(storage, "_ensure_collection_loaded"):
+            storage._create_collection_if_not_exist()
+
+        assert storage.final_namespace in collections
+        assert old_name not in collections
+        client.create_collection.assert_not_called()
+
+    def test_inplace_recovery_precedes_legacy_migration(self):
+        # A suffixed target was migrated in-place and interrupted after Step 3:
+        # final renamed to _old, the completed copy sits in _temp, and the old
+        # unsuffixed legacy backup still exists. Recovery MUST win over the
+        # legacy migration, which would otherwise drop _temp and overwrite the
+        # target with stale legacy data (losing every write since the split).
+        storage = _make_model_storage()
+        final = storage.final_namespace
+        legacy = storage.legacy_namespace
+        temp = f"{final}_temp"
+        old = f"{final}_old"
+        collections = {legacy, old, temp}
+        client = _wire_collection_state(storage, collections)
+
+        with patch.object(storage, "_ensure_collection_loaded"):
+            with patch.object(storage, "_migrate_collection_schema") as migrate:
+                storage._create_collection_if_not_exist()
+
+        migrate.assert_not_called()
+        assert final in collections  # _temp promoted to the target
+        assert temp not in collections
+        assert legacy in collections  # untouched
+        client.create_collection.assert_not_called()
+        drops = [call.args[0] for call in client.drop_collection.call_args_list]
+        assert temp not in drops
+
+    def test_partial_suffix_temp_with_legacy_is_remigrated_not_promoted(self):
+        # An aborted suffix copy (legacy -> final) left a PARTIAL _temp while
+        # legacy is intact and there is no _old. The partial temp must be
+        # treated as scratch and re-migrated from legacy, never promoted.
+        storage = _make_model_storage()
+        final = storage.final_namespace
+        legacy = storage.legacy_namespace
+        temp = f"{final}_temp"
+        collections = {legacy, temp}
+        client = _wire_collection_state(storage, collections)
+
+        with patch.object(storage, "_has_vector_field", return_value=True):
+            with patch.object(storage, "_check_vector_dimension"):
+                with patch.object(storage, "_migrate_collection_schema") as migrate:
+                    with patch.object(storage, "_ensure_collection_loaded"):
+                        storage._create_collection_if_not_exist()
+
+        migrate.assert_called_once_with(
+            source_collection_name=legacy,
+            target_collection_name=final,
+        )
+        # Recovery (which promotes/restores via rename) must not have run.
+        client.rename_collection.assert_not_called()
+
+
+def _fail_specific_rename_once(client, collections, fail_source, fail_target, error):
+    """Wrap the wired rename so one specific rename raises `error` on first call."""
+    state = {"n": 0}
+
+    def rename(source, target):
+        if source == fail_source and target == fail_target:
+            state["n"] += 1
+            if state["n"] == 1:
+                raise error
+        collections.discard(source)
+        collections.add(target)
+
+    client.rename_collection.side_effect = rename
+
+
+@pytest.mark.offline
+class TestInPlaceCommitWindowRecovery:
+    """The in-place commit window (source vacated, temp not yet promoted) must
+    treat temp/_old as recoverable state, never as scratch to be dropped."""
+
+    def test_step4_connection_failure_retry_promotes_temp(self):
+        storage = _make_model_storage()
+        final = storage.final_namespace
+        temp = f"{final}_temp"
+        old = f"{final}_old"
+        collections = {final}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+        # Step 4 (temp -> final) fails once with a connection error; the retry
+        # must promote the surviving temp copy, not drop and re-copy.
+        _fail_specific_rename_once(
+            client,
+            collections,
+            temp,
+            final,
+            MilvusException(code=2, message="Fail connecting to server on host:19530"),
+        )
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch.object(storage, "_rebuild_milvus_client") as rebuild:
+                with patch("lightrag.kg.milvus_impl.time.sleep") as sleep:
+                    storage._migrate_collection_schema()
+
+        rebuild.assert_called_once()
+        sleep.assert_called_once()
+        assert final in collections
+        assert temp not in collections
+        assert old in collections  # backup preserved
+        assert storage.final_namespace == final
+        drops = [call.args[0] for call in client.drop_collection.call_args_list]
+        assert temp not in drops  # never treated as scratch in the commit window
+
+    def test_drop_source_fallback_then_step4_failure_recovers_temp(self):
+        # rename(source -> _old) fails for a non-connection reason, so the
+        # drop-source fallback runs (NO _old backup); Step 4 then fails with a
+        # connection error. Without recovery this is total data loss; the retry
+        # must promote the only surviving copy (temp).
+        storage = _make_model_storage()
+        final = storage.final_namespace
+        temp = f"{final}_temp"
+        old = f"{final}_old"
+        collections = {final}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+        state = {"temp_to_final": 0}
+
+        def rename(source, target):
+            if source == final and target == old:
+                raise RuntimeError("rename to _old unsupported")  # forces drop-source
+            if source == temp and target == final:
+                state["temp_to_final"] += 1
+                if state["temp_to_final"] == 1:
+                    raise MilvusException(
+                        code=2, message="Fail connecting to server on host:19530"
+                    )
+            collections.discard(source)
+            collections.add(target)
+
+        client.rename_collection.side_effect = rename
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with patch.object(storage, "_rebuild_milvus_client") as rebuild:
+                with patch("lightrag.kg.milvus_impl.time.sleep"):
+                    storage._migrate_collection_schema()
+
+        rebuild.assert_called_once()
+        assert final in collections
+        assert temp not in collections
+        assert old not in collections
+        assert storage.final_namespace == final
+        drops = [call.args[0] for call in client.drop_collection.call_args_list]
+        assert temp not in drops
+
+    def test_non_retryable_commit_failure_keeps_temp_for_startup_recovery(self):
+        # A non-retryable Step 4 failure must still preserve temp (the source is
+        # already vacated) so startup recovery can finish the commit.
+        storage = _make_model_storage()
+        final = storage.final_namespace
+        temp = f"{final}_temp"
+        old = f"{final}_old"
+        collections = {final}
+        client = _wire_collection_state(storage, collections)
+        _wire_fresh_iterator_per_attempt(client)
+
+        def rename(source, target):
+            if source == temp and target == final:
+                raise RuntimeError("non-retryable rename failure")
+            collections.discard(source)
+            collections.add(target)
+
+        client.rename_collection.side_effect = rename
+
+        with patch.object(storage, "_create_indexes_after_collection"):
+            with pytest.raises(RuntimeError, match="Iterator-based migration failed"):
+                storage._migrate_collection_schema()
+
+        assert temp in collections  # preserved as recovery state
+        assert old in collections
+        drops = [call.args[0] for call in client.drop_collection.call_args_list]
+        assert temp not in drops

tests/kg/mongo_impl/test_client_manager_lifecycle.py

@@ -0,0 +1,100 @@
+"""Tests for ClientManager connection lifecycle (no MongoDB instance required).
+
+These tests verify that ClientManager properly closes the underlying
+AsyncMongoClient when all references are released, and keeps it alive
+while references remain.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+
+pytest.importorskip(
+    "pymongo",
+    reason="pymongo is required for Mongo ClientManager tests",
+)
+
+from lightrag.kg.mongo_impl import ClientManager
+
+
+class TestClientManagerLifecycle:
+    """Verify ClientManager connection open/close behavior."""
+
+    def _reset_manager(self):
+        """Reset ClientManager class state between tests."""
+        ClientManager._instances = {"client": None, "db": None, "ref_count": 0}
+
+    def teardown_method(self):
+        self._reset_manager()
+
+    @pytest.mark.asyncio
+    async def test_release_client_closes_connection_when_ref_count_zero(self):
+        """When ref_count drops to 0, the MongoClient should be closed and cleared."""
+        mock_client = AsyncMock()
+        mock_db = MagicMock()
+
+        ClientManager._instances = {
+            "client": mock_client,
+            "db": mock_db,
+            "ref_count": 1,
+        }
+
+        await ClientManager.release_client(mock_db)
+
+        mock_client.close.assert_awaited_once()
+        assert ClientManager._instances["client"] is None
+        assert ClientManager._instances["db"] is None
+        assert ClientManager._instances["ref_count"] == 0
+
+    @pytest.mark.asyncio
+    async def test_release_client_keeps_connection_with_multiple_refs(self):
+        """When other references exist, the MongoClient must NOT be closed."""
+        mock_client = AsyncMock()
+        mock_db = MagicMock()
+
+        ClientManager._instances = {
+            "client": mock_client,
+            "db": mock_db,
+            "ref_count": 3,
+        }
+
+        await ClientManager.release_client(mock_db)
+
+        mock_client.close.assert_not_awaited()
+        assert ClientManager._instances["ref_count"] == 2
+        assert ClientManager._instances["client"] is mock_client
+        assert ClientManager._instances["db"] is mock_db
+
+    @pytest.mark.asyncio
+    async def test_release_client_noop_for_wrong_db(self):
+        """Releasing a db that is not the tracked instance should do nothing."""
+        mock_client = AsyncMock()
+        mock_db = MagicMock()
+        other_db = MagicMock()
+
+        ClientManager._instances = {
+            "client": mock_client,
+            "db": mock_db,
+            "ref_count": 1,
+        }
+
+        await ClientManager.release_client(other_db)
+
+        mock_client.close.assert_not_awaited()
+        assert ClientManager._instances["ref_count"] == 1
+
+    @pytest.mark.asyncio
+    async def test_release_client_noop_for_none(self):
+        """Releasing None should be a safe no-op."""
+        mock_client = AsyncMock()
+        mock_db = MagicMock()
+
+        ClientManager._instances = {
+            "client": mock_client,
+            "db": mock_db,
+            "ref_count": 1,
+        }
+
+        await ClientManager.release_client(None)
+
+        mock_client.close.assert_not_awaited()
+        assert ClientManager._instances["ref_count"] == 1

tests/kg/mongo_impl/test_mongo_deferred_embedding.py

@@ -0,0 +1,567 @@
+"""Unit tests for MongoVectorDBStorage's deferred-embedding flush pipeline.
+
+All tests use mocks — no running MongoDB instance required.
+Mirrors the structure of tests/kg/opensearch_impl/test_opensearch_storage.py's
+TestVectorStorageBatching to keep behaviour aligned across backends.
+"""
+
+import asyncio
+import os
+
+import numpy as np
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+pytest.importorskip(
+    "pymongo",
+    reason="pymongo is required for Mongo storage tests",
+)
+
+from pymongo import UpdateOne  # type: ignore
+
+from lightrag.kg.mongo_impl import MongoVectorDBStorage
+
+pytestmark = pytest.mark.offline
+
+
+# ---------------------------------------------------------------------------
+# Fixtures and helpers
+# ---------------------------------------------------------------------------
+
+
+class MockEmbeddingFunc:
+    def __init__(self, dim=8):
+        self.embedding_dim = dim
+        self.max_token_size = 512
+        self.model_name = "mock-embed"
+
+    async def __call__(self, texts, **kwargs):
+        return np.random.rand(len(texts), self.embedding_dim).astype(np.float32)
+
+
+class CountingEmbeddingFunc(MockEmbeddingFunc):
+    def __init__(self, dim=8, fail_times=0):
+        super().__init__(dim=dim)
+        self.fail_times = fail_times
+        self.call_count = 0
+        self.batches: list[list[str]] = []
+        self.texts: list[str] = []
+
+    async def __call__(self, texts, **kwargs):
+        self.call_count += 1
+        batch = list(texts)
+        self.batches.append(batch)
+        self.texts.extend(batch)
+        if self.fail_times > 0:
+            self.fail_times -= 1
+            raise RuntimeError("embedding failed")
+        return await super().__call__(texts, **kwargs)
+
+
+class _AsyncCursor:
+    def __init__(self, docs):
+        self._docs = list(docs)
+
+    async def to_list(self, length=None):
+        return list(self._docs)
+
+
+@pytest.fixture(autouse=True)
+def patch_namespace_lock(monkeypatch):
+    """Cache real asyncio.Locks per (namespace, workspace) for shared semantics.
+
+    Also unconditionally clears ``MONGODB_WORKSPACE`` so tests are insulated
+    from shell-level env leakage: ``final_namespace`` depends on this var,
+    and a leaked value (e.g. ``space2``) silently collapses distinct
+    workspaces into one namespace. Tests that need an override set it
+    explicitly via ``patch.dict(os.environ, ...)``.
+    """
+    monkeypatch.delenv("MONGODB_WORKSPACE", raising=False)
+    cache: dict[tuple[str, str | None], asyncio.Lock] = {}
+
+    def factory(namespace, workspace=None, enable_logging=False):
+        key = (namespace, workspace or "")
+        lock = cache.get(key)
+        if lock is None:
+            lock = asyncio.Lock()
+            cache[key] = lock
+        return lock
+
+    with patch("lightrag.kg.mongo_impl.get_namespace_lock", side_effect=factory):
+        yield cache
+
+
+def _make_storage(
+    embed_func,
+    *,
+    namespace="entities",
+    workspace="test",
+    meta_fields=None,
+):
+    if meta_fields is None:
+        meta_fields = {"content", "entity_name", "src_id", "tgt_id"}
+    storage = MongoVectorDBStorage(
+        namespace=namespace,
+        workspace=workspace,
+        global_config={
+            "embedding_batch_num": 10,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=embed_func,
+        meta_fields=meta_fields,
+    )
+    # Wire a fake AsyncCollection (the only Mongo surface our code touches).
+    storage._data = MagicMock()
+    storage._data.bulk_write = AsyncMock()
+    storage._data.delete_many = AsyncMock(return_value=MagicMock(deleted_count=0))
+    storage._data.find_one = AsyncMock(return_value=None)
+    storage._data.find = MagicMock(return_value=_AsyncCursor([]))
+    storage.db = MagicMock()  # non-None so finalize releases it
+
+    from lightrag.kg.mongo_impl import get_namespace_lock
+
+    storage._flush_lock = get_namespace_lock(
+        namespace=storage.final_namespace, workspace=""
+    )
+    return storage
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_upsert_buffers_without_embedding():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}, "v2": {"content": "world"}})
+
+    assert embed.call_count == 0
+    assert set(s._pending_vector_docs.keys()) == {"v1", "v2"}
+    assert s._pending_vector_docs["v1"].vector is None
+    s._data.bulk_write.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_index_done_callback_triggers_flush():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}, "v2": {"content": "world"}})
+    await s.index_done_callback()
+
+    assert embed.call_count == 1
+    s._data.bulk_write.assert_called_once()
+    ops, kwargs = (
+        s._data.bulk_write.call_args.args[0],
+        s._data.bulk_write.call_args.kwargs,
+    )
+    assert kwargs.get("ordered") is False
+    assert all(isinstance(op, UpdateOne) for op in ops)
+    assert len(ops) == 2
+    assert s._pending_vector_docs == {}
+
+
+@pytest.mark.asyncio
+async def test_repeated_upsert_same_id_embeds_once():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "first"}})
+    await s.upsert({"v1": {"content": "second"}})
+    await s.upsert({"v1": {"content": "third"}})
+    await s.index_done_callback()
+
+    assert embed.call_count == 1
+    assert embed.texts == ["third"]
+    s._data.bulk_write.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_deferred_embeddings_respect_batch_size():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._max_batch_size = 2
+
+    await s.upsert({f"v{i}": {"content": f"doc {i}"} for i in range(5)})
+    await s.index_done_callback()
+
+    assert embed.call_count == 3
+    assert [len(b) for b in embed.batches] == [2, 2, 1]
+
+
+@pytest.mark.asyncio
+async def test_get_vectors_by_ids_lazy_embed_then_reuse_in_flush():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}})
+
+    vectors = await s.get_vectors_by_ids(["v1"])
+    assert "v1" in vectors
+    assert embed.call_count == 1
+    assert s._pending_vector_docs["v1"].vector is not None
+
+    await s.index_done_callback()
+    assert embed.call_count == 1
+    s._data.bulk_write.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_flush_failure_keeps_buffer_no_double_embed_on_retry():
+    embed = CountingEmbeddingFunc(fail_times=1)
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}})
+
+    with pytest.raises(RuntimeError, match="embedding failed"):
+        await s.index_done_callback()
+
+    assert "v1" in s._pending_vector_docs
+    assert s._pending_vector_docs["v1"].vector is None
+    s._data.bulk_write.assert_not_called()
+
+    await s.index_done_callback()
+    assert embed.call_count == 2
+    s._data.bulk_write.assert_called_once()
+    assert s._pending_vector_docs == {}
+
+
+@pytest.mark.asyncio
+async def test_server_write_failure_keeps_buffer():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._data.bulk_write.side_effect = RuntimeError("mongo down")
+
+    await s.upsert({"v1": {"content": "hello"}})
+
+    with pytest.raises(RuntimeError, match="mongo down"):
+        await s.index_done_callback()
+
+    assert "v1" in s._pending_vector_docs
+    assert s._pending_vector_docs["v1"].vector is not None
+
+    s._data.bulk_write.side_effect = None
+    await s.index_done_callback()
+    assert embed.call_count == 1
+    assert s._pending_vector_docs == {}
+
+
+@pytest.mark.asyncio
+async def test_finalize_raises_when_buffer_unflushed_and_still_releases_client():
+    """finalize() must release the Mongo client even when the flush fails."""
+    from lightrag.kg.mongo_impl import ClientManager
+
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._data.bulk_write.side_effect = RuntimeError("mongo down")
+
+    await s.upsert({"v1": {"content": "hello"}})
+
+    with patch.object(ClientManager, "release_client", new=AsyncMock()) as rel:
+        with pytest.raises(RuntimeError, match="finalize.*flush raised") as exc_info:
+            await s.finalize()
+        rel.assert_awaited_once()
+
+    # Operator-diagnostic counts must appear in the message so the buffered
+    # data loss is auditable from the log alone (1 upsert pre-loaded, 0 deletes).
+    msg = str(exc_info.value)
+    assert "1 pending upserts" in msg
+    assert "0 pending deletes" in msg
+
+    # Client references cleared so a second finalize doesn't release twice.
+    assert s.db is None
+
+
+@pytest.mark.asyncio
+async def test_delete_then_upsert_same_id_keeps_upsert():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.delete(["v1"])
+    assert "v1" in s._pending_vector_deletes
+
+    await s.upsert({"v1": {"content": "hello"}})
+    assert "v1" in s._pending_vector_docs
+    assert "v1" not in s._pending_vector_deletes
+
+    await s.index_done_callback()
+    ops = s._data.bulk_write.call_args.args[0]
+    assert all(isinstance(op, UpdateOne) for op in ops)
+
+
+@pytest.mark.asyncio
+async def test_upsert_then_delete_same_id_keeps_delete():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello"}})
+    await s.delete(["v1"])
+
+    assert "v1" not in s._pending_vector_docs
+    assert "v1" in s._pending_vector_deletes
+
+    await s.index_done_callback()
+    s._data.bulk_write.assert_not_called()
+    s._data.delete_many.assert_awaited_once_with({"_id": {"$in": ["v1"]}})
+    assert s._pending_vector_deletes == set()
+
+
+@pytest.mark.asyncio
+async def test_flush_uses_update_one_and_delete_many_mix():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"u1": {"content": "u1"}, "u2": {"content": "u2"}})
+    await s.delete(["d1", "d2"])
+
+    await s.index_done_callback()
+    ops = s._data.bulk_write.call_args.args[0]
+    assert sum(isinstance(op, UpdateOne) for op in ops) == 2
+    s._data.delete_many.assert_awaited_once()
+    delete_filter = s._data.delete_many.await_args.args[0]
+    assert set(delete_filter["_id"]["$in"]) == {"d1", "d2"}
+
+
+@pytest.mark.asyncio
+async def test_flush_chunks_deletes_by_record_count():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._max_delete_records_per_batch = 2
+
+    await s.delete([f"d{i}" for i in range(5)])
+
+    await s.index_done_callback()
+
+    assert s._data.delete_many.await_count == 3
+    chunk_sizes = sorted(
+        len(call.args[0]["_id"]["$in"]) for call in s._data.delete_many.await_args_list
+    )
+    assert chunk_sizes == [1, 2, 2]
+    assert s._pending_vector_deletes == set()
+
+
+@pytest.mark.asyncio
+async def test_delete_entity_relation_raises_on_server_failure():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._data.find = MagicMock(
+        return_value=_AsyncCursor([{"_id": "rel1"}, {"_id": "rel2"}])
+    )
+    s._data.delete_many = AsyncMock(side_effect=RuntimeError("mongo delete failed"))
+
+    with pytest.raises(RuntimeError, match="mongo delete failed"):
+        await s.delete_entity_relation("X")
+
+
+@pytest.mark.asyncio
+async def test_delete_entity_relation_prunes_pending_buffer():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert(
+        {
+            "rel-A-B": {"content": "A → B", "src_id": "A", "tgt_id": "B"},
+            "rel-C-D": {"content": "C → D", "src_id": "C", "tgt_id": "D"},
+        }
+    )
+    s._data.find = MagicMock(return_value=_AsyncCursor([]))
+
+    await s.delete_entity_relation("A")
+    assert "rel-A-B" not in s._pending_vector_docs
+    assert "rel-C-D" in s._pending_vector_docs
+
+
+@pytest.mark.asyncio
+async def test_get_by_id_buffer_excludes_vector():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.upsert({"v1": {"content": "hello", "entity_name": "E1"}})
+    doc = await s.get_by_id("v1")
+    assert doc is not None
+    assert doc["id"] == "v1"
+    assert doc.get("entity_name") == "E1"
+    assert "vector" not in doc
+    s._data.find_one.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_get_by_id_fallback_projects_out_vector():
+    """Server-side find_one must request projection={'vector': 0}."""
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    s._data.find_one = AsyncMock(
+        return_value={"_id": "v9", "entity_name": "X", "created_at": 0}
+    )
+
+    doc = await s.get_by_id("v9")
+    assert doc is not None
+    assert "vector" not in doc
+    args, kwargs = s._data.find_one.call_args.args, s._data.find_one.call_args.kwargs
+    # projection is positional arg #2 in Mongo's API.
+    projection = args[1] if len(args) > 1 else kwargs.get("projection")
+    assert projection == {"vector": 0}
+
+
+@pytest.mark.asyncio
+async def test_get_by_ids_fallback_projects_out_vector():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    s._data.find = MagicMock(
+        return_value=_AsyncCursor(
+            [{"_id": "a", "entity_name": "A"}, {"_id": "b", "entity_name": "B"}]
+        )
+    )
+    docs = await s.get_by_ids(["a", "b"])
+    assert len(docs) == 2
+    assert all("vector" not in d for d in docs if d)
+    args, kwargs = s._data.find.call_args.args, s._data.find.call_args.kwargs
+    projection = args[1] if len(args) > 1 else kwargs.get("projection")
+    assert projection == {"vector": 0}
+
+
+@pytest.mark.asyncio
+async def test_get_by_id_returns_none_for_pending_delete():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    await s.delete(["v1"])
+    assert await s.get_by_id("v1") is None
+    s._data.find_one.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_env_workspace_override_shares_flush_lock(patch_namespace_lock):
+    cache = patch_namespace_lock
+    embed = CountingEmbeddingFunc()
+
+    with patch.dict(os.environ, {"MONGODB_WORKSPACE": "shared_ws"}, clear=False):
+        a = _make_storage(embed, workspace="caller_a")
+        b = _make_storage(embed, workspace="caller_b")
+        assert a.final_namespace == b.final_namespace == "shared_ws_entities"
+        assert a._flush_lock is b._flush_lock
+        assert len([k for k in cache if k[0] == "shared_ws_entities"]) == 1
+
+
+@pytest.mark.asyncio
+async def test_same_workspace_param_shares_flush_lock(patch_namespace_lock):
+    """Plain ctor path (no MONGODB_WORKSPACE env): same workspace → shared lock.
+
+    Companion to ``test_env_workspace_override_shares_flush_lock``; together
+    they cover both ways two instances can land on the same final_namespace.
+    The autouse fixture clears MONGODB_WORKSPACE so this exercises the plain
+    constructor path, not the env-override path.
+    """
+    cache = patch_namespace_lock
+    embed = CountingEmbeddingFunc()
+    a = _make_storage(embed, workspace="caller")
+    b = _make_storage(embed, workspace="caller")
+    assert a.final_namespace == b.final_namespace == "caller_entities"
+    assert a._flush_lock is b._flush_lock
+    assert len([k for k in cache if k[0] == "caller_entities"]) == 1
+
+
+@pytest.mark.asyncio
+async def test_distinct_namespaces_get_independent_locks():
+    embed = CountingEmbeddingFunc()
+    a = _make_storage(embed, workspace="a")
+    b = _make_storage(embed, workspace="b")
+    assert a.final_namespace != b.final_namespace
+    assert a._flush_lock is not b._flush_lock
+
+
+@pytest.mark.asyncio
+async def test_concurrent_upsert_and_flush_serialize_on_lock():
+    """upsert() and index_done_callback() racing on the same namespace must
+    not corrupt the buffer or split a single doc across two embed calls.
+
+    Drives a slow embed (asyncio.Event-gated) so the flush genuinely holds
+    the lock while a second coroutine attempts upsert mid-flight. Asserts:
+      - the late upsert lands in the buffer (not silently dropped)
+      - it is *not* embedded by the in-flight flush (still pending after)
+      - a follow-up flush picks it up cleanly with exactly one extra embed
+
+    Note: we replace ``s.embedding_func`` directly (not ``patch.object`` on
+    ``embed.__call__``) because Python dispatches ``embed(...)`` through
+    ``type(embed).__call__``, bypassing any instance-level patch.
+    """
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+
+    embed_gate = asyncio.Event()
+    flush_entered = asyncio.Event()
+    original_embed = s.embedding_func
+
+    async def gated_embed(texts, **kwargs):
+        flush_entered.set()
+        await embed_gate.wait()
+        return await original_embed(texts, **kwargs)
+
+    await s.upsert({"v1": {"content": "first"}})
+    s.embedding_func = gated_embed
+    try:
+        flush_task = asyncio.create_task(s.index_done_callback())
+        await flush_entered.wait()  # flush is now holding _flush_lock
+
+        # This upsert must wait on _flush_lock; schedule it concurrently.
+        late_upsert = asyncio.create_task(s.upsert({"v2": {"content": "late"}}))
+        # Give the event loop a chance to actually start late_upsert and
+        # confirm it is blocked on the lock (still no v2 in buffer).
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert "v2" not in s._pending_vector_docs
+        assert not late_upsert.done()
+
+        embed_gate.set()
+        await flush_task
+        await late_upsert
+    finally:
+        s.embedding_func = original_embed
+
+    # Flush embedded v1 only; v2 arrived after the docs_to_embed snapshot.
+    assert embed.call_count == 1
+    assert embed.batches == [["first"]]
+    assert "v1" not in s._pending_vector_docs  # flushed
+    assert "v2" in s._pending_vector_docs  # still buffered
+    s._data.bulk_write.assert_called_once()
+
+    # Next flush picks up the late upsert without re-embedding v1.
+    await s.index_done_callback()
+    assert embed.call_count == 2
+    assert embed.batches[-1] == ["late"]
+    assert s._pending_vector_docs == {}
+    assert s._data.bulk_write.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_drop_clears_pending_buffers():
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    with patch.object(s, "create_vector_index_if_not_exists", new=AsyncMock()):
+        await s.upsert({"v1": {"content": "hello"}})
+        await s.delete(["v2"])
+        assert s._pending_vector_docs and s._pending_vector_deletes
+
+        result = await s.drop()
+        assert result["status"] == "success"
+        assert s._pending_vector_docs == {}
+        assert s._pending_vector_deletes == set()
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_drop_pending_index_ops_clears_buffers():
+    """On an internal-error abort the pipeline calls drop_pending_index_ops to
+    discard buffered upserts/deletes without flushing them (PR #3187)."""
+    embed = CountingEmbeddingFunc()
+    s = _make_storage(embed)
+    await s.upsert({"v1": {"content": "x"}, "v2": {"content": "y"}})
+    s._pending_vector_deletes.add("old-id")
+    assert s._pending_vector_docs
+    await s.drop_pending_index_ops()
+    assert not s._pending_vector_docs
+    assert not s._pending_vector_deletes

tests/kg/mongo_impl/test_mongo_storage.py

@@ -0,0 +1,656 @@
+import pytest
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock, patch
+
+pytest.importorskip(
+    "pymongo",
+    reason="pymongo is required for Mongo storage tests",
+)
+
+from pymongo.errors import PyMongoError, BulkWriteError, DuplicateKeyError
+
+from lightrag.kg.mongo_impl import (
+    MongoDocStatusStorage,
+    MongoGraphStorage,
+    _canonical_edge_endpoints,
+)
+
+pytestmark = pytest.mark.offline
+
+
+class _AsyncCursor:
+    def __init__(self, docs):
+        self._docs = list(docs)
+
+    def limit(self, n: int):
+        self._docs = self._docs[:n]
+        return self
+
+    def __aiter__(self):
+        self._iter = iter(self._docs)
+        return self
+
+    async def __anext__(self):
+        try:
+            return next(self._iter)
+        except StopIteration:
+            raise StopAsyncIteration
+
+
+class TestMongoGraphStorage:
+    def _make_storage(self):
+        storage = MongoGraphStorage.__new__(MongoGraphStorage)
+        storage.workspace = "test"
+        storage.global_config = {"max_graph_nodes": 1000}
+        storage._edge_collection_name = "test_edges"
+        storage.collection = SimpleNamespace()
+        storage.edge_collection = SimpleNamespace()
+        return storage
+
+    @pytest.mark.asyncio
+    async def test_get_knowledge_graph_all_backfills_isolated_nodes_when_truncated(
+        self,
+    ):
+        storage = self._make_storage()
+        storage.collection.count_documents = AsyncMock(return_value=5)
+        storage.edge_collection.aggregate = AsyncMock(
+            return_value=_AsyncCursor(
+                [{"_id": "A", "degree": 1}, {"_id": "B", "degree": 1}]
+            )
+        )
+
+        def collection_find_side_effect(query, projection=None):
+            if query == {"_id": {"$nin": ["A", "B"]}}:
+                return _AsyncCursor(
+                    [
+                        {"_id": "C", "entity_type": "person"},
+                        {"_id": "D", "entity_type": "person"},
+                        {"_id": "E", "entity_type": "person"},
+                    ]
+                )
+            if query == {"_id": {"$in": ["A", "B", "C", "D"]}}:
+                return _AsyncCursor(
+                    [
+                        {"_id": "B", "entity_type": "person"},
+                        {"_id": "D", "entity_type": "person"},
+                        {"_id": "A", "entity_type": "person"},
+                        {"_id": "C", "entity_type": "person"},
+                    ]
+                )
+            raise AssertionError(f"Unexpected node query: {query}")
+
+        storage.collection.find = Mock(side_effect=collection_find_side_effect)
+        storage.edge_collection.find = Mock(
+            return_value=_AsyncCursor(
+                [
+                    {
+                        "source_node_id": "A",
+                        "target_node_id": "B",
+                        "relationship": "knows",
+                    }
+                ]
+            )
+        )
+
+        result = await storage.get_knowledge_graph_all_by_degree(
+            max_depth=2, max_nodes=4
+        )
+
+        assert result.is_truncated is True
+        assert [node.id for node in result.nodes] == ["A", "B", "C", "D"]
+        assert len(result.edges) == 1
+        assert result.edges[0].source == "A"
+        assert result.edges[0].target == "B"
+
+
+class TestMongoEdgeKey:
+    """Canonical edge-endpoint writes, duplicate-key retries, and the migration."""
+
+    def _make_storage(self):
+        s = MongoGraphStorage.__new__(MongoGraphStorage)
+        s.workspace = "test"
+        s.namespace = "chunk_entity_relation"
+        s.global_config = {}
+        s._edge_collection_name = "test_edges"
+        s._max_upsert_payload_bytes = 16 * 1024 * 1024
+        s._max_upsert_records_per_batch = 128
+        s.collection = SimpleNamespace(update_one=AsyncMock())
+        s.edge_collection = SimpleNamespace()
+        return s
+
+    def test_canonical_edge_endpoints_are_direction_independent(self):
+        assert _canonical_edge_endpoints("B", "A") == _canonical_edge_endpoints(
+            "A", "B"
+        )
+        assert _canonical_edge_endpoints("B", "A") == ("A", "B")
+
+    def test_canonical_endpoints_never_collide_across_delimiter(self):
+        # Distinct pairs that a delimiter-joined key could conflate must stay
+        # distinct as separate (lo, hi) fields.
+        assert _canonical_edge_endpoints("A\x1fB", "C") != _canonical_edge_endpoints(
+            "A", "B\x1fC"
+        )
+
+    @pytest.mark.asyncio
+    async def test_upsert_edge_filters_and_sets_canonical_endpoints(self):
+        s = self._make_storage()
+        s.edge_collection.update_one = AsyncMock()
+        await s.upsert_edge("B", "A", {"weight": 1.0, "source_id": "c1<SEP>c2"})
+
+        args, kwargs = s.edge_collection.update_one.call_args
+        filt, update = args[0], args[1]
+        lo, hi = _canonical_edge_endpoints("B", "A")
+        assert filt == {"edge_lo": lo, "edge_hi": hi}
+        assert update["$set"]["edge_lo"] == lo
+        assert update["$set"]["edge_hi"] == hi
+        assert update["$set"]["source_node_id"] == "B"
+        assert update["$set"]["target_node_id"] == "A"
+        assert update["$set"]["source_ids"] == ["c1", "c2"]
+        assert kwargs.get("upsert") is True
+
+    @pytest.mark.asyncio
+    async def test_upsert_edge_retries_once_on_duplicate_key(self):
+        s = self._make_storage()
+        s.edge_collection.update_one = AsyncMock(
+            side_effect=[DuplicateKeyError("E11000 dup"), None]
+        )
+        await s.upsert_edge("A", "B", {"weight": 1.0})
+        assert s.edge_collection.update_one.await_count == 2
+
+    @pytest.mark.asyncio
+    async def test_upsert_edge_reraises_on_persistent_duplicate(self):
+        s = self._make_storage()
+        s.edge_collection.update_one = AsyncMock(
+            side_effect=DuplicateKeyError("E11000 dup")
+        )
+        with pytest.raises(DuplicateKeyError):
+            await s.upsert_edge("A", "B", {"weight": 1.0})
+        assert s.edge_collection.update_one.await_count == 2
+
+    @pytest.mark.asyncio
+    async def test_upsert_edges_batch_dedupes_reciprocal_and_sets_endpoints(self):
+        s = self._make_storage()
+        calls = []
+
+        async def fake_bulk(collection, ops, **kwargs):
+            calls.append((collection, ops))
+
+        with patch(
+            "lightrag.kg.mongo_impl._run_batched_bulk_write",
+            new=AsyncMock(side_effect=fake_bulk),
+        ):
+            await s.upsert_edges_batch(
+                [("A", "B", {"weight": 1.0}), ("B", "A", {"weight": 2.0})]
+            )
+
+        # Last call is the edge bulk (first is the node-placeholder bulk).
+        edge_collection, edge_ops = calls[-1]
+        assert edge_collection is s.edge_collection
+        assert len(edge_ops) == 1  # reciprocal pair collapsed to one op
+        op, _bytes, _logid = edge_ops[0]
+        lo, hi = _canonical_edge_endpoints("A", "B")
+        assert op._filter == {"edge_lo": lo, "edge_hi": hi}
+        assert op._doc["$set"]["edge_lo"] == lo
+        assert op._doc["$set"]["edge_hi"] == hi
+        assert op._doc["$set"]["weight"] == 2.0  # last-write-wins
+
+    @pytest.mark.asyncio
+    async def test_upsert_edges_batch_retries_on_duplicate_bulk_error(self):
+        s = self._make_storage()
+        seq = []
+
+        async def fake_bulk(collection, ops, **kwargs):
+            seq.append(collection)
+            # Fail the first edge bulk with an all-11000 BulkWriteError.
+            if collection is s.edge_collection and seq.count(s.edge_collection) == 1:
+                raise BulkWriteError({"writeErrors": [{"code": 11000}]})
+
+        with patch(
+            "lightrag.kg.mongo_impl._run_batched_bulk_write",
+            new=AsyncMock(side_effect=fake_bulk),
+        ):
+            await s.upsert_edges_batch([("A", "B", {"weight": 1.0})])
+
+        # node bulk + edge bulk (raises) + edge bulk (retry succeeds)
+        assert seq.count(s.edge_collection) == 2
+
+    @pytest.mark.asyncio
+    async def test_upsert_edges_batch_reraises_non_duplicate_bulk_error(self):
+        s = self._make_storage()
+
+        async def fake_bulk(collection, ops, **kwargs):
+            if collection is s.edge_collection:
+                raise BulkWriteError({"writeErrors": [{"code": 121}]})
+
+        with patch(
+            "lightrag.kg.mongo_impl._run_batched_bulk_write",
+            new=AsyncMock(side_effect=fake_bulk),
+        ):
+            with pytest.raises(BulkWriteError):
+                await s.upsert_edges_batch([("A", "B", {"weight": 1.0})])
+
+    @pytest.mark.asyncio
+    async def test_upsert_edges_batch_surfaces_non_dup_error_hidden_behind_dup(self):
+        """Under ordered=True the bulk aborts at the first failing op, so a
+        non-11000 error sitting *behind* a leading 11000 race is not reported on
+        the first pass. It must still surface — on the retry, where the leading
+        dup has resolved to an update and the real error now leads."""
+        s = self._make_storage()
+        edge_calls = []
+
+        async def fake_bulk(collection, ops, **kwargs):
+            if collection is s.edge_collection:
+                edge_calls.append(collection)
+                if len(edge_calls) == 1:
+                    # First pass: ordered abort reports only the leading dup race.
+                    raise BulkWriteError({"writeErrors": [{"code": 11000}]})
+                # Retry: dup resolved to an update, the hidden error now leads.
+                raise BulkWriteError({"writeErrors": [{"code": 121}]})
+
+        with patch(
+            "lightrag.kg.mongo_impl._run_batched_bulk_write",
+            new=AsyncMock(side_effect=fake_bulk),
+        ):
+            with pytest.raises(BulkWriteError) as exc:
+                await s.upsert_edges_batch([("A", "B", {"weight": 1.0})])
+        # Retried once (dup-only first pass), then the non-dup error re-raised.
+        assert len(edge_calls) == 2
+        assert exc.value.details["writeErrors"][0]["code"] == 121
+
+    @pytest.mark.asyncio
+    async def test_upsert_edges_batch_reraises_write_concern_error(self):
+        """A writeConcern-only BulkWriteError (empty writeErrors) is a durability
+        failure — it must surface, not be masked by the duplicate-key retry."""
+        s = self._make_storage()
+        edge_calls = []
+
+        async def fake_bulk(collection, ops, **kwargs):
+            if collection is s.edge_collection:
+                edge_calls.append(collection)
+                raise BulkWriteError(
+                    {"writeErrors": [], "writeConcernErrors": [{"code": 64}]}
+                )
+
+        with patch(
+            "lightrag.kg.mongo_impl._run_batched_bulk_write",
+            new=AsyncMock(side_effect=fake_bulk),
+        ):
+            with pytest.raises(BulkWriteError):
+                await s.upsert_edges_batch([("A", "B", {"weight": 1.0})])
+        assert len(edge_calls) == 1  # not retried
+
+    @pytest.mark.asyncio
+    async def test_edge_migration_skips_when_index_exists(self):
+        s = self._make_storage()
+        s.edge_collection.list_indexes = AsyncMock(
+            return_value=SimpleNamespace(
+                to_list=AsyncMock(return_value=[{"name": "test_edge_endpoints_unique"}])
+            )
+        )
+        s.edge_collection.aggregate = AsyncMock()
+        s.edge_collection.create_index = AsyncMock()
+
+        await s.create_edge_indexes_and_migrate_if_not_exists()
+
+        s.edge_collection.aggregate.assert_not_awaited()
+        s.edge_collection.create_index.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_edge_migration_dedupes_backfills_and_builds_unique_index(self):
+        s = self._make_storage()
+        s.edge_collection.list_indexes = AsyncMock(
+            return_value=SimpleNamespace(
+                to_list=AsyncMock(return_value=[{"name": "_id_"}])
+            )
+        )
+        # One duplicate group for edge {A,B}: two docs with distinct provenance
+        # AND distinct relation payload (description/keywords/weight).
+        group = {
+            "_id": {"lo": "A", "hi": "B"},
+            "count": 2,
+            "docs": [
+                {
+                    "_id": 1,
+                    "source_id": "c1",
+                    "source_ids": ["c1"],
+                    "file_path": "f1",
+                    "description": "d1",
+                    "keywords": "alpha,beta",
+                    "weight": 1.0,
+                    "created_at": 10,
+                },
+                {
+                    "_id": 2,
+                    "source_id": "c2",
+                    "source_ids": ["c2"],
+                    "file_path": "f2",
+                    "description": "d2",
+                    "keywords": "beta,gamma",
+                    "weight": 2.0,
+                    "created_at": 20,
+                },
+            ],
+        }
+        s.edge_collection.estimated_document_count = AsyncMock(return_value=4)
+        s.edge_collection.aggregate = AsyncMock(return_value=_AsyncCursor([group]))
+        s.edge_collection.update_one = AsyncMock()
+        s.edge_collection.delete_many = AsyncMock()
+        s.edge_collection.update_many = AsyncMock(
+            return_value=SimpleNamespace(modified_count=3)
+        )
+        s.edge_collection.create_index = AsyncMock()
+
+        with patch("lightrag.kg.mongo_impl.logger") as mock_logger:
+            await s.create_edge_indexes_and_migrate_if_not_exists()
+
+        # OpenSearch-aligned start/complete log wording.
+        info_lines = [c.args[0] for c in mock_logger.info.call_args_list]
+        assert any(
+            line.startswith("[test] Starting canonical edge migration for ")
+            and "(~4 edges to scan)" in line
+            for line in info_lines
+        )
+        assert any(
+            "Canonical edge migration complete for" in line
+            and "scanned 4, deduped 1, backfilled 3" in line
+            for line in info_lines
+        )
+
+        # Survivor is the newest (created_at 20 → _id 2); the full relation
+        # payload is merged in before the duplicate is deleted (no evidence lost).
+        surv_filter, surv_update = s.edge_collection.update_one.call_args[0]
+        set_fields = surv_update["$set"]
+        assert surv_filter == {"_id": 2}
+        assert set_fields["source_ids"] == ["c1", "c2"]
+        assert set_fields["file_path"] == "f1<SEP>f2"
+        assert set_fields["description"] == "d1<SEP>d2"  # distinct descriptions joined
+        assert set_fields["keywords"] == "alpha,beta,gamma"  # comma set-union, sorted
+        assert set_fields["weight"] == 3.0  # summed (1.0 + 2.0), idempotent
+        # The other duplicate is deleted.
+        assert s.edge_collection.delete_many.call_args[0][0] == {"_id": {"$in": [1]}}
+        # Compound unique partial index built as the completion flag.
+        ci_args = s.edge_collection.create_index.call_args.args
+        ci_kwargs = s.edge_collection.create_index.call_args.kwargs
+        assert ci_args[0] == [("edge_lo", 1), ("edge_hi", 1)]
+        assert ci_kwargs["name"] == "test_edge_endpoints_unique"
+        assert ci_kwargs["unique"] is True
+        assert set(ci_kwargs["partialFilterExpression"]) == {"edge_lo", "edge_hi"}
+
+    async def _run_dedupe(self, s, group):
+        """Run _dedupe_legacy_edges over a single group; return the survivor $set."""
+        s.edge_collection.aggregate = AsyncMock(return_value=_AsyncCursor([group]))
+        s.edge_collection.update_one = AsyncMock()
+        s.edge_collection.delete_many = AsyncMock()
+        await s._dedupe_legacy_edges()
+        if s.edge_collection.update_one.call_args is None:
+            return None
+        return s.edge_collection.update_one.call_args[0][1]["$set"]
+
+    @pytest.mark.asyncio
+    async def test_dedupe_coerces_string_weights(self):
+        """Legacy string weights (graph API is dict[str, str]) must not crash the
+        migration; they coerce to float and sum."""
+        s = self._make_storage()
+        group = {
+            "_id": {"lo": "A", "hi": "B"},
+            "count": 2,
+            "docs": [
+                {"_id": 1, "source_ids": ["c1"], "weight": "1.0", "created_at": 10},
+                {"_id": 2, "source_ids": ["c2"], "weight": "2.0", "created_at": 20},
+            ],
+        }
+        set_fields = await self._run_dedupe(s, group)
+        assert set_fields["weight"] == 3.0  # coerced floats summed, no TypeError
+
+    @pytest.mark.asyncio
+    async def test_dedupe_merge_is_idempotent_on_retry(self):
+        """A fail-fast retry that re-processes an already-merged survivor must
+        produce the same fields (no double-summed weight, no grown description)."""
+        s = self._make_storage()
+        original = [
+            {
+                "_id": 1,
+                "source_id": "c1",
+                "source_ids": ["c1"],
+                "file_path": "f1",
+                "description": "d1",
+                "keywords": "alpha,beta",
+                "weight": 1.0,
+                "created_at": 10,
+            },
+            {
+                "_id": 2,
+                "source_id": "c2",
+                "source_ids": ["c2"],
+                "file_path": "f2",
+                "description": "d2",
+                "keywords": "beta,gamma",
+                "weight": 2.0,
+                "created_at": 20,
+            },
+        ]
+        first = await self._run_dedupe(
+            s, {"_id": {"lo": "A", "hi": "B"}, "count": 2, "docs": list(original)}
+        )
+
+        # Retry: survivor (_id 2) already carries the merged fields, the other
+        # duplicate is still present (the previous delete never committed).
+        survivor_merged = {"_id": 2, "created_at": 20, **first}
+        second = await self._run_dedupe(
+            s,
+            {
+                "_id": {"lo": "A", "hi": "B"},
+                "count": 2,
+                "docs": [original[0], survivor_merged],
+            },
+        )
+
+        # Summed once (1.0 + 2.0); the retry must NOT re-add the duplicate's
+        # weight (its source_ids are already folded into the survivor).
+        assert second["weight"] == first["weight"] == 3.0
+        assert second["description"] == first["description"] == "d1<SEP>d2"
+        assert second["source_ids"] == first["source_ids"] == ["c1", "c2"]
+        assert second["file_path"] == first["file_path"] == "f1<SEP>f2"
+        assert second["keywords"] == first["keywords"] == "alpha,beta,gamma"
+
+
+class TestMongoEdgeReadCanonicalFilter:
+    """Exact single-edge reads/deletes use the canonical (edge_lo, edge_hi)
+    filter so they hit the compound unique index instead of the bidirectional
+    ``$or``. Per-node scans (node_degree/get_node_edges/delete_node) intentionally
+    keep ``$or`` — edge_hi is not a compound-index prefix, so they gain nothing."""
+
+    def _make_storage(self, max_delete_records_per_batch=1000):
+        s = MongoGraphStorage.__new__(MongoGraphStorage)
+        s.workspace = "test"
+        s.namespace = "chunk_entity_relation"
+        s._edge_collection_name = "test_edges"
+        s._max_delete_records_per_batch = max_delete_records_per_batch
+        s.edge_collection = SimpleNamespace()
+        return s
+
+    @pytest.mark.asyncio
+    async def test_has_edge_uses_canonical_filter(self):
+        s = self._make_storage()
+        s.edge_collection.find_one = AsyncMock(return_value={"_id": "x"})
+        assert await s.has_edge("B", "A") is True
+
+        filt, projection = s.edge_collection.find_one.call_args[0]
+        lo, hi = _canonical_edge_endpoints("B", "A")
+        assert filt == {"edge_lo": lo, "edge_hi": hi}
+        assert projection == {"_id": 1}
+
+    @pytest.mark.asyncio
+    async def test_has_edge_direction_independent(self):
+        s = self._make_storage()
+        s.edge_collection.find_one = AsyncMock(return_value=None)
+        await s.has_edge("A", "B")
+        forward = s.edge_collection.find_one.call_args[0][0]
+        await s.has_edge("B", "A")
+        backward = s.edge_collection.find_one.call_args[0][0]
+        assert forward == backward
+
+    @pytest.mark.asyncio
+    async def test_get_edge_uses_canonical_filter(self):
+        s = self._make_storage()
+        s.edge_collection.find_one = AsyncMock(return_value={"weight": 1.0})
+        await s.get_edge("B", "A")
+        filt = s.edge_collection.find_one.call_args[0][0]
+        lo, hi = _canonical_edge_endpoints("B", "A")
+        assert filt == {"edge_lo": lo, "edge_hi": hi}
+
+    @pytest.mark.asyncio
+    async def test_remove_edges_uses_canonical_pairs_and_collapses_reciprocals(self):
+        s = self._make_storage()
+        s.edge_collection.delete_many = AsyncMock()
+        # (A,B) and its reciprocal (B,A) must collapse to a single clause.
+        await s.remove_edges([("A", "B"), ("B", "A"), ("C", "D")])
+
+        query = s.edge_collection.delete_many.call_args[0][0]
+        lo_ab, hi_ab = _canonical_edge_endpoints("A", "B")
+        lo_cd, hi_cd = _canonical_edge_endpoints("C", "D")
+        assert query == {
+            "$or": [
+                {"edge_lo": lo_ab, "edge_hi": hi_ab},
+                {"edge_lo": lo_cd, "edge_hi": hi_cd},
+            ]
+        }
+
+    @pytest.mark.asyncio
+    async def test_remove_edges_empty_is_noop(self):
+        s = self._make_storage()
+        s.edge_collection.delete_many = AsyncMock()
+        await s.remove_edges([])
+        s.edge_collection.delete_many.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_remove_edges_chunks_large_or_by_record_cap(self):
+        # 5 distinct edges with a cap of 2 → 3 delete_many calls (2 + 2 + 1),
+        # and every input edge is covered exactly once across the chunks.
+        s = self._make_storage(max_delete_records_per_batch=2)
+        s.edge_collection.delete_many = AsyncMock()
+        edges = [(f"n{i}", f"n{i + 1}") for i in range(5)]
+        await s.remove_edges(edges)
+
+        calls = s.edge_collection.delete_many.call_args_list
+        assert [len(c[0][0]["$or"]) for c in calls] == [2, 2, 1]
+        clauses = [clause for c in calls for clause in c[0][0]["$or"]]
+        expected = [
+            {"edge_lo": lo, "edge_hi": hi}
+            for lo, hi in (_canonical_edge_endpoints(src, tgt) for src, tgt in edges)
+        ]
+        assert clauses == expected
+
+    @pytest.mark.asyncio
+    async def test_remove_edges_non_positive_cap_disables_chunking(self):
+        # A non-positive cap means "no record-count splitting": one delete_many.
+        s = self._make_storage(max_delete_records_per_batch=0)
+        s.edge_collection.delete_many = AsyncMock()
+        edges = [(f"n{i}", f"n{i + 1}") for i in range(5)]
+        await s.remove_edges(edges)
+
+        assert s.edge_collection.delete_many.await_count == 1
+        assert len(s.edge_collection.delete_many.call_args[0][0]["$or"]) == 5
+
+
+class TestMongoDocStatusLookup:
+    """Cover the Mongo-native overrides for basename / content_hash lookups."""
+
+    def _make_storage(self):
+        storage = MongoDocStatusStorage.__new__(MongoDocStatusStorage)
+        storage.workspace = "test"
+        storage.global_config = {}
+        storage._collection_name = "test_doc_status"
+        storage._data = SimpleNamespace()
+        return storage
+
+    @pytest.mark.asyncio
+    async def test_get_doc_by_file_basename_returns_tuple_on_hit(self):
+        storage = self._make_storage()
+        storage._data.find_one = AsyncMock(
+            return_value={
+                "_id": "doc-1",
+                "file_path": "report.pdf",
+                "status": "processed",
+            }
+        )
+
+        result = await storage.get_doc_by_file_basename("report.pdf")
+
+        assert result is not None
+        doc_id, doc = result
+        assert doc_id == "doc-1"
+        assert doc["file_path"] == "report.pdf"
+        storage._data.find_one.assert_awaited_once_with({"file_path": "report.pdf"})
+
+    @pytest.mark.asyncio
+    async def test_get_doc_by_file_basename_empty_returns_none_without_query(self):
+        storage = self._make_storage()
+        storage._data.find_one = AsyncMock()
+
+        assert await storage.get_doc_by_file_basename("") is None
+        storage._data.find_one.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_get_doc_by_file_basename_unknown_source_sentinel(self):
+        # Lookup for the sentinel must not match real rows that happen to have
+        # file_path == "unknown_source".
+        storage = self._make_storage()
+        storage._data.find_one = AsyncMock()
+
+        assert await storage.get_doc_by_file_basename("unknown_source") is None
+        storage._data.find_one.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_get_doc_by_file_basename_miss_returns_none(self):
+        storage = self._make_storage()
+        storage._data.find_one = AsyncMock(return_value=None)
+
+        assert await storage.get_doc_by_file_basename("missing.pdf") is None
+
+    @pytest.mark.asyncio
+    async def test_get_doc_by_content_hash_returns_tuple_on_hit(self):
+        storage = self._make_storage()
+        storage._data.find_one = AsyncMock(
+            return_value={
+                "_id": "doc-1",
+                "file_path": "report.pdf",
+                "content_hash": "abc123",
+                "status": "processed",
+            }
+        )
+
+        result = await storage.get_doc_by_content_hash("abc123")
+
+        assert result is not None
+        doc_id, doc = result
+        assert doc_id == "doc-1"
+        assert doc["content_hash"] == "abc123"
+        storage._data.find_one.assert_awaited_once_with({"content_hash": "abc123"})
+
+    @pytest.mark.asyncio
+    async def test_get_doc_by_content_hash_empty_returns_none_without_query(self):
+        # Empty hash must short-circuit so it cannot match legacy rows missing
+        # the field via accidental coercion.
+        storage = self._make_storage()
+        storage._data.find_one = AsyncMock()
+
+        assert await storage.get_doc_by_content_hash("") is None
+        storage._data.find_one.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_get_doc_by_content_hash_miss_returns_none(self):
+        storage = self._make_storage()
+        storage._data.find_one = AsyncMock(return_value=None)
+
+        assert await storage.get_doc_by_content_hash("zzz999") is None
+
+    @pytest.mark.asyncio
+    async def test_lookup_swallows_pymongo_error_and_returns_none(self):
+        # PyMongoError must not propagate to the caller; the dedup path treats
+        # a storage failure as "no match" and the error is logged instead.
+        storage = self._make_storage()
+        storage._data.find_one = AsyncMock(side_effect=PyMongoError("boom"))
+
+        assert await storage.get_doc_by_file_basename("report.pdf") is None
+        assert await storage.get_doc_by_content_hash("abc123") is None

tests/kg/nano_impl/test_atomic_write_nano.py

@@ -0,0 +1,99 @@
+"""Atomicity coverage for the NanoVectorDB save path.
+
+``NanoVectorDB.save()`` is third-party and writes via plain
+``open(self.storage_file, "w") + json.dump``. ``NanoVectorDBStorage`` wraps
+that call by swapping ``storage_file`` to a per-writer tmp under
+``atomic_write``. The contract we have to preserve is:
+
+1. On success: the destination file is updated, no tmp residue remains.
+2. On failure: the destination is unchanged, ``client.storage_file`` is
+   restored to the real path (so subsequent reads don't keep pointing at
+   the tmp).
+"""
+
+import json
+import os
+from unittest.mock import patch
+
+import pytest
+
+nano_vectordb = pytest.importorskip("nano_vectordb")  # noqa: F841
+
+from lightrag.file_atomic import atomic_write  # noqa: E402
+from nano_vectordb import NanoVectorDB  # noqa: E402
+
+
+def _make_client(tmp_path) -> tuple[NanoVectorDB, str]:
+    target = str(tmp_path / "vdb_test.json")
+    client = NanoVectorDB(4, storage_file=target)
+    return client, target
+
+
+def _save_atomic(client: NanoVectorDB, target: str) -> None:
+    """Mirrors the production callback's save lambda — kept inline so the
+    test exercises the exact swap-and-restore pattern."""
+
+    def _do(tmp: str) -> None:
+        original = client.storage_file
+        client.storage_file = tmp
+        try:
+            client.save()
+        finally:
+            client.storage_file = original
+
+    atomic_write(target, _do)
+
+
+@pytest.mark.offline
+def test_nano_save_atomic_publishes_file_and_restores_storage_file(tmp_path):
+    client, target = _make_client(tmp_path)
+    _save_atomic(client, target)
+
+    assert os.path.exists(target)
+    assert client.storage_file == target
+    # Sanity check the payload is valid JSON written by NanoVectorDB.
+    json.load(open(target))
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"Unexpected tmp residue: {leftovers}"
+
+
+@pytest.mark.offline
+def test_nano_save_atomic_replace_crash_preserves_prior(tmp_path):
+    client, target = _make_client(tmp_path)
+    _save_atomic(client, target)
+    original_payload = open(target).read()
+
+    with patch(
+        "lightrag.file_atomic.os.replace",
+        side_effect=OSError("simulated crash"),
+    ):
+        with pytest.raises(OSError, match="simulated crash"):
+            _save_atomic(client, target)
+
+    # Destination unchanged.
+    assert open(target).read() == original_payload
+    # storage_file must have been restored even though the outer atomic_write
+    # raised — otherwise NanoVectorDB would silently start writing to a path
+    # that no longer exists.
+    assert client.storage_file == target
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"Python-exception path must clean tmp, got {leftovers}"
+
+
+@pytest.mark.offline
+def test_nano_save_failure_inside_save_restores_storage_file(tmp_path):
+    """If ``NanoVectorDB.save()`` itself raises (e.g. encoding failure), the
+    inner ``finally`` must restore ``storage_file`` before the exception
+    bubbles up through ``atomic_write``."""
+    client, target = _make_client(tmp_path)
+
+    # Patch the third-party save to blow up.
+    with patch.object(NanoVectorDB, "save", side_effect=RuntimeError("boom")):
+        with pytest.raises(RuntimeError, match="boom"):
+            _save_atomic(client, target)
+
+    assert client.storage_file == target, (
+        "Inner save failure must still restore storage_file to the real path"
+    )
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"save failure must clean tmp, got {leftovers}"

tests/kg/nano_impl/test_nano_deferred_embedding.py

@@ -0,0 +1,450 @@
+"""Deferred-embedding coverage for ``NanoVectorDBStorage``.
+
+The storage no longer embeds eagerly in ``upsert``: it buffers a pending doc
+and embeds once per id at flush time (``index_done_callback`` / ``finalize``).
+These tests pin that contract using a counting mock embedding function — no
+live model or network. They mirror the protocol proven for
+``OpenSearchVectorDBStorage`` (issue #2785).
+"""
+
+import numpy as np
+import pytest
+
+nano_vectordb = pytest.importorskip("nano_vectordb")  # noqa: F841
+
+from lightrag.kg.nano_vector_db_impl import NanoVectorDBStorage  # noqa: E402
+from lightrag.kg.shared_storage import (  # noqa: E402
+    initialize_share_data,
+    finalize_share_data,
+)
+from lightrag.utils import EmbeddingFunc  # noqa: E402
+
+DIM = 8
+
+
+@pytest.fixture(autouse=True)
+def _shared_data():
+    finalize_share_data()
+    initialize_share_data()
+    yield
+    finalize_share_data()
+
+
+class _CountingEmbed:
+    """Async embedding callable that records how many texts it embedded and how
+    many times it was invoked (one invocation == one batch)."""
+
+    def __init__(self, dim: int = DIM):
+        self.dim = dim
+        self.call_count = 0
+        self.embedded_texts: list[str] = []
+
+    async def __call__(self, texts, **kwargs):
+        self.call_count += 1
+        self.embedded_texts.extend(texts)
+        # Deterministic per-text vector so duplicates are still 1-1.
+        return np.array(
+            [
+                np.full(self.dim, (abs(hash(t)) % 97) + 1, dtype=np.float32)
+                for t in texts
+            ]
+        )
+
+
+def _make_storage(tmp_path, embed: _CountingEmbed) -> NanoVectorDBStorage:
+    return NanoVectorDBStorage(
+        namespace="test_vectors",
+        workspace="ws",
+        global_config={
+            "working_dir": str(tmp_path),
+            "embedding_batch_num": 32,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=EmbeddingFunc(embedding_dim=DIM, max_token_size=512, func=embed),
+        meta_fields={"content"},
+    )
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_upsert_defers_embedding_to_index_done_callback(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert(
+        {
+            "id1": {"content": "alpha"},
+            "id2": {"content": "beta"},
+        }
+    )
+    assert embed.call_count == 0, "upsert must not embed"
+    assert len(storage._client) == 0, "nothing should be materialized yet"
+
+    await storage.index_done_callback()
+    assert embed.call_count == 1, "flush should embed in a single batch"
+    assert sorted(embed.embedded_texts) == ["alpha", "beta"]
+    assert len(storage._client) == 2
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_repeated_upserts_same_id_embed_once_per_flush(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "v1"}})
+    await storage.upsert({"id1": {"content": "v2"}})
+    await storage.upsert({"id1": {"content": "v3"}})
+
+    await storage.index_done_callback()
+
+    assert embed.call_count == 1
+    assert embed.embedded_texts == ["v3"], "only the latest content is embedded"
+    assert len(storage._client) == 1
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_get_vectors_caches_and_flush_reuses(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    vecs = await storage.get_vectors_by_ids(["id1"])
+    assert "id1" in vecs and len(vecs["id1"]) == DIM
+    assert embed.call_count == 1, "get_vectors_by_ids embeds pending lazily"
+
+    # Flush must reuse the cached vector, not re-embed.
+    await storage.index_done_callback()
+    assert embed.call_count == 1, "flush should reuse the cached temp vector"
+    assert len(storage._client) == 1
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_reupsert_after_get_vectors_clears_cached_vector(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "old"}})
+    await storage.get_vectors_by_ids(["id1"])  # caches a temp vector for "old"
+    assert embed.call_count == 1
+
+    # New content version must clear the cached vector and re-embed at flush.
+    await storage.upsert({"id1": {"content": "new"}})
+    await storage.index_done_callback()
+
+    assert embed.call_count == 2
+    assert embed.embedded_texts == ["old", "new"]
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_delete_cancels_pending_and_removes_materialized(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    # Materialize id1; leave id2 only as a pending (unflushed) upsert.
+    await storage.upsert({"id1": {"content": "alpha"}})
+    await storage.index_done_callback()
+    await storage.upsert({"id2": {"content": "beta"}})
+
+    await storage.delete(["id1", "id2"])
+
+    assert "id2" not in storage._pending_upserts, "delete cancels pending upsert"
+    assert len(storage._client) == 0, "delete removes the materialized row immediately"
+    assert await storage.get_by_id("id1") is None
+    assert await storage.get_by_id("id2") is None
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_stale_client_reload_still_flushes_pending_upsert(tmp_path):
+    embed = _CountingEmbed()
+    writer = _make_storage(tmp_path, embed)
+    stale_writer = _make_storage(tmp_path, embed)
+    await writer.initialize()
+    await stale_writer.initialize()
+
+    await writer.upsert({"id1": {"content": "alpha"}})
+    assert await writer.index_done_callback() is True
+    assert stale_writer.storage_updated.value is True
+
+    await stale_writer.upsert({"id2": {"content": "beta"}})
+    assert await stale_writer.index_done_callback() is True
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    rows = await reader.get_by_ids(["id1", "id2"])
+    assert [row["id"] for row in rows] == ["id1", "id2"]
+    assert stale_writer._pending_upserts == {}
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_delete_reloads_stale_client_before_mutating(tmp_path):
+    embed = _CountingEmbed()
+    writer = _make_storage(tmp_path, embed)
+    stale_deleter = _make_storage(tmp_path, embed)
+    await writer.initialize()
+    await stale_deleter.initialize()
+
+    await writer.upsert({"id1": {"content": "alpha"}})
+    assert await writer.index_done_callback() is True
+    assert stale_deleter.storage_updated.value is True
+
+    await stale_deleter.delete(["id1"])
+    assert stale_deleter.storage_updated.value is False
+    assert await stale_deleter.index_done_callback() is True
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    assert await reader.get_by_id("id1") is None
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_finalize_reloads_stale_client_before_flushing(tmp_path):
+    embed = _CountingEmbed()
+    writer = _make_storage(tmp_path, embed)
+    stale_finalizer = _make_storage(tmp_path, embed)
+    await writer.initialize()
+    await stale_finalizer.initialize()
+
+    await writer.upsert({"id1": {"content": "alpha"}})
+    assert await writer.index_done_callback() is True
+    assert stale_finalizer.storage_updated.value is True
+
+    await stale_finalizer.upsert({"id2": {"content": "beta"}})
+    await stale_finalizer.finalize()
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    rows = await reader.get_by_ids(["id1", "id2"])
+    assert [row["id"] for row in rows] == ["id1", "id2"]
+    assert stale_finalizer._pending_upserts == {}
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_read_your_writes_and_query_after_flush(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    # Before flush: read paths see the pending row, query does not.
+    hit = await storage.get_by_id("id1")
+    assert hit is not None and hit["id"] == "id1" and hit["content"] == "alpha"
+    by_ids = await storage.get_by_ids(["id1", "missing"])
+    assert by_ids[0]["id"] == "id1" and by_ids[1] is None
+    assert await storage.query("alpha", top_k=5) == [], "query ignores unflushed data"
+
+    # After flush: query returns the row.
+    await storage.index_done_callback()
+    results = await storage.query("alpha", top_k=5)
+    assert any(r["id"] == "id1" for r in results)
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_finalize_flushes_pending(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+    await storage.finalize()
+
+    assert embed.call_count == 1
+    assert storage._pending_upserts == {}
+    assert len(storage._client) == 1
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_delete_entity_relation_cancels_pending(tmp_path):
+    embed = _CountingEmbed()
+    storage = NanoVectorDBStorage(
+        namespace="test_relations",
+        workspace="ws",
+        global_config={
+            "working_dir": str(tmp_path),
+            "embedding_batch_num": 32,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=EmbeddingFunc(embedding_dim=DIM, max_token_size=512, func=embed),
+        meta_fields={"content", "src_id", "tgt_id"},
+    )
+    await storage.initialize()
+
+    # Materialize r1 (A->B), leave r2 (A->C) and r3 (X->Y) as pending.
+    await storage.upsert({"r1": {"content": "rel1", "src_id": "A", "tgt_id": "B"}})
+    await storage.index_done_callback()
+    await storage.upsert(
+        {
+            "r2": {"content": "rel2", "src_id": "A", "tgt_id": "C"},
+            "r3": {"content": "rel3", "src_id": "X", "tgt_id": "Y"},
+        }
+    )
+
+    await storage.delete_entity_relation("A")
+
+    assert "r2" not in storage._pending_upserts, "incident pending entry cancelled"
+    assert "r3" in storage._pending_upserts, "unrelated pending entry preserved"
+    assert len(storage._client) == 0, "materialized A->B removed"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_flush_embedding_failure_raises_and_keeps_pending(tmp_path):
+    class _FailingEmbed:
+        def __init__(self):
+            self.call_count = 0
+
+        async def __call__(self, texts, **kwargs):
+            self.call_count += 1
+            raise RuntimeError("embed boom")
+
+    embed = _FailingEmbed()
+    storage = NanoVectorDBStorage(
+        namespace="test_vectors",
+        workspace="ws",
+        global_config={
+            "working_dir": str(tmp_path),
+            "embedding_batch_num": 32,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.2},
+        },
+        embedding_func=EmbeddingFunc(embedding_dim=DIM, max_token_size=512, func=embed),
+        meta_fields={"content"},
+    )
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    with pytest.raises(RuntimeError, match="embed boom"):
+        await storage.index_done_callback()
+
+    assert "id1" in storage._pending_upserts, "pending preserved for retry"
+    assert len(storage._client) == 0, "nothing materialized on embed failure"
+    # Embed failure happens before self._client.upsert in _flush_pending_locked,
+    # so _client_dirty must NOT be set. (A save-stage failure would leave it True
+    # — see test_finalize_retries_save_after_flush_failure.)
+    assert storage._client_dirty is False
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_drop_discards_pending_without_embedding(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+    assert "id1" in storage._pending_upserts
+
+    result = await storage.drop()
+
+    assert result["status"] == "success"
+    assert storage._pending_upserts == {}, "drop discards buffered upserts"
+    assert embed.call_count == 0, "drop must not embed"
+    assert storage._client_dirty is False
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_finalize_retries_save_after_flush_failure(tmp_path):
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    original_save = storage._save_to_disk_locked
+    save_calls = 0
+
+    def fail_once():
+        nonlocal save_calls
+        save_calls += 1
+        if save_calls == 1:
+            raise OSError("boom")
+        original_save()
+
+    storage._save_to_disk_locked = fail_once
+
+    with pytest.raises(OSError, match="boom"):
+        await storage.finalize()
+
+    assert storage._pending_upserts == {}
+    assert storage._client_dirty is True
+
+    await storage.finalize()
+
+    assert save_calls == 2
+    assert storage._client_dirty is False
+
+    reader = _make_storage(tmp_path, embed)
+    await reader.initialize()
+    hit = await reader.get_by_id("id1")
+    assert hit is not None and hit["id"] == "id1"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_drop_pending_index_ops_clears_buffer(tmp_path):
+    """An internal-error abort calls drop_pending_index_ops to discard the
+    not-yet-flushed buffer without materializing anything."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    await storage.upsert({"id1": {"content": "alpha"}, "id2": {"content": "beta"}})
+    assert storage._pending_upserts, "upsert buffers, does not flush"
+
+    await storage.drop_pending_index_ops()
+
+    assert storage._pending_upserts == {}
+    assert embed.call_count == 0, "drop must not embed"
+    assert len(storage._client) == 0, "nothing was materialized"
+
+
+@pytest.mark.offline
+@pytest.mark.asyncio
+async def test_drop_pending_does_not_rollback_materialized(tmp_path):
+    """drop_pending_index_ops discards ONLY the pending buffer; records already
+    materialized into self._client by a flush whose save then failed
+    (``_client_dirty=True``) are intentionally NOT rolled back."""
+    embed = _CountingEmbed()
+    storage = _make_storage(tmp_path, embed)
+    await storage.initialize()
+
+    # Flush id1 into the in-memory client, then fail the save so it stays
+    # materialized-but-unsaved (dirty) and the pending buffer is emptied.
+    await storage.upsert({"id1": {"content": "alpha"}})
+
+    def fail_save():
+        raise OSError("save boom")
+
+    storage._save_to_disk_locked = fail_save
+    with pytest.raises(OSError, match="save boom"):
+        await storage.index_done_callback()
+    assert storage._pending_upserts == {}, "flush succeeded so pending is empty"
+    assert storage._client_dirty is True
+    assert len(storage._client) == 1, "id1 materialized, not saved"
+
+    # A new pending op arrives, then the batch aborts and drops pending.
+    await storage.upsert({"id2": {"content": "beta"}})
+    assert "id2" in storage._pending_upserts
+
+    await storage.drop_pending_index_ops()
+
+    assert storage._pending_upserts == {}, "pending id2 dropped"
+    assert len(storage._client) == 1, "materialized id1 NOT rolled back"
+    assert storage._client_dirty is True, "still dirty for a later save retry"

tests/kg/neo4j_impl/test_neo4j_fulltext_index.py

@@ -0,0 +1,361 @@
+#!/usr/bin/env python
+"""
+Test Neo4j full-text index functionality, specifically:
+1. Workspace-specific index naming
+2. Legacy index migration
+3. search_labels functionality with workspace-specific indexes
+"""
+
+import asyncio
+import os
+import sys
+import pytest
+import numpy as np
+
+# Add the project root directory to the Python path
+sys.path.append(
+    os.path.dirname(
+        os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    )
+)
+
+from lightrag.kg.shared_storage import initialize_share_data
+
+
+# Mock embedding function that returns random vectors
+async def mock_embedding_func(texts):
+    return np.random.rand(len(texts), 10)  # Return 10-dimensional random vectors
+
+
+@pytest.fixture
+async def neo4j_storage():
+    """
+    Initialize Neo4j storage for testing.
+    Requires Neo4j to be running and configured via environment variables.
+    """
+    # Check if Neo4j is configured
+    if not os.getenv("NEO4J_URI"):
+        pytest.skip("Neo4j not configured (NEO4J_URI not set)")
+
+    from lightrag.kg.neo4j_impl import Neo4JStorage
+
+    # Initialize shared_storage for locks
+    initialize_share_data()
+
+    global_config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.5},
+        "working_dir": os.environ.get("WORKING_DIR", "./rag_storage"),
+    }
+
+    storage = Neo4JStorage(
+        namespace="test_fulltext_index",
+        workspace="test_workspace",
+        global_config=global_config,
+        embedding_func=mock_embedding_func,
+    )
+
+    # Initialize the connection
+    await storage.initialize()
+
+    # Clean up any existing data
+    await storage.drop()
+
+    yield storage
+
+    # Cleanup
+    await storage.drop()
+    await storage.finalize()
+
+
+@pytest.mark.integration
+@pytest.mark.requires_db
+async def test_fulltext_index_creation(neo4j_storage):
+    """
+    Test that the full-text index is created with the workspace-specific name.
+    """
+    storage = neo4j_storage
+    workspace_label = storage._get_workspace_label()
+    expected_index_name = f"entity_id_fulltext_idx_{workspace_label}"
+
+    # Query Neo4j to check if the index exists
+    async with storage._driver.session(database=storage._DATABASE) as session:
+        result = await session.run("SHOW FULLTEXT INDEXES")
+        indexes = await result.data()
+        await result.consume()
+
+        # Check if the workspace-specific index exists
+        index_names = [idx["name"] for idx in indexes]
+        assert expected_index_name in index_names, (
+            f"Expected index '{expected_index_name}' not found. Found indexes: {index_names}"
+        )
+
+        # Check if the legacy index doesn't exist (should be migrated if it was there)
+        legacy_index_name = "entity_id_fulltext_idx"
+        if legacy_index_name in index_names:
+            # If legacy index exists, it should be for a different workspace
+            # or it means migration didn't happen
+            print(
+                f"Warning: Legacy index '{legacy_index_name}' still exists alongside '{expected_index_name}'"
+            )
+
+
+@pytest.mark.integration
+@pytest.mark.requires_db
+async def test_search_labels_with_workspace_index(neo4j_storage):
+    """
+    Test that search_labels uses the workspace-specific index and returns results.
+    """
+    storage = neo4j_storage
+
+    # Insert test nodes
+    test_nodes = [
+        {
+            "entity_id": "Artificial Intelligence",
+            "description": "AI field",
+            "keywords": "AI,ML,DL",
+            "entity_type": "Technology",
+        },
+        {
+            "entity_id": "Machine Learning",
+            "description": "ML subfield",
+            "keywords": "supervised,unsupervised",
+            "entity_type": "Technology",
+        },
+        {
+            "entity_id": "Deep Learning",
+            "description": "DL subfield",
+            "keywords": "neural networks",
+            "entity_type": "Technology",
+        },
+        {
+            "entity_id": "Natural Language Processing",
+            "description": "NLP field",
+            "keywords": "text,language",
+            "entity_type": "Technology",
+        },
+    ]
+
+    for node_data in test_nodes:
+        await storage.upsert_node(node_data["entity_id"], node_data)
+
+    # Give the index time to become consistent (eventually consistent index)
+    await asyncio.sleep(2)
+
+    # Test search_labels
+    results = await storage.search_labels("Learning", limit=10)
+
+    # Should find nodes with "Learning" in them
+    assert len(results) > 0, "search_labels should return results for 'Learning'"
+    assert any("Learning" in result for result in results), (
+        "Results should contain 'Learning'"
+    )
+
+    # Test case-insensitive search
+    results_lower = await storage.search_labels("learning", limit=10)
+    assert len(results_lower) > 0, "search_labels should be case-insensitive"
+
+    # Test partial match
+    results_partial = await storage.search_labels("Intelli", limit=10)
+    assert len(results_partial) > 0, (
+        "search_labels should support partial matching with wildcard"
+    )
+    assert any("Intelligence" in result for result in results_partial), (
+        "Should find 'Artificial Intelligence'"
+    )
+
+
+@pytest.mark.integration
+@pytest.mark.requires_db
+async def test_search_labels_with_hyphen(neo4j_storage):
+    """
+    Regression test: searching a label fragment containing a hyphen (or other
+    Lucene reserved character) must still return matches via the full-text index
+    instead of being misparsed into an empty result.
+    """
+    storage = neo4j_storage
+
+    test_nodes = [
+        {"entity_id": "tb-alpha", "entity_type": "Test"},
+        {"entity_id": "tb-beta", "entity_type": "Test"},
+        {"entity_id": "tbcd", "entity_type": "Test"},
+    ]
+
+    for node_data in test_nodes:
+        await storage.upsert_node(node_data["entity_id"], node_data)
+
+    # Give the eventually-consistent index time to catch up.
+    await asyncio.sleep(2)
+
+    # Plain prefix returns all three (baseline that already worked).
+    results = await storage.search_labels("tb", limit=10)
+    for expected in ("tb-alpha", "tb-beta", "tbcd"):
+        assert expected in results, f"'{expected}' should match 'tb', got {results}"
+
+    # The reported bug: a trailing hyphen previously cleared the dropdown.
+    results_hyphen = await storage.search_labels("tb-", limit=10)
+    assert "tb-alpha" in results_hyphen, (
+        f"'tb-' should match 'tb-alpha', got {results_hyphen}"
+    )
+    assert "tb-beta" in results_hyphen, (
+        f"'tb-' should match 'tb-beta', got {results_hyphen}"
+    )
+
+    # An exact hyphenated label resolves to itself.
+    results_exact = await storage.search_labels("tb-alpha", limit=10)
+    assert "tb-alpha" in results_exact, (
+        f"'tb-alpha' should match itself, got {results_exact}"
+    )
+
+
+@pytest.mark.integration
+@pytest.mark.requires_db
+async def test_search_labels_chinese_text(neo4j_storage):
+    """
+    Test that search_labels works with Chinese text using the CJK analyzer.
+    """
+    storage = neo4j_storage
+
+    # Insert Chinese test nodes
+    chinese_nodes = [
+        {
+            "entity_id": "人工智能",
+            "description": "人工智能领域",
+            "keywords": "AI,机器学习",
+            "entity_type": "技术",
+        },
+        {
+            "entity_id": "机器学习",
+            "description": "机器学习子领域",
+            "keywords": "监督学习,无监督学习",
+            "entity_type": "技术",
+        },
+        {
+            "entity_id": "深度学习",
+            "description": "深度学习子领域",
+            "keywords": "神经网络",
+            "entity_type": "技术",
+        },
+    ]
+
+    for node_data in chinese_nodes:
+        await storage.upsert_node(node_data["entity_id"], node_data)
+
+    # Give the index time to become consistent
+    await asyncio.sleep(2)
+
+    # Test Chinese text search
+    results = await storage.search_labels("学习", limit=10)
+
+    # Should find nodes with "学习" in them
+    assert len(results) > 0, "search_labels should return results for Chinese text"
+    assert any("学习" in result for result in results), (
+        "Results should contain Chinese characters '学习'"
+    )
+
+
+@pytest.mark.integration
+@pytest.mark.requires_db
+async def test_search_labels_fallback_to_contains(neo4j_storage):
+    """
+    Test that search_labels falls back to CONTAINS search if the index fails.
+    This can happen with older Neo4j versions or if the index is not yet available.
+    """
+    storage = neo4j_storage
+
+    # Insert test nodes
+    test_nodes = [
+        {
+            "entity_id": "Test Node Alpha",
+            "description": "Test node",
+            "keywords": "test",
+            "entity_type": "Test",
+        },
+        {
+            "entity_id": "Test Node Beta",
+            "description": "Test node",
+            "keywords": "test",
+            "entity_type": "Test",
+        },
+    ]
+
+    for node_data in test_nodes:
+        await storage.upsert_node(node_data["entity_id"], node_data)
+
+    # Even if the full-text index is not available, CONTAINS should work
+    results = await storage.search_labels("Alpha", limit=10)
+
+    # Should find the node using fallback CONTAINS search
+    assert len(results) > 0, "Fallback CONTAINS search should return results"
+    assert "Test Node Alpha" in results, "Should find 'Test Node Alpha'"
+
+
+@pytest.mark.integration
+@pytest.mark.requires_db
+async def test_multiple_workspaces_have_separate_indexes(neo4j_storage):
+    """
+    Test that different workspaces have their own separate indexes.
+    """
+    from lightrag.kg.neo4j_impl import Neo4JStorage
+
+    # Create storage for workspace 1
+    storage1 = neo4j_storage
+
+    # Create storage for workspace 2
+    global_config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.5},
+        "working_dir": os.environ.get("WORKING_DIR", "./rag_storage"),
+    }
+
+    storage2 = Neo4JStorage(
+        namespace="test_fulltext_index",
+        workspace="test_workspace_2",
+        global_config=global_config,
+        embedding_func=mock_embedding_func,
+    )
+
+    await storage2.initialize()
+    await storage2.drop()
+
+    try:
+        # Check that both workspaces have their own indexes
+        async with storage1._driver.session(database=storage1._DATABASE) as session:
+            result = await session.run("SHOW FULLTEXT INDEXES")
+            indexes = await result.data()
+            await result.consume()
+
+            index_names = [idx["name"] for idx in indexes]
+            workspace1_index = (
+                f"entity_id_fulltext_idx_{storage1._get_workspace_label()}"
+            )
+            workspace2_index = (
+                f"entity_id_fulltext_idx_{storage2._get_workspace_label()}"
+            )
+
+            assert workspace1_index in index_names, (
+                f"Workspace 1 index '{workspace1_index}' should exist"
+            )
+            assert workspace2_index in index_names, (
+                f"Workspace 2 index '{workspace2_index}' should exist"
+            )
+
+    finally:
+        # Clean up: drop the fulltext index created for workspace 2 to prevent accumulation
+        try:
+            async with storage2._driver.session(database=storage2._DATABASE) as session:
+                index_name = storage2._get_fulltext_index_name(
+                    storage2._get_workspace_label()
+                )
+                drop_query = f"DROP INDEX {index_name} IF EXISTS"
+                result = await session.run(drop_query)
+                await result.consume()
+        except Exception:
+            pass  # Ignore errors during cleanup
+        await storage2.drop()
+        await storage2.finalize()
+
+
+if __name__ == "__main__":
+    # Run tests with pytest
+    pytest.main([__file__, "-v", "--run-integration"])

tests/kg/neo4j_impl/test_search_labels_sanitize.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+"""
+Unit tests for Neo4JStorage._sanitize_fulltext_query.
+
+These run without a live Neo4j instance and guard the regression where a label
+search containing Lucene reserved characters (e.g. "tb-") was misparsed by the
+full-text query parser ('-' as NOT) and silently returned nothing.
+"""
+
+import os
+import sys
+
+import pytest
+
+# Add the project root directory to the Python path
+sys.path.append(
+    os.path.dirname(
+        os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    )
+)
+
+from lightrag.kg.neo4j_impl import Neo4JStorage
+
+
+@pytest.mark.parametrize(
+    "raw, expected",
+    [
+        # Reserved characters are replaced with spaces and whitespace collapsed.
+        ("tb-", "tb"),
+        ("tb-foo", "tb foo"),
+        ("node:x", "node x"),
+        ("a (b)", "a b"),
+        ("C++", "C"),
+        ("foo && bar", "foo bar"),
+        ("name*", "name"),
+        # Composed entirely of reserved characters -> empty (caller returns []).
+        ("---", ""),
+        ("+++", ""),
+        ("()[]{}", ""),
+        # No reserved characters -> unchanged (modulo whitespace collapsing).
+        ("machine learning", "machine learning"),
+        ("  spaced   out  ", "spaced out"),
+        ("学习", "学习"),
+        ("机器-学习", "机器 学习"),
+    ],
+)
+def test_sanitize_fulltext_query(raw, expected):
+    assert Neo4JStorage._sanitize_fulltext_query(raw) == expected

tests/kg/neo4j_impl/test_workspace_label_injection.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python
+"""
+Unit tests guarding against Cypher injection through the workspace label.
+
+The workspace label is interpolated into two different syntactic contexts in
+``neo4j_impl.py``:
+
+1. Backtick-quoted identifiers, e.g. ``MATCH (n:`{label}`)`` — here backticks
+   must be doubled and all other characters (including single quotes) are
+   literal.
+2. The APOC ``labelFilter`` config, which lives inside a single-quoted Cypher
+   string literal. Cypher string literals are NOT escaped by doubling quotes
+   (that is SQL); they use backslash escaping. Rather than escape by hand, the
+   label is bound as a query parameter so Neo4j handles it safely.
+
+These tests run without a live Neo4j instance.
+"""
+
+import os
+import sys
+
+import pytest
+
+# Add the project root directory to the Python path
+sys.path.append(
+    os.path.dirname(
+        os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    )
+)
+
+from lightrag.kg.neo4j_impl import Neo4JStorage
+
+
+def _make_storage(workspace: str) -> Neo4JStorage:
+    """Create a Neo4JStorage with a given workspace (no connection)."""
+    return Neo4JStorage(
+        namespace="test",
+        global_config={},
+        embedding_func=None,
+        workspace=workspace if workspace else None,
+    )
+
+
+# ---------------------------------------------------------------------------
+# _get_workspace_label: backtick-identifier context
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "workspace, expected",
+    [
+        # Normal names pass through unchanged.
+        ("my-project", "my-project"),
+        ("base", "base"),
+        ("workspace_123", "workspace_123"),
+        # Whitespace is stripped; empty falls back to "base".
+        ("  trimmed  ", "trimmed"),
+        ("", "base"),
+        ("   ", "base"),
+        # Backticks are doubled for identifier context.
+        ("ws`name", "ws``name"),
+        ("a`b`c", "a``b``c"),
+        ("ws`} RETURN 0", "ws``} RETURN 0"),
+        # Single quotes are NOT touched here: inside backticks they are literal,
+        # so doubling them would corrupt the label name.
+        ("ws'name", "ws'name"),
+        ("team'a", "team'a"),
+        ("base' OR 1=1", "base' OR 1=1"),
+        # Mixed: only backticks are escaped.
+        ("ws'`name", "ws'``name"),
+    ],
+)
+def test_get_workspace_label(workspace, expected):
+    storage = _make_storage(workspace)
+    assert storage._get_workspace_label() == expected
+
+
+def test_label_safe_in_backtick_identifier():
+    """The returned label cannot break out of a backtick-quoted identifier."""
+    storage = _make_storage("ws`} OR 1=1")
+    identifier = f"`{storage._get_workspace_label()}`"
+    # Every ` in the label was doubled to ``, so the quoting stays balanced.
+    assert identifier.count("`") % 2 == 0
+
+
+# ---------------------------------------------------------------------------
+# _get_raw_workspace_label: the actual label name, bound as a parameter
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "workspace, expected",
+    [
+        ("my-project", "my-project"),
+        ("  trimmed  ", "trimmed"),
+        ("", "base"),
+        ("   ", "base"),
+        # No escaping of any kind: the raw value is bound, never interpolated.
+        ("ws`name", "ws`name"),
+        ("team'a", "team'a"),
+        ("x'); CALL apoc.util.sleep(1);", "x'); CALL apoc.util.sleep(1);"),
+    ],
+)
+def test_get_raw_workspace_label(workspace, expected):
+    storage = _make_storage(workspace)
+    assert storage._get_raw_workspace_label() == expected
+
+
+# ---------------------------------------------------------------------------
+# get_knowledge_graph: labelFilter must be parameterized, never interpolated
+# ---------------------------------------------------------------------------
+
+
+class _FakeResult:
+    def __init__(self, record):
+        self._record = record
+
+    async def single(self):
+        return self._record
+
+    async def consume(self):
+        return None
+
+
+class _FakeSession:
+    """Records every (query, params) pair passed to ``run``."""
+
+    def __init__(self):
+        self.calls = []
+
+    async def run(self, query, params=None):
+        self.calls.append((query, params or {}))
+        # total_nodes within limit -> the full result is used directly and the
+        # truncated/limited branch is never taken (a single run call).
+        return _FakeResult({"total_nodes": 0, "node_info": [], "relationships": []})
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        return False
+
+
+class _FakeDriver:
+    def __init__(self, session):
+        self._session = session
+
+    def session(self, **kwargs):
+        return self._session
+
+
+@pytest.mark.asyncio
+async def test_label_filter_is_parameterized_not_interpolated():
+    """A malicious workspace must reach APOC labelFilter as a bound parameter.
+
+    If the label were string-interpolated (the old behaviour), a single quote
+    would close the string literal and inject Cypher. Binding it as a parameter
+    makes interpolation — and therefore injection — impossible by construction.
+    """
+    payload = "x'); CALL apoc.util.sleep(1);"
+    storage = _make_storage(payload)
+    session = _FakeSession()
+    storage._driver = _FakeDriver(session)
+    storage._DATABASE = None
+
+    await storage.get_knowledge_graph(node_label="some-entity", max_depth=2)
+
+    # The subgraphAll query is the one carrying labelFilter.
+    subgraph_calls = [(q, p) for q, p in session.calls if "apoc.path.subgraphAll" in q]
+    assert subgraph_calls, "expected an apoc.path.subgraphAll query"
+
+    for query, params in subgraph_calls:
+        # labelFilter is bound, not interpolated: there is no single-quoted
+        # string literal for the payload to break out of.
+        assert "labelFilter: $label_filter" in query
+        assert "labelFilter: '" not in query
+        # The raw (un-escaped) label is passed as the bound value.
+        assert params.get("label_filter") == payload

tests/kg/networkx_impl/test_networkx_atomic_write.py

@@ -0,0 +1,194 @@
+"""Regression tests for the atomic `write_nx_graph` path in NetworkXStorage.
+
+After the migration to the shared ``lightrag.file_atomic`` helper, two
+failure modes have distinct tmp-residue semantics:
+
+- Python-level exceptions (``write_fn`` raised, ``os.replace`` raised, ...):
+  ``atomic_write``'s ``finally`` removes the tmp best-effort and the prior
+  snapshot remains on disk.
+- Process-level kills (SIGKILL, OOM, hard reboot): ``finally`` does not run
+  and the tmp survives as an orphan; the startup reaper handles it.
+
+Covers concerns introduced when migrating from direct in-place write to
+temp-file + os.replace:
+
+1. Concurrent writers (multi-thread) all succeed and the final file parses.
+2. A crash between ``write_fn(tmp)`` and ``os.replace`` leaves the prior
+   snapshot intact on disk (atomicity guarantee, Python-exception path).
+3. The startup orphan-tmp reaper deletes tmp files older than the threshold
+   without touching tmp files that may belong to a live concurrent writer.
+   This is what salvages real SIGKILL/OOM residue.
+4. The rename preserves the destination file's existing permission bits
+   instead of widening them to the umask default.
+5. The orphan reaper handles glob metacharacters in file_name correctly
+   (workspace/namespace strings can legitimately contain '[', '*', '?').
+"""
+
+import os
+import stat
+import sys
+import threading
+import time
+from unittest.mock import patch
+
+import networkx as nx
+import pytest
+
+from lightrag.file_atomic import TMP_REAP_AGE_SECONDS, reap_orphan_tmp_files
+from lightrag.kg.networkx_impl import NetworkXStorage
+
+
+@pytest.mark.offline
+def test_write_nx_graph_concurrent_writers_no_race(tmp_path):
+    """5 threads racing on the same destination must all succeed; the final
+    file must parse, and no per-writer .tmp sibling may be left behind."""
+    dst = str(tmp_path / "graph_concurrent.graphml")
+    errors: list[BaseException] = []
+    barrier = threading.Barrier(5)
+
+    def writer(tid: int) -> None:
+        try:
+            g = nx.Graph()
+            g.add_node(f"node-{tid}", label=f"label-{tid}")
+            barrier.wait()
+            NetworkXStorage.write_nx_graph(g, dst, workspace=f"w{tid}")
+        except BaseException as exc:
+            errors.append(exc)
+
+    threads = [threading.Thread(target=writer, args=(i,)) for i in range(5)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert errors == [], f"Concurrent writers raised: {errors}"
+
+    # Final file must parse cleanly.
+    loaded = nx.read_graphml(dst)
+    assert loaded.number_of_nodes() == 1
+
+    # No per-writer tmp residue.
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"Unexpected tmp residue: {leftovers}"
+
+
+@pytest.mark.offline
+def test_write_nx_graph_crash_preserves_prior_snapshot(tmp_path):
+    """If ``os.replace`` raises after the tmp is written, the destination
+    must still hold the prior snapshot — never a torn/partial file. Because
+    the failure is a Python-level exception, ``atomic_write`` runs its
+    ``finally`` and removes the tmp; the startup-reaper path is exercised
+    separately in ``test_reap_orphan_tmp_files_respects_age_threshold``."""
+    dst = str(tmp_path / "graph_crash.graphml")
+
+    # Commit v1.
+    v1 = nx.Graph()
+    v1.add_node("v1-node")
+    NetworkXStorage.write_nx_graph(v1, dst, workspace="crash")
+    assert os.path.exists(dst)
+
+    # Attempt v2 with os.replace raising. Patch the symbol in file_atomic
+    # because the rename now lives in the shared helper, not in
+    # networkx_impl.
+    v2 = nx.Graph()
+    v2.add_node("v2-node")
+    with patch(
+        "lightrag.file_atomic.os.replace",
+        side_effect=OSError("simulated crash"),
+    ):
+        with pytest.raises(OSError, match="simulated crash"):
+            NetworkXStorage.write_nx_graph(v2, dst, workspace="crash")
+
+    # Destination must still be v1 — atomicity guarantee.
+    reloaded = nx.read_graphml(dst)
+    assert "v1-node" in reloaded.nodes
+    assert "v2-node" not in reloaded.nodes
+
+    # Python-exception path: tmp must have been cleaned up by atomic_write's
+    # finally. Real SIGKILL/OOM residue is the reaper's job, not this one.
+    leftovers = [p for p in os.listdir(tmp_path) if ".tmp." in p]
+    assert leftovers == [], f"Python-exception path must clean tmp, got {leftovers}"
+
+
+@pytest.mark.offline
+def test_reap_orphan_tmp_files_respects_age_threshold(tmp_path):
+    """`reap_orphan_tmp_files` must delete tmp siblings older than the
+    threshold and leave younger ones (potentially belonging to a live
+    concurrent writer) untouched."""
+    dst = str(tmp_path / "graph_reap.graphml")
+    old_tmp = f"{dst}.tmp.111.222.333"
+    young_tmp = f"{dst}.tmp.444.555.666"
+    unrelated = str(tmp_path / "graph_other.graphml.tmp.999")
+
+    for p in (old_tmp, young_tmp, unrelated):
+        with open(p, "w") as fh:
+            fh.write("partial xml")
+
+    # Age old_tmp past the threshold; leave young_tmp and unrelated fresh.
+    aged_mtime = time.time() - (TMP_REAP_AGE_SECONDS + 60)
+    os.utime(old_tmp, (aged_mtime, aged_mtime))
+
+    reap_orphan_tmp_files(dst, workspace="reap")
+
+    assert not os.path.exists(old_tmp), "Aged orphan should have been reaped"
+    assert os.path.exists(young_tmp), "Fresh tmp may be in-flight — must not be reaped"
+    assert os.path.exists(unrelated), (
+        "Unrelated path must not be touched by this reaper"
+    )
+
+
+@pytest.mark.offline
+@pytest.mark.skipif(sys.platform == "win32", reason="POSIX chmod semantics")
+def test_write_nx_graph_preserves_existing_file_mode(tmp_path):
+    """If the destination already exists with restrictive mode (e.g. 0600),
+    the atomic write must not silently widen it to umask defaults — the
+    inode swap done by os.replace would otherwise inherit the tmp file's
+    fresh permissions."""
+    dst = str(tmp_path / "graph_mode.graphml")
+
+    g = nx.Graph()
+    g.add_node("a")
+    NetworkXStorage.write_nx_graph(g, dst, workspace="mode")
+    os.chmod(dst, 0o600)
+    assert stat.S_IMODE(os.stat(dst).st_mode) == 0o600
+
+    # Second write — mode must survive.
+    g.add_node("b")
+    NetworkXStorage.write_nx_graph(g, dst, workspace="mode")
+    assert stat.S_IMODE(os.stat(dst).st_mode) == 0o600, (
+        "atomic write must preserve dst permissions across the rename"
+    )
+
+    # And content was actually updated.
+    assert nx.read_graphml(dst).number_of_nodes() == 2
+
+
+@pytest.mark.offline
+def test_reap_orphan_tmp_files_handles_glob_metacharacters(tmp_path):
+    """`file_name` is composed from `working_dir` and `namespace`, both of
+    which can legitimately contain glob metacharacters on POSIX. The reaper
+    must match literally — not (a) miss the real orphan because '[v2]' was
+    parsed as a character class, nor (b) widen its pattern to match
+    unrelated tmp files belonging to a sibling storage."""
+    dst = str(tmp_path / "graph_[v2].graphml")
+    real_orphan = f"{dst}.tmp.111.222.333"
+    decoy = str(tmp_path / "graph_v.graphml.tmp.unrelated")
+
+    for p in (real_orphan, decoy):
+        with open(p, "w") as fh:
+            fh.write("partial xml")
+
+    # Age both past the threshold so age is not the discriminator here —
+    # the only thing protecting the decoy is correct pattern escaping.
+    aged_mtime = time.time() - (TMP_REAP_AGE_SECONDS + 60)
+    for p in (real_orphan, decoy):
+        os.utime(p, (aged_mtime, aged_mtime))
+
+    reap_orphan_tmp_files(dst, workspace="meta")
+
+    assert not os.path.exists(real_orphan), (
+        "Reaper must match the real orphan even when path contains '['"
+    )
+    assert os.path.exists(decoy), (
+        "Reaper must not match tmp files belonging to unrelated paths"
+    )

tests/kg/networkx_impl/test_networkx_index_done.py

@@ -0,0 +1,65 @@
+"""NetworkXStorage.index_done_callback must RAISE on a graph-save failure
+(PR #3187), not swallow it and return False.
+
+A swallowed save error would let _insert_done's _flush_one treat the flush as
+successful (it only detects failures via exceptions), so the document would be
+marked PROCESSED with the graph changes unpersisted — silent data loss. This
+locks the raise-on-failure behavior that aligns NetworkX with the other
+backends (faiss/nano raise too).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from lightrag.kg.networkx_impl import NetworkXStorage
+from lightrag.kg.shared_storage import finalize_share_data, initialize_share_data
+from lightrag.utils import EmbeddingFunc
+
+pytestmark = pytest.mark.offline
+
+
+@pytest.fixture(autouse=True)
+def _shared_data():
+    finalize_share_data()
+    initialize_share_data()
+    yield
+    finalize_share_data()
+
+
+async def _embed(texts):
+    return np.random.rand(len(texts), 8)
+
+
+def _make_storage(tmp_path) -> NetworkXStorage:
+    return NetworkXStorage(
+        namespace="test_graph",
+        workspace="ws",
+        global_config={
+            "working_dir": str(tmp_path),
+            "embedding_batch_num": 10,
+            "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.5},
+        },
+        embedding_func=EmbeddingFunc(embedding_dim=8, max_token_size=512, func=_embed),
+    )
+
+
+@pytest.mark.asyncio
+async def test_index_done_callback_raises_on_save_failure(tmp_path, monkeypatch):
+    storage = _make_storage(tmp_path)
+    await storage.initialize()
+    try:
+        await storage.upsert_node("n1", {"entity_id": "n1", "description": "x"})
+
+        def boom(graph, file_name, workspace):
+            raise OSError("save boom")
+
+        # write_nx_graph is invoked as NetworkXStorage.write_nx_graph(...).
+        monkeypatch.setattr(NetworkXStorage, "write_nx_graph", staticmethod(boom))
+
+        # Must surface the error (NOT return False / swallow it).
+        with pytest.raises(OSError, match="save boom"):
+            await storage.index_done_callback()
+    finally:
+        await storage.finalize()

tests/kg/opensearch_impl/test_cwe89_opensearch_injection.py

@@ -0,0 +1,233 @@
+"""
+PoC test: CWE-89 OpenSearch injection via unsanitized entity names in query construction.
+
+The test validates that:
+1. Wildcard special characters (*, ?) in user input to search_labels are escaped
+   before being used in OpenSearch wildcard queries, preventing DoS via expensive
+   wildcard patterns.
+2. PPL escape handles control characters and additional metacharacters beyond
+   just backslash and single-quote.
+
+Run with: pytest tests/kg/opensearch_impl/test_cwe89_opensearch_injection.py -v
+"""
+
+import re
+import pytest
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock, patch
+import numpy as np
+
+pytest.importorskip(
+    "opensearchpy",
+    reason="opensearchpy is required for OpenSearch storage tests",
+)
+
+from lightrag.kg.opensearch_impl import (
+    OpenSearchGraphStorage,
+    ClientManager,
+)
+
+pytestmark = pytest.mark.offline
+
+
+@asynccontextmanager
+async def _mock_lock():
+    yield
+
+
+def _mock_lock_factory():
+    return _mock_lock()
+
+
+@pytest.fixture(autouse=True)
+def patch_data_init_lock():
+    with patch(
+        "lightrag.kg.opensearch_impl.get_data_init_lock", side_effect=_mock_lock_factory
+    ):
+        yield
+
+
+class MockEmbeddingFunc:
+    def __init__(self, dim=128):
+        self.embedding_dim = dim
+        self.max_token_size = 512
+        self.model_name = "mock-embed"
+
+    async def __call__(self, texts, **kwargs):
+        return np.random.rand(len(texts), self.embedding_dim).astype(np.float32)
+
+
+@pytest.fixture
+def global_config():
+    return {
+        "embedding_batch_num": 10,
+        "max_graph_nodes": 1000,
+    }
+
+
+@pytest.fixture
+def embed_func():
+    return MockEmbeddingFunc()
+
+
+def _make_client():
+    from opensearchpy import AsyncOpenSearch
+
+    client = AsyncMock(spec=AsyncOpenSearch)
+    # opensearchpy decorates client methods with @query_params, which hides their
+    # coroutine nature from inspect.iscoroutinefunction. AsyncMock(spec=...) then
+    # creates plain MagicMocks for them, so callers awaiting client.search(...) hit
+    # "object dict can't be used in 'await' expression". Force AsyncMock explicitly.
+    client.search = AsyncMock()
+    client.indices = AsyncMock()
+    client.indices.exists = AsyncMock(return_value=True)
+    client.indices.refresh = AsyncMock()
+    client.transport = AsyncMock()
+    return client
+
+
+@pytest.fixture
+def graph_storage(global_config, embed_func):
+    with patch.object(ClientManager, "get_client") as mock_get:
+        client = _make_client()
+        mock_get.return_value = client
+
+        storage = OpenSearchGraphStorage(
+            namespace="test_graph",
+            global_config=global_config,
+            embedding_func=embed_func,
+        )
+        storage.client = client
+        storage._indices_ready = True
+        storage._ppl_graphlookup_available = True
+        yield storage
+
+
+class TestWildcardInjection:
+    """Test that wildcard special chars are escaped in search_labels."""
+
+    @pytest.mark.asyncio
+    async def test_wildcard_chars_escaped_in_search_labels(self, graph_storage):
+        """Wildcard metacharacters *, ? in user input must be escaped."""
+        client = graph_storage.client
+
+        # Setup mock to return empty results
+        client.search.return_value = {"hits": {"hits": []}}
+
+        # Malicious query with wildcard chars that could cause expensive patterns
+        malicious_query = "test*?foo"
+        await graph_storage.search_labels(malicious_query)
+
+        # Inspect the query body that was sent to OpenSearch
+        assert client.search.called, "search should have been called"
+        call_kwargs = client.search.call_args
+        body = call_kwargs.kwargs.get("body") or call_kwargs[1].get("body")
+
+        # Extract the wildcard clause
+        should_clauses = body["query"]["bool"]["should"]
+        wildcard_clause = None
+        for clause in should_clauses:
+            if "wildcard" in clause:
+                wildcard_clause = clause["wildcard"]["entity_id"]["value"]
+                break
+
+        assert wildcard_clause is not None, "wildcard clause should exist"
+
+        # The wildcard value should NOT contain unescaped * or ? from the user input
+        # The outer * wrapping is fine (those are the intentional wildcards),
+        # but the inner user-provided * and ? must be escaped
+        # Expected: *test\*\?foo* (with the user's * and ? escaped)
+        inner_value = wildcard_clause[1:-1]  # strip leading and trailing *
+        assert "\\*" in inner_value, (
+            f"User's '*' should be escaped as '\\*' in wildcard, got: {wildcard_clause}"
+        )
+        assert "\\?" in inner_value, (
+            f"User's '?' should be escaped as '\\?' in wildcard, got: {wildcard_clause}"
+        )
+
+    @pytest.mark.asyncio
+    async def test_wildcard_heavy_pattern_not_exploitable(self, graph_storage):
+        """A series of ? chars should be escaped, not passed raw to OpenSearch."""
+        client = graph_storage.client
+        client.search.return_value = {"hits": {"hits": []}}
+
+        # Attack: many single-char wildcards cause exponential matching
+        attack_query = "?" * 50
+        await graph_storage.search_labels(attack_query)
+
+        call_kwargs = client.search.call_args
+        body = call_kwargs.kwargs.get("body") or call_kwargs[1].get("body")
+
+        should_clauses = body["query"]["bool"]["should"]
+        wildcard_clause = None
+        for clause in should_clauses:
+            if "wildcard" in clause:
+                wildcard_clause = clause["wildcard"]["entity_id"]["value"]
+                break
+
+        # None of the user's ? should appear as unescaped wildcards
+        # The value between the outer * delimiters should have all ? escaped
+        inner = wildcard_clause[1:-1]
+        # Count unescaped ? (i.e., ? not preceded by \)
+        unescaped_q = re.findall(r"(?<!\\)\?", inner)
+        assert len(unescaped_q) == 0, (
+            f"Found {len(unescaped_q)} unescaped '?' in wildcard pattern: {wildcard_clause}"
+        )
+
+    @pytest.mark.asyncio
+    async def test_backslash_escaped_in_wildcard(self, graph_storage):
+        """Backslashes in user input must be double-escaped for the wildcard query."""
+        client = graph_storage.client
+        client.search.return_value = {"hits": {"hits": []}}
+
+        attack_query = "test\\*"
+        await graph_storage.search_labels(attack_query)
+
+        call_kwargs = client.search.call_args
+        body = call_kwargs.kwargs.get("body") or call_kwargs[1].get("body")
+
+        should_clauses = body["query"]["bool"]["should"]
+        wildcard_clause = None
+        for clause in should_clauses:
+            if "wildcard" in clause:
+                wildcard_clause = clause["wildcard"]["entity_id"]["value"]
+                break
+
+        # The backslash should be escaped first, then the * — so we get \\\\\\*
+        # In the final pattern between outer *...*, user's \ becomes \\ and * becomes \*
+        inner = wildcard_clause[1:-1]
+        assert "\\\\" in inner or "\\*" in inner, (
+            f"Backslash and * from user should be escaped in wildcard: {wildcard_clause}"
+        )
+
+
+class TestPPLInjection:
+    """Test that PPL string escape handles additional metacharacters."""
+
+    def test_escape_ppl_basic_quote(self, graph_storage):
+        """Single quotes should be escaped."""
+        result = graph_storage._escape_ppl("it's a test")
+        assert "'" not in result.replace("\\'", ""), (
+            f"Unescaped quote found in: {result}"
+        )
+
+    def test_escape_ppl_backslash(self, graph_storage):
+        """Backslashes should be escaped."""
+        result = graph_storage._escape_ppl("test\\path")
+        assert result == "test\\\\path"
+
+    def test_escape_ppl_newline_and_control_chars(self, graph_storage):
+        """Newlines and control characters should be escaped/stripped."""
+        result = graph_storage._escape_ppl("line1\nline2\rline3\t")
+        # Control chars should either be stripped or escaped — no raw newlines
+        assert "\n" not in result, f"Raw newline in PPL literal: {repr(result)}"
+        assert "\r" not in result, f"Raw carriage return in PPL literal: {repr(result)}"
+
+    def test_escape_ppl_pipe_in_quotes_safe(self, graph_storage):
+        """Pipe character inside a quoted string literal poses no injection risk."""
+        result = graph_storage._escape_ppl("entity | stats count()")
+        assert isinstance(result, str)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/kg/opensearch_impl/test_llm_cache_tools_opensearch.py

@@ -0,0 +1,185 @@
+"""
+Offline tests for OpenSearch support in LLM cache tools.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+pytest.importorskip(
+    "opensearchpy",
+    reason="opensearchpy is required for OpenSearch tool tests",
+)
+
+from lightrag.tools.clean_llm_query_cache import CleanupStats, CleanupTool
+from lightrag.tools.migrate_llm_cache import MigrationTool
+
+pytestmark = pytest.mark.offline
+
+
+class FakeOpenSearchStorage:
+    def __init__(self, batches, workspace="test-workspace"):
+        self._batches = batches
+        self.workspace = workspace
+        self.deleted_batches = []
+
+    async def _iter_raw_docs(self, batch_size=1000):
+        for batch in self._batches:
+            yield batch
+
+    async def delete(self, ids):
+        self.deleted_batches.append(list(ids))
+
+
+def _flatten(batches):
+    return [item for batch in batches for item in batch]
+
+
+class TestCleanupToolOpenSearch:
+    @pytest.mark.asyncio
+    async def test_count_query_caches_opensearch(self):
+        tool = CleanupTool()
+        storage = FakeOpenSearchStorage(
+            [
+                [
+                    {"_id": "mix:query:1", "_source": {}},
+                    {"_id": "mix:keywords:1", "_source": {}},
+                    {"_id": "default:extract:1", "_source": {}},
+                ],
+                [
+                    {"_id": "hybrid:query:1", "_source": {}},
+                    {"_id": "local:keywords:1", "_source": {}},
+                    {"_id": "other:key:1", "_source": {}},
+                ],
+            ]
+        )
+
+        counts = await tool.count_query_caches(storage, "OpenSearchKVStorage")
+
+        assert counts["mix"] == {"query": 1, "keywords": 1}
+        assert counts["hybrid"] == {"query": 1, "keywords": 0}
+        assert counts["local"] == {"query": 0, "keywords": 1}
+        assert counts["global"] == {"query": 0, "keywords": 0}
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("cleanup_type", "expected_ids"),
+        [
+            (
+                "all",
+                [
+                    "mix:query:1",
+                    "mix:keywords:1",
+                    "global:query:1",
+                    "local:keywords:1",
+                ],
+            ),
+            ("query", ["mix:query:1", "global:query:1"]),
+            ("keywords", ["mix:keywords:1", "local:keywords:1"]),
+        ],
+    )
+    async def test_delete_query_caches_opensearch(self, cleanup_type, expected_ids):
+        tool = CleanupTool()
+        tool.batch_size = 2
+        storage = FakeOpenSearchStorage(
+            [
+                [
+                    {"_id": "mix:query:1", "_source": {}},
+                    {"_id": "mix:keywords:1", "_source": {}},
+                ],
+                [
+                    {"_id": "global:query:1", "_source": {}},
+                    {"_id": "local:keywords:1", "_source": {}},
+                    {"_id": "default:extract:1", "_source": {}},
+                ],
+            ]
+        )
+        stats = CleanupStats()
+
+        await tool.delete_query_caches(
+            storage, "OpenSearchKVStorage", cleanup_type, stats
+        )
+
+        assert _flatten(storage.deleted_batches) == expected_ids
+        assert all(len(batch) <= 2 for batch in storage.deleted_batches)
+        assert stats.successfully_deleted == len(expected_ids)
+        assert stats.successful_batches == len(storage.deleted_batches)
+
+    def test_check_config_ini_for_storage_opensearch(self, tmp_path, monkeypatch):
+        monkeypatch.chdir(tmp_path)
+        (tmp_path / "config.ini").write_text("[opensearch]\nhosts = localhost:9200\n")
+
+        assert CleanupTool().check_config_ini_for_storage("OpenSearchKVStorage")
+
+    def test_get_storage_class_opensearch(self):
+        cleanup_cls = CleanupTool().get_storage_class("OpenSearchKVStorage")
+        migrate_cls = MigrationTool().get_storage_class("OpenSearchKVStorage")
+
+        assert cleanup_cls.__name__ == "OpenSearchKVStorage"
+        assert migrate_cls.__name__ == "OpenSearchKVStorage"
+
+
+class TestMigrationToolOpenSearch:
+    @pytest.mark.asyncio
+    async def test_count_and_stream_default_caches_opensearch(self):
+        tool = MigrationTool()
+        storage = FakeOpenSearchStorage(
+            [
+                [
+                    {"_id": "default:extract:1", "_source": {"return": "a"}},
+                    {"_id": "mix:query:1", "_source": {"return": "ignored"}},
+                ],
+                [
+                    {"_id": "default:summary:1", "_source": {"return": "b"}},
+                    {"_id": "default:extract:2", "_source": {"return": "c"}},
+                ],
+            ]
+        )
+
+        count = await tool.count_default_caches(storage, "OpenSearchKVStorage")
+        streamed = [
+            batch
+            async for batch in tool.stream_default_caches(
+                storage, "OpenSearchKVStorage", batch_size=2
+            )
+        ]
+
+        assert count == 3
+        assert streamed == [
+            {
+                "default:extract:1": {"return": "a"},
+                "default:summary:1": {"return": "b"},
+            },
+            {"default:extract:2": {"return": "c"}},
+        ]
+
+    def test_count_available_storage_types_includes_opensearch(
+        self, tmp_path, monkeypatch
+    ):
+        monkeypatch.chdir(tmp_path)
+        (tmp_path / "config.ini").write_text("[opensearch]\nhosts = localhost:9200\n")
+
+        with patch.dict("os.environ", {}, clear=True):
+            assert MigrationTool().count_available_storage_types() == 2
+
+    @pytest.mark.asyncio
+    async def test_setup_storage_returns_effective_workspace(self, monkeypatch):
+        tool = MigrationTool()
+        fake_storage = SimpleNamespace(workspace="forced-workspace")
+
+        monkeypatch.setattr(tool, "check_env_vars", lambda _: True)
+        monkeypatch.setattr(
+            tool, "initialize_storage", AsyncMock(return_value=fake_storage)
+        )
+        monkeypatch.setattr(tool, "count_default_caches", AsyncMock(return_value=3))
+
+        with patch("builtins.input", return_value="5"):
+            storage, storage_name, workspace, total_count = await tool.setup_storage(
+                "Source", use_streaming=True
+            )
+
+        assert storage is fake_storage
+        assert storage_name == "OpenSearchKVStorage"
+        assert workspace == "forced-workspace"
+        assert total_count == 3

tests/kg/postgres_impl/test_doc_full_parse_engine_text.py

@@ -0,0 +1,91 @@
+"""LIGHTRAG_DOC_FULL.parse_engine must be TEXT (Phase 2).
+
+The parse_engine field may carry an encoded engine-parameter directive
+(``mineru(page_range=1-3,language=en)``) longer than the original VARCHAR(32),
+so the CREATE DDL uses TEXT and the migration widens an existing VARCHAR(32)
+column.  (A real overflow round-trip against live Postgres is the separate
+@pytest.mark.integration test; these mocked assertions cannot run real SQL.)
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock
+
+import pytest
+
+from lightrag.kg.postgres_impl import TABLES, PostgreSQLDB
+
+
+def test_create_ddl_uses_text_for_parse_engine():
+    ddl = TABLES["LIGHTRAG_DOC_FULL"]["ddl"]
+    assert "parse_engine TEXT" in ddl
+    assert "parse_engine VARCHAR(32)" not in ddl
+
+
+@pytest.mark.asyncio
+async def test_migration_widens_existing_varchar_column_to_text():
+    db = PostgreSQLDB.__new__(PostgreSQLDB)  # skip __init__/connection setup
+
+    executed: list[str] = []
+
+    async def fake_query(sql, *args, **kwargs):
+        # All columns already present -> ADD COLUMN loop is a no-op; the
+        # parse_engine column reports the legacy VARCHAR type so the widening
+        # ALTER fires.
+        if "information_schema.columns" in sql and "column_name = ANY" in sql:
+            return [
+                {"column_name": c}
+                for c in (
+                    "sidecar_location",
+                    "parse_format",
+                    "content_hash",
+                    "process_options",
+                    "chunk_options",
+                    "parse_engine",
+                )
+            ]
+        if "information_schema.columns" in sql:  # the parse_engine type probe
+            return {"data_type": "character varying"}
+        return None
+
+    async def fake_execute(sql, *args, **kwargs):
+        executed.append(sql)
+
+    db.query = AsyncMock(side_effect=fake_query)
+    db.execute = AsyncMock(side_effect=fake_execute)
+
+    await db._migrate_doc_full_add_pipeline_fields()
+
+    assert any("ALTER COLUMN parse_engine TYPE TEXT" in sql for sql in executed), (
+        executed
+    )
+
+
+@pytest.mark.asyncio
+async def test_migration_skips_alter_when_already_text():
+    db = PostgreSQLDB.__new__(PostgreSQLDB)
+    executed: list[str] = []
+
+    async def fake_query(sql, *args, **kwargs):
+        if "information_schema.columns" in sql and "column_name = ANY" in sql:
+            return [
+                {"column_name": c}
+                for c in (
+                    "sidecar_location",
+                    "parse_format",
+                    "content_hash",
+                    "process_options",
+                    "chunk_options",
+                    "parse_engine",
+                )
+            ]
+        if "information_schema.columns" in sql:
+            return {"data_type": "text"}
+        return None
+
+    db.query = AsyncMock(side_effect=fake_query)
+    db.execute = AsyncMock(side_effect=lambda sql, *a, **k: executed.append(sql))
+
+    await db._migrate_doc_full_add_pipeline_fields()
+
+    assert not any("ALTER COLUMN parse_engine" in sql for sql in executed)

tests/kg/postgres_impl/test_postgres_age_quote_fix.py

@@ -0,0 +1,252 @@
+"""
+Unit tests for PGGraphStorage.get_nodes_edges_batch and get_node_edges
+with special characters in entity names.
+
+Verifies the fix for KeyError when entity names contain double quotes (PR #2872)
+and the follow-up Option C refactor to parameterized Cypher queries.
+
+The root cause: AGE returns the original un-escaped entity_id, but the edges_norm
+dict was previously keyed with the normalized (escaped) ID, causing a KeyError on lookup.
+
+The Option C fix: use $node_ids / $entity_id parameters instead of string interpolation,
+eliminating the need for _normalize_node_id in these read paths entirely.
+"""
+
+import json
+import pytest
+from unittest.mock import MagicMock, patch
+
+from lightrag.kg.postgres_impl import PGGraphStorage
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def make_graph_storage() -> PGGraphStorage:
+    """Construct a PGGraphStorage instance with a mocked _query method."""
+    storage = PGGraphStorage.__new__(PGGraphStorage)
+    storage.workspace = "test_ws"
+    storage.namespace = "test_graph"
+    storage.graph_name = "test_graph"
+    storage.db = MagicMock()
+    return storage
+
+
+# ---------------------------------------------------------------------------
+# _normalize_node_id (still used by write paths: remove_nodes, upsert_node, etc.)
+# ---------------------------------------------------------------------------
+
+
+def test_normalize_plain_id():
+    assert PGGraphStorage._normalize_node_id("Alice") == "Alice"
+
+
+def test_normalize_double_quote():
+    assert PGGraphStorage._normalize_node_id('John "Smith"') == 'John \\"Smith\\"'
+
+
+def test_normalize_backslash():
+    assert PGGraphStorage._normalize_node_id("C:\\path") == "C:\\\\path"
+
+
+def test_normalize_both_special_chars():
+    assert (
+        PGGraphStorage._normalize_node_id('say \\"hello\\"')
+        == 'say \\\\\\"hello\\\\\\"'
+    )
+
+
+# ---------------------------------------------------------------------------
+# get_node_edges — parameterized query (Option C)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_get_node_edges_passes_original_id_as_parameter():
+    """entity_id must be passed as a JSON parameter, not interpolated into Cypher."""
+    storage = make_graph_storage()
+    entity = 'John "Smith"'
+    captured_params: list[dict] = []
+
+    async def fake_query(sql, **kwargs):
+        if kwargs.get("params"):
+            captured_params.append(json.loads(list(kwargs["params"].values())[0]))
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.get_node_edges(entity)
+
+    assert len(captured_params) == 1
+    assert captured_params[0]["entity_id"] == entity
+
+
+@pytest.mark.asyncio
+async def test_get_node_edges_cypher_uses_parameter_syntax():
+    """The SQL sent to _query must use $1::agtype, not a hardcoded escaped string."""
+    storage = make_graph_storage()
+    entity = 'John "Smith"'
+    captured_sql: list[str] = []
+
+    async def fake_query(sql, **kwargs):
+        captured_sql.append(sql)
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.get_node_edges(entity)
+
+    assert len(captured_sql) == 1
+    assert "$1::agtype" in captured_sql[0]
+    # Entity name must NOT appear literally in the SQL string
+    assert entity not in captured_sql[0]
+    assert '\\"' not in captured_sql[0]
+
+
+@pytest.mark.asyncio
+async def test_get_node_edges_returns_edges():
+    storage = make_graph_storage()
+
+    async def fake_query(_sql, **_kwargs):
+        return [
+            {"source_id": "Alice", "connected_id": "Bob"},
+            {"source_id": "Alice", "connected_id": None},
+        ]
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        result = await storage.get_node_edges("Alice")
+
+    assert result == [("Alice", "Bob")]
+
+
+# ---------------------------------------------------------------------------
+# get_nodes_edges_batch — parameterized query (Option C)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_get_nodes_edges_batch_passes_original_ids_as_parameter():
+    """node_ids batch must be passed as a JSON parameter, not interpolated."""
+    storage = make_graph_storage()
+    entities = ['John "Smith"', "Alice", "O\\Brien"]
+    captured_params: list[dict] = []
+
+    async def fake_query(_sql, **kwargs):
+        if kwargs.get("params"):
+            captured_params.append(json.loads(list(kwargs["params"].values())[0]))
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.get_nodes_edges_batch(entities)
+
+    assert len(captured_params) == 2  # outgoing + incoming
+    assert captured_params[0]["node_ids"] == entities
+    assert captured_params[1]["node_ids"] == entities
+
+
+@pytest.mark.asyncio
+async def test_get_nodes_edges_batch_cypher_uses_parameter_syntax():
+    """The SQL must use $1::agtype, not hardcoded escaped entity names."""
+    storage = make_graph_storage()
+    entity = 'John "Smith"'
+    captured_sql: list[str] = []
+
+    async def fake_query(sql, **_kwargs):
+        captured_sql.append(sql)
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.get_nodes_edges_batch([entity])
+
+    assert len(captured_sql) == 2
+    for sql in captured_sql:
+        assert "$1::agtype" in sql
+        assert entity not in sql
+        assert '\\"' not in sql
+
+
+@pytest.mark.asyncio
+async def test_get_nodes_edges_batch_with_quoted_entity():
+    """
+    AGE returns the original un-escaped node_id in query results.
+    The result dict must be keyed by the original ID.
+    """
+    storage = make_graph_storage()
+    entity = 'John "Smith"'
+
+    async def fake_query(sql, **_kwargs):
+        if "OPTIONAL MATCH (n:base)-[]->" in sql:
+            return [{"node_id": entity, "connected_id": "Alice"}]
+        if "OPTIONAL MATCH (n:base)<-[]-" in sql:
+            return [{"node_id": entity, "connected_id": "Bob"}]
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        result = await storage.get_nodes_edges_batch([entity])
+
+    assert entity in result
+    assert (entity, "Alice") in result[entity]
+    assert ("Bob", entity) in result[entity]
+
+
+@pytest.mark.asyncio
+async def test_get_nodes_edges_batch_plain_entity():
+    """Entity names without special chars still work correctly."""
+    storage = make_graph_storage()
+    entity = "Alice"
+
+    async def fake_query(sql, **_kwargs):
+        if "OPTIONAL MATCH (n:base)-[]->" in sql:
+            return [{"node_id": entity, "connected_id": "Bob"}]
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        result = await storage.get_nodes_edges_batch([entity])
+
+    assert entity in result
+    assert (entity, "Bob") in result[entity]
+
+
+@pytest.mark.asyncio
+async def test_get_nodes_edges_batch_no_results():
+    """Nodes with no edges return an empty list, not a KeyError."""
+    storage = make_graph_storage()
+    entity = 'Entity "X"'
+
+    async def fake_query(_sql, **_kwargs):
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        result = await storage.get_nodes_edges_batch([entity])
+
+    assert entity in result
+    assert result[entity] == []
+
+
+@pytest.mark.asyncio
+async def test_get_nodes_edges_batch_deduplication():
+    """Duplicate input IDs are deduplicated; each maps to the same edge list."""
+    storage = make_graph_storage()
+    entity = 'Dup "Entity"'
+
+    async def fake_query(sql, **_kwargs):
+        if "OPTIONAL MATCH (n:base)-[]->" in sql:
+            return [{"node_id": entity, "connected_id": "Other"}]
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        result = await storage.get_nodes_edges_batch([entity, entity])
+
+    assert result[entity] == [(entity, "Other")]
+
+
+@pytest.mark.asyncio
+async def test_get_nodes_edges_batch_empty_input():
+    """Empty input returns empty dict without calling _query."""
+    storage = make_graph_storage()
+
+    with patch.object(storage, "_query") as mock_q:
+        result = await storage.get_nodes_edges_batch([])
+
+    assert result == {}
+    mock_q.assert_not_called()

tests/kg/postgres_impl/test_postgres_batch_limits.py

@@ -0,0 +1,160 @@
+"""Unit tests for the PostgreSQL batch-limit helpers.
+
+Covers the module-level batching primitives shared by the PG non-graph write
+paths (mirrors the MONGO_* / OPENSEARCH_* contract):
+
+- ``_resolve_pg_batch_limits``: env parsing, defaults, non-positive disable.
+- ``_estimate_record_bytes``: per-field byte estimation for asyncpg tuples.
+- ``_chunk_by_budget``: record-count + payload-byte splitting.
+"""
+
+import json
+from unittest.mock import patch
+
+import numpy as np
+
+from lightrag.kg.postgres_impl import (
+    DEFAULT_PG_DELETE_MAX_RECORDS_PER_BATCH,
+    DEFAULT_PG_UPSERT_MAX_PAYLOAD_BYTES,
+    DEFAULT_PG_UPSERT_MAX_RECORDS_PER_BATCH,
+    _chunk_by_budget,
+    _estimate_record_bytes,
+    _resolve_pg_batch_limits,
+)
+
+
+# ---------------------------------------------------------------------------
+# _resolve_pg_batch_limits
+# ---------------------------------------------------------------------------
+
+
+class TestResolvePgBatchLimits:
+    def test_defaults(self):
+        with patch.dict("os.environ", {}, clear=True):
+            payload, upserts, deletes = _resolve_pg_batch_limits()
+        assert payload == DEFAULT_PG_UPSERT_MAX_PAYLOAD_BYTES
+        assert upserts == DEFAULT_PG_UPSERT_MAX_RECORDS_PER_BATCH
+        assert deletes == DEFAULT_PG_DELETE_MAX_RECORDS_PER_BATCH
+
+    def test_default_record_cap_keeps_kv_historical_200(self):
+        # The PG-wide upsert record default deliberately stays at 200 (KV's
+        # historical batch size), not Mongo/OS's 128.
+        assert DEFAULT_PG_UPSERT_MAX_RECORDS_PER_BATCH == 200
+
+    def test_env_override(self):
+        env = {
+            "POSTGRES_UPSERT_MAX_PAYLOAD_BYTES": "12345",
+            "POSTGRES_UPSERT_MAX_RECORDS_PER_BATCH": "7",
+            "POSTGRES_DELETE_MAX_RECORDS_PER_BATCH": "9",
+        }
+        with patch.dict("os.environ", env, clear=True):
+            assert _resolve_pg_batch_limits() == (12345, 7, 9)
+
+    def test_non_positive_disables_and_warns(self):
+        env = {
+            "POSTGRES_UPSERT_MAX_PAYLOAD_BYTES": "0",
+            "POSTGRES_UPSERT_MAX_RECORDS_PER_BATCH": "-1",
+            "POSTGRES_DELETE_MAX_RECORDS_PER_BATCH": "0",
+        }
+        with patch.dict("os.environ", env, clear=True):
+            with patch("lightrag.kg.postgres_impl.logger") as mock_logger:
+                payload, upserts, deletes = _resolve_pg_batch_limits()
+        assert (payload, upserts, deletes) == (0, -1, 0)
+        warnings = [c.args[0] for c in mock_logger.warning.call_args_list]
+        assert sum("non-positive" in msg for msg in warnings) == 3
+
+
+# ---------------------------------------------------------------------------
+# _estimate_record_bytes
+# ---------------------------------------------------------------------------
+
+
+class TestEstimateRecordBytes:
+    def test_ndarray_uses_nbytes(self):
+        vec = np.zeros(128, dtype=np.float32)
+        # 128 * 4 bytes (nbytes) plus the 2-byte "id" str.
+        assert _estimate_record_bytes(("id", vec)) == vec.nbytes + 2
+
+    def test_str_uses_utf8_length(self):
+        assert _estimate_record_bytes(("abc",)) == 3
+
+    def test_multibyte_str(self):
+        # Each CJK char is 3 bytes in UTF-8.
+        assert _estimate_record_bytes(("中文",)) == 6
+
+    def test_bytes_and_none(self):
+        assert _estimate_record_bytes((b"abcd", None)) == 4
+
+    def test_dict_and_list_use_compact_json(self):
+        payload = {"a": 1, "b": [1, 2, 3]}
+        expected = len(
+            json.dumps(
+                payload, ensure_ascii=False, separators=(",", ":"), default=str
+            ).encode("utf-8")
+        )
+        assert _estimate_record_bytes((payload,)) == expected
+
+    def test_scalar_constant(self):
+        # int + float + None -> 16 + 16 + 0
+        assert _estimate_record_bytes((1, 2.0, None)) == 32
+
+    def test_mixed_tuple(self):
+        vec = np.ones(4, dtype=np.float32)  # 16 bytes
+        record = ("ws", "id", "hello", vec, None, 42)
+        # 2 + 2 + 5 + 16 + 0 + 16
+        assert _estimate_record_bytes(record) == 41
+
+
+# ---------------------------------------------------------------------------
+# _chunk_by_budget
+# ---------------------------------------------------------------------------
+
+
+class TestChunkByBudget:
+    def test_empty(self):
+        assert _chunk_by_budget([], lambda x: 1, 100, 100) == []
+
+    def test_split_by_record_count(self):
+        items = list(range(7))
+        batches = _chunk_by_budget(
+            items, lambda x: 1, max_payload_bytes=0, max_records_per_batch=3
+        )
+        sizes = [len(b) for b, _ in batches]
+        assert sizes == [3, 3, 1]
+
+    def test_split_by_payload_bytes(self):
+        # Each item ~10 bytes; a 25-byte budget fits 2 per batch (2 overhead +
+        # 10 + 1 + 10 = 23 <= 25; adding a third would exceed).
+        items = ["x" * 10, "y" * 10, "z" * 10, "w" * 10]
+        batches = _chunk_by_budget(
+            items, lambda s: len(s), max_payload_bytes=25, max_records_per_batch=0
+        )
+        sizes = [len(b) for b, _ in batches]
+        assert sizes == [2, 2]
+
+    def test_both_dimensions_record_cap_wins(self):
+        items = ["x" * 5] * 10
+        # record cap 4 binds before the generous byte budget.
+        batches = _chunk_by_budget(
+            items, lambda s: len(s), max_payload_bytes=10_000, max_records_per_batch=4
+        )
+        sizes = [len(b) for b, _ in batches]
+        assert sizes == [4, 4, 2]
+
+    def test_oversized_single_item_becomes_own_batch(self):
+        items = ["small", "X" * 1000, "small2"]
+        batches = _chunk_by_budget(
+            items, lambda s: len(s), max_payload_bytes=50, max_records_per_batch=0
+        )
+        # The 1000-byte item is emitted alone rather than raising.
+        sizes = [len(b) for b, _ in batches]
+        assert sizes == [1, 1, 1]
+        assert batches[1][0] == ["X" * 1000]
+
+    def test_non_positive_disables_both_dimensions(self):
+        items = list(range(100))
+        batches = _chunk_by_budget(
+            items, lambda x: 999, max_payload_bytes=0, max_records_per_batch=0
+        )
+        assert len(batches) == 1
+        assert len(batches[0][0]) == 100

tests/kg/postgres_impl/test_postgres_client_manager.py

@@ -0,0 +1,73 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from lightrag.kg.postgres_impl import ClientManager
+
+
+@pytest.fixture(autouse=True)
+def reset_client_manager_state() -> None:
+    ClientManager._instances = {"db": None, "ref_count": 0, "vector_signature": None}
+
+
+def test_pg_vector_storage_enables_vector() -> None:
+    config = ClientManager.get_config("PGVectorStorage")
+    assert config["enable_vector"] is True
+
+
+def test_non_pg_vector_storage_disables_vector() -> None:
+    config = ClientManager.get_config("NanoVectorDBStorage")
+    assert config["enable_vector"] is False
+
+
+def test_milvus_storage_disables_vector() -> None:
+    config = ClientManager.get_config("MilvusVectorDBStorage")
+    assert config["enable_vector"] is False
+
+
+def test_qdrant_storage_disables_vector() -> None:
+    config = ClientManager.get_config("QdrantVectorDBStorage")
+    assert config["enable_vector"] is False
+
+
+def test_none_vector_storage_defaults_to_true() -> None:
+    # Backward compatibility: when vector_storage is unknown (None), default to True.
+    config = ClientManager.get_config(None)
+    assert config["enable_vector"] is True
+
+
+def test_no_args_defaults_to_true() -> None:
+    # Backward compatibility: calling without arguments preserves prior behavior.
+    config = ClientManager.get_config()
+    assert config["enable_vector"] is True
+
+
+@pytest.mark.asyncio
+async def test_get_client_reuses_shared_pool_for_same_vector_settings() -> None:
+    db = MagicMock()
+    db.initdb = AsyncMock()
+    db.check_tables = AsyncMock()
+
+    with patch("lightrag.kg.postgres_impl.PostgreSQLDB", return_value=db) as db_cls:
+        first = await ClientManager.get_client("PGVectorStorage")
+        second = await ClientManager.get_client("PGVectorStorage")
+
+    assert first is db
+    assert second is db
+    assert ClientManager._instances["ref_count"] == 2
+    db_cls.assert_called_once()
+    db.initdb.assert_awaited_once()
+    db.check_tables.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_get_client_rejects_conflicting_vector_storage_settings() -> None:
+    db = MagicMock()
+    db.initdb = AsyncMock()
+    db.check_tables = AsyncMock()
+
+    with patch("lightrag.kg.postgres_impl.PostgreSQLDB", return_value=db):
+        await ClientManager.get_client("NanoVectorDBStorage")
+
+        with pytest.raises(RuntimeError, match="process-wide"):
+            await ClientManager.get_client("PGVectorStorage")

tests/kg/postgres_impl/test_postgres_cypher_injection.py

@@ -0,0 +1,362 @@
+"""
+Unit tests for Cypher injection prevention in PGGraphStorage write paths.
+
+Verifies that upsert_node and upsert_edge keep entity IDs parameterized while
+rendering property maps as safely escaped Cypher literals, which is required by
+Apache AGE because ``SET ... += $props`` is not supported.
+"""
+
+import json
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from lightrag.kg.postgres_impl import PGGraphStorage
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def make_graph_storage() -> PGGraphStorage:
+    """Construct a PGGraphStorage instance with a mocked db."""
+    storage = PGGraphStorage.__new__(PGGraphStorage)
+    storage.workspace = "test_ws"
+    storage.namespace = "test_graph"
+    storage.graph_name = "test_graph"
+    storage.db = MagicMock()
+    return storage
+
+
+class _FakeConnection:
+    """Captures statements + args passed to a fake asyncpg connection."""
+
+    def __init__(self):
+        self.calls: list[dict] = []
+
+    def transaction(self):
+        return _FakeTransaction()
+
+    async def execute(self, sql, *args):
+        self.calls.append({"sql": sql, "args": args})
+        return ""
+
+
+class _FakeTransaction:
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+
+async def _capture_upsert_edge(storage: PGGraphStorage, src: str, tgt: str, edge_data):
+    """Invoke upsert_edge against a fake connection and return the captured calls."""
+    conn = _FakeConnection()
+
+    async def fake_run_with_retry(operation, **_kwargs):
+        return await operation(conn)
+
+    storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
+    await storage.upsert_edge(src, tgt, edge_data)
+    return conn.calls
+
+
+# ---------------------------------------------------------------------------
+# upsert_node — parameterized Cypher
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_upsert_node_uses_parameterized_cypher():
+    """upsert_node must pass entity_id as a Cypher parameter, not interpolate it."""
+    storage = make_graph_storage()
+    captured_calls: list[dict] = []
+
+    async def fake_query(sql, **kwargs):
+        captured_calls.append({"sql": sql, **kwargs})
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.upsert_node(
+            "Alice", {"entity_id": "Alice", "description": "A person"}
+        )
+
+    assert len(captured_calls) == 1
+    call = captured_calls[0]
+    assert "$1::agtype" in call["sql"]
+    assert '"Alice"' not in call["sql"].replace("$1::agtype", "")
+    assert "params" in call
+    params = json.loads(call["params"]["params"])
+    assert params["entity_id"] == "Alice"
+    assert "props" not in params
+    assert '`description`: "A person"' in call["sql"]
+
+
+@pytest.mark.asyncio
+async def test_upsert_node_injection_payload_in_entity_id():
+    """A Cypher injection payload in entity_id must be treated as data, not code."""
+    storage = make_graph_storage()
+    injection = 'test"}) RETURN n; MATCH (m) DETACH DELETE m; //'
+    captured_calls: list[dict] = []
+
+    async def fake_query(sql, **kwargs):
+        captured_calls.append({"sql": sql, **kwargs})
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.upsert_node(
+            injection, {"entity_id": injection, "description": "malicious"}
+        )
+
+    call = captured_calls[0]
+    # The injection payload must NOT appear in the SQL string
+    assert "DETACH DELETE" not in call["sql"]
+    assert injection not in call["sql"]
+    # It must be safely contained in the JSON parameter
+    params = json.loads(call["params"]["params"])
+    assert params["entity_id"] == injection
+
+
+@pytest.mark.asyncio
+async def test_upsert_node_special_chars_in_properties():
+    """Property values with special characters are safely escaped in Cypher."""
+    storage = make_graph_storage()
+    captured_calls: list[dict] = []
+
+    async def fake_query(sql, **kwargs):
+        captured_calls.append({"sql": sql, **kwargs})
+        return []
+
+    node_data = {
+        "entity_id": "test_node",
+        "description": 'He said "hello" and used a backslash \\',
+        "notes": "Line1\nLine2\tTabbed",
+        "formula": "x < 5 && y > 3",
+    }
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.upsert_node("test_node", node_data)
+
+    call = captured_calls[0]
+    assert (
+        '`description`: "He said \\"hello\\" and used a backslash \\\\"' in call["sql"]
+    )
+    assert '`notes`: "Line1\\nLine2\\tTabbed"' in call["sql"]
+    assert '`formula`: "x < 5 && y > 3"' in call["sql"]
+
+
+@pytest.mark.asyncio
+async def test_upsert_node_unicode_entity_id():
+    """Unicode entity names are safely parameterized."""
+    storage = make_graph_storage()
+    captured_calls: list[dict] = []
+
+    async def fake_query(sql, **kwargs):
+        captured_calls.append({"sql": sql, **kwargs})
+        return []
+
+    unicode_id = "\u4e2d\u6587\u5b9e\u4f53"  # Chinese characters
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.upsert_node(
+            unicode_id, {"entity_id": unicode_id, "description": "\u63cf\u8ff0"}
+        )
+
+    call = captured_calls[0]
+    params = json.loads(call["params"]["params"])
+    assert params["entity_id"] == unicode_id
+    assert '`description`: "描述"' in call["sql"]
+
+
+@pytest.mark.asyncio
+async def test_upsert_node_dollar_signs_in_entity_id():
+    """Dollar signs in entity_id don't break dollar-quoting of the Cypher template."""
+    storage = make_graph_storage()
+    captured_calls: list[dict] = []
+
+    async def fake_query(sql, **kwargs):
+        captured_calls.append({"sql": sql, **kwargs})
+        return []
+
+    dollar_id = "price is $100 or $$200$$"
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.upsert_node(
+            dollar_id, {"entity_id": dollar_id, "description": "has dollars"}
+        )
+
+    call = captured_calls[0]
+    # The dollar signs are in the params, not the SQL template
+    params = json.loads(call["params"]["params"])
+    assert params["entity_id"] == dollar_id
+
+
+@pytest.mark.asyncio
+async def test_upsert_node_escapes_backticks_in_property_keys():
+    """Backticks in property keys must be escaped before inlining the map."""
+    storage = make_graph_storage()
+    captured_calls: list[dict] = []
+
+    async def fake_query(sql, **kwargs):
+        captured_calls.append({"sql": sql, **kwargs})
+        return []
+
+    with patch.object(storage, "_query", side_effect=fake_query):
+        await storage.upsert_node(
+            "node",
+            {"entity_id": "node", "danger`key": 'value "quoted"'},
+        )
+
+    assert '`danger``key`: "value \\"quoted\\""' in captured_calls[0]["sql"]
+
+
+@pytest.mark.asyncio
+async def test_upsert_node_requires_entity_id():
+    """upsert_node still raises ValueError when entity_id is missing."""
+    storage = make_graph_storage()
+    with pytest.raises(ValueError, match="entity_id"):
+        await storage.upsert_node("test", {"description": "no entity_id"})
+
+
+# ---------------------------------------------------------------------------
+# upsert_edge — parameterized Cypher
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_upsert_edge_uses_parameterized_cypher():
+    """upsert_edge must pass entity IDs as Cypher parameters."""
+    storage = make_graph_storage()
+    calls = await _capture_upsert_edge(
+        storage, "Alice", "Bob", {"weight": "1.0", "description": "knows"}
+    )
+
+    # Three statements: per-edge lock, graph-wide shared lock, then cypher.
+    assert len(calls) == 3
+
+    lock_sql = calls[0]["sql"]
+    # Raw node IDs are positional params on the lock, never interpolated.
+    assert "Alice" not in lock_sql
+    assert "Bob" not in lock_sql
+    # graph_name flows as $1, the endpoint pair as $2/$3.
+    assert calls[0]["args"] == ("test_graph", "Alice", "Bob")
+
+    cypher_call = calls[2]
+    cypher_sql = cypher_call["sql"]
+    assert "$1::agtype" in cypher_sql
+    assert '"Alice"' not in cypher_sql.replace("$1::agtype", "")
+    assert '"Bob"' not in cypher_sql.replace("$1::agtype", "")
+    # Cypher params arrive as a single positional agtype JSON arg.
+    params = json.loads(cypher_call["args"][0])
+    assert params["src_id"] == "Alice"
+    assert params["tgt_id"] == "Bob"
+    assert "props" not in params
+    assert '`weight`: "1.0"' in cypher_sql
+    assert '`description`: "knows"' in cypher_sql
+
+
+@pytest.mark.asyncio
+async def test_upsert_edge_injection_payload():
+    """Injection payloads in edge entity IDs are safely parameterized."""
+    storage = make_graph_storage()
+    injection_src = 'src"}) MATCH (x) DETACH DELETE x; //'
+    injection_tgt = 'tgt"})-[r]-() DELETE r; //'
+    calls = await _capture_upsert_edge(
+        storage, injection_src, injection_tgt, {"description": "edge"}
+    )
+
+    # Injection payloads must never appear in either SQL template — they only
+    # flow through positional params.
+    for call in calls:
+        assert "DETACH DELETE" not in call["sql"]
+        assert "DELETE r" not in call["sql"]
+        assert injection_src not in call["sql"]
+        assert injection_tgt not in call["sql"]
+
+    # Lock statement passes graph_name + raw IDs as positional params.
+    assert calls[0]["args"] == ("test_graph", injection_src, injection_tgt)
+
+    # Cypher params arrive as a single positional agtype JSON arg (3rd statement,
+    # after the per-edge and graph-wide-shared locks).
+    params = json.loads(calls[2]["args"][0])
+    assert params["src_id"] == injection_src
+    assert params["tgt_id"] == injection_tgt
+
+
+@pytest.mark.asyncio
+async def test_upsert_edge_unicode_entity_ids():
+    """Unicode entity IDs in edges are safely parameterized."""
+    storage = make_graph_storage()
+    src = "\u5317\u4eac"
+    tgt = "\u4e0a\u6d77"
+    calls = await _capture_upsert_edge(
+        storage, src, tgt, {"description": "\u8def\u7ebf"}
+    )
+
+    # Lock statement carries graph_name + raw IDs as positional params, not
+    # interpolated.
+    assert calls[0]["args"] == ("test_graph", src, tgt)
+    assert src not in calls[0]["sql"]
+    assert tgt not in calls[0]["sql"]
+
+    # Cypher params parsed from the positional agtype JSON arg (3rd statement).
+    cypher_sql = calls[2]["sql"]
+    params = json.loads(calls[2]["args"][0])
+    assert params["src_id"] == src
+    assert params["tgt_id"] == tgt
+    assert '`description`: "路线"' in cypher_sql
+
+
+# ---------------------------------------------------------------------------
+# _normalize_node_id — defence-in-depth for remaining interpolation paths
+# ---------------------------------------------------------------------------
+
+
+def test_normalize_node_id_strips_null_bytes():
+    """Null bytes are stripped to prevent string truncation."""
+    assert PGGraphStorage._normalize_node_id("before\x00after") == "beforeafter"
+
+
+def test_normalize_node_id_escapes_backslash_and_quote():
+    """Backslashes and double quotes are escaped."""
+    assert PGGraphStorage._normalize_node_id('a\\"b') == 'a\\\\\\"b'
+
+
+def test_normalize_node_id_injection_payload():
+    """Injection payload is escaped so it cannot break out of Cypher string."""
+    payload = 'test"}) RETURN n; MATCH (m) DETACH DELETE m; //'
+    normalized = PGGraphStorage._normalize_node_id(payload)
+    # The double quote must be escaped
+    assert '\\"' in normalized
+    # The escaped string must not contain an unescaped double quote
+    # (remove all escaped quotes and check no raw ones remain)
+    unescaped = normalized.replace('\\"', "")
+    assert '"' not in unescaped
+
+
+# ---------------------------------------------------------------------------
+# _query write path passes params to db.execute
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_query_write_path_passes_params():
+    """When readonly=False, _query must forward params to db.execute."""
+    storage = make_graph_storage()
+    captured_execute_kwargs: list[dict] = []
+
+    async def fake_execute(sql, **kwargs):
+        captured_execute_kwargs.append(kwargs)
+        return None
+
+    storage.db.execute = fake_execute
+
+    test_params = {"params": json.dumps({"entity_id": "test"})}
+    await storage._query(
+        "SELECT 1",
+        readonly=False,
+        upsert=True,
+        params=test_params,
+    )
+
+    assert len(captured_execute_kwargs) == 1
+    assert captured_execute_kwargs[0]["data"] == test_params

tests/kg/postgres_impl/test_postgres_delete_set_ids.py

@@ -0,0 +1,98 @@
+"""Regression test for Issue #3286.
+
+``adelete_by_doc_id`` / ``_purge_doc_chunks_and_kg`` historically passed
+``chunk_ids`` as a ``set`` to ``PGKVStorage.delete`` / ``PGDocStatusStorage.delete``.
+After delete-batching was introduced, the chunked path slices the id collection
+(``ids[i : i + chunk]``); slicing a ``set`` raises ``'set' object is not
+subscriptable``, which the broad ``except Exception`` in ``delete`` then swallows
+as a log line -- so the records were silently NOT deleted.
+
+These tests feed a ``set`` straight into ``delete`` with a positive batch cap
+(forcing the chunked path) and assert that every id is actually handed to the
+SQL layer. Asserting "no exception" alone is insufficient because the bug is
+swallowed; we must assert the records really got deleted.
+"""
+
+import pytest
+
+pytest.importorskip("asyncpg")
+
+from lightrag.kg.postgres_impl import (  # noqa: E402
+    PGDocStatusStorage,
+    PGKVStorage,
+)
+from lightrag.namespace import NameSpace  # noqa: E402
+
+
+class _FakeTransaction:
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        return False
+
+
+class _RecordingConnection:
+    """Captures the id-array argument of each ``execute`` call."""
+
+    def __init__(self):
+        self.deleted_id_batches: list[list[str]] = []
+
+    def transaction(self):
+        return _FakeTransaction()
+
+    async def execute(self, sql, workspace, id_slice):
+        # The chunked path slices ids; a slice of a list is a list. Record it.
+        assert isinstance(id_slice, list)
+        self.deleted_id_batches.append(id_slice)
+
+
+class _FakeDB:
+    def __init__(self, connection):
+        self._connection = connection
+
+    async def _run_with_retry(self, operation):
+        # Mirror the real helper: invoke the closure against a live connection.
+        return await operation(self._connection)
+
+
+def _make_storage(cls, namespace, batch_cap):
+    """Build a storage instance exercising only the attributes ``delete`` reads."""
+    storage = object.__new__(cls)
+    storage.namespace = namespace
+    storage.workspace = "test_ws"
+    storage._max_delete_records_per_batch = batch_cap
+    connection = _RecordingConnection()
+    storage.db = _FakeDB(connection)
+    return storage, connection
+
+
+@pytest.mark.asyncio
+async def test_pgkv_delete_accepts_set_and_chunks_all_ids():
+    ids = {f"chunk-{i}" for i in range(5)}
+    # Positive cap < len(ids) forces the chunked slicing path that broke on sets.
+    storage, connection = _make_storage(
+        PGKVStorage, NameSpace.KV_STORE_TEXT_CHUNKS, batch_cap=2
+    )
+
+    await storage.delete(ids)
+
+    flattened = [cid for batch in connection.deleted_id_batches for cid in batch]
+    assert set(flattened) == ids
+    assert len(flattened) == len(ids)  # no duplicates / drops
+    assert all(len(batch) <= 2 for batch in connection.deleted_id_batches)
+
+
+@pytest.mark.asyncio
+async def test_pgdocstatus_delete_accepts_set_and_chunks_all_ids():
+    ids = {f"doc-{i}" for i in range(5)}
+    storage, connection = _make_storage(
+        PGDocStatusStorage, NameSpace.DOC_STATUS, batch_cap=2
+    )
+
+    await storage.delete(ids)
+
+    flattened = [did for batch in connection.deleted_id_batches for did in batch]
+    assert set(flattened) == ids
+    assert len(flattened) == len(ids)
+    assert all(len(batch) <= 2 for batch in connection.deleted_id_batches)

tests/kg/postgres_impl/test_postgres_doc_status_lookup.py

@@ -0,0 +1,166 @@
+"""Unit tests for PGDocStatusStorage database-native overrides.
+
+Covers the PG-specific implementations of:
+  * get_doc_by_file_basename
+  * get_doc_by_content_hash
+
+Both override the base-class full-table scan with indexed SQL queries.
+"""
+
+from datetime import datetime
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+
+from lightrag.kg.postgres_impl import PGDocStatusStorage
+from lightrag.namespace import NameSpace
+
+
+def _make_storage():
+    storage = PGDocStatusStorage.__new__(PGDocStatusStorage)
+    storage.namespace = NameSpace.DOC_STATUS
+    storage.workspace = "test_ws"
+    storage.global_config = {"embedding_batch_num": 10}
+    storage.db = MagicMock()
+    storage.db.query = AsyncMock()
+    return storage
+
+
+def _row(**overrides):
+    base = {
+        "id": "doc-1",
+        "content_summary": "summary",
+        "content_length": 12,
+        "chunks_count": 1,
+        "status": "processed",
+        "file_path": "report.pdf",
+        "chunks_list": "[]",
+        "metadata": "{}",
+        "error_msg": None,
+        "track_id": None,
+        "content_hash": "abc123",
+        "created_at": datetime(2024, 1, 1, 0, 0, 0),
+        "updated_at": datetime(2024, 1, 1, 0, 0, 0),
+    }
+    base.update(overrides)
+    return base
+
+
+# ---------------------------------------------------------------------------
+# get_doc_by_file_basename
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_file_basename_empty_returns_none():
+    storage = _make_storage()
+    assert await storage.get_doc_by_file_basename("") is None
+    storage.db.query.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_file_basename_unknown_source_returns_none():
+    storage = _make_storage()
+    # normalize_document_file_path returns "unknown_source" for None-ish inputs
+    assert await storage.get_doc_by_file_basename("unknown_source") is None
+    storage.db.query.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_file_basename_exact_match():
+    storage = _make_storage()
+    storage.db.query.return_value = [_row(file_path="report.pdf")]
+
+    result = await storage.get_doc_by_file_basename("report.pdf")
+
+    assert result is not None
+    doc_id, doc = result
+    assert doc_id == "doc-1"
+    assert doc["file_path"] == "report.pdf"
+    assert doc["content_hash"] == "abc123"
+
+    call = storage.db.query.call_args
+    sql = call.args[0]
+    params = call.args[1]
+    assert "LIGHTRAG_DOC_STATUS" in sql
+    assert "workspace=$1" in sql
+    assert params[0] == "test_ws"
+    assert params[1] == "report.pdf"
+    assert params == ["test_ws", "report.pdf"]
+    assert "LIKE" not in sql
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_file_basename_orders_stably_for_canonical_rows():
+    storage = _make_storage()
+    storage.db.query.return_value = [_row(id="doc-exact", file_path="report.pdf")]
+
+    result = await storage.get_doc_by_file_basename("report.pdf")
+
+    assert result is not None
+    assert result[0] == "doc-exact"
+
+    sql = storage.db.query.call_args.args[0]
+    assert "file_path = $2" in sql
+    assert "created_at ASC" in sql
+    assert "id ASC" in sql
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_file_basename_uses_exact_match_for_like_metacharacters():
+    storage = _make_storage()
+    storage.db.query.return_value = []
+
+    await storage.get_doc_by_file_basename("100%_off.pdf")
+
+    sql = storage.db.query.call_args.args[0]
+    params = storage.db.query.call_args.args[1]
+    assert "LIKE" not in sql
+    assert params == ["test_ws", "100%_off.pdf"]
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_file_basename_no_match_returns_none():
+    storage = _make_storage()
+    storage.db.query.return_value = []
+    assert await storage.get_doc_by_file_basename("missing.pdf") is None
+
+
+# ---------------------------------------------------------------------------
+# get_doc_by_content_hash
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_content_hash_empty_returns_none():
+    storage = _make_storage()
+    assert await storage.get_doc_by_content_hash("") is None
+    storage.db.query.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_content_hash_match():
+    storage = _make_storage()
+    storage.db.query.return_value = [_row(content_hash="hash-abc")]
+
+    result = await storage.get_doc_by_content_hash("hash-abc")
+
+    assert result is not None
+    doc_id, doc = result
+    assert doc_id == "doc-1"
+    assert doc["content_hash"] == "hash-abc"
+
+    call = storage.db.query.call_args
+    sql = call.args[0]
+    params = call.args[1]
+    assert "content_hash=$2" in sql
+    # Stable ordering for repeatability across re-runs / replicas
+    assert "ORDER BY created_at ASC, id ASC" in sql
+    assert "LIMIT 1" in sql
+    assert params == ["test_ws", "hash-abc"]
+
+
+@pytest.mark.asyncio
+async def test_get_doc_by_content_hash_no_match_returns_none():
+    storage = _make_storage()
+    storage.db.query.return_value = []
+    assert await storage.get_doc_by_content_hash("nope") is None

tests/kg/postgres_impl/test_postgres_get_knowledge_graph_all.py

@@ -0,0 +1,189 @@
+"""
+Unit tests for PGGraphStorage.get_knowledge_graph("*") — the full-graph view.
+
+These cover the undirected fix for the degree-based node selection used when the
+graph is truncated to ``max_nodes``: the previous implementation counted only
+outgoing edges (``OPTIONAL MATCH (n)-[r]->()``), so a node that is mostly an edge
+*target* was under-ranked and could be dropped on truncation. The fix selects
+top-degree nodes via native SQL that counts both ``start_id`` and ``end_id``
+(undirected degree), LEFT JOIN-ing the base vertex table so isolated nodes are
+preserved when the graph is not truncated.
+
+The subgraph edge read intentionally stays directed: ``a`` iterates over every
+selected node, so ``(a)-[r]->(b)`` already captures every edge whose endpoints
+are both selected. These tests guard against an accidental regression there.
+
+All tests mock ``PGGraphStorage._query`` and inspect the SQL it receives.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch
+
+from lightrag.kg.postgres_impl import PGGraphStorage
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def make_graph_storage() -> PGGraphStorage:
+    """Construct a PGGraphStorage instance with a mocked _query method."""
+    storage = PGGraphStorage.__new__(PGGraphStorage)
+    storage.workspace = "test_ws"
+    storage.namespace = "test_graph"
+    storage.graph_name = "test_graph"
+    storage.global_config = {"max_graph_nodes": 1000}
+    storage.db = MagicMock()
+    return storage
+
+
+class _QueryCapture:
+    """Dispatch the three _query calls of the '*' branch by SQL content."""
+
+    def __init__(self, *, total_nodes, degree_rows, subgraph_rows):
+        self._total_nodes = total_nodes
+        self._degree_rows = degree_rows
+        self._subgraph_rows = subgraph_rows
+        self.count_sql = None
+        self.degree_sql = None
+        self.degree_params = None
+        self.subgraph_sql = None
+
+    def as_side_effect(self):
+        """Return an ``async def`` so AsyncMock awaits it (a callable instance is not)."""
+
+        async def fake_query(query, **kwargs):
+            if "count(distinct n)" in query:
+                self.count_sql = query
+                return [{"total_nodes": self._total_nodes}]
+            if "node_degrees" in query:
+                self.degree_sql = query
+                self.degree_params = kwargs.get("params")
+                return self._degree_rows
+            # subgraph read
+            self.subgraph_sql = query
+            return self._subgraph_rows
+
+        return fake_query
+
+
+def _node(node_id, entity_id):
+    return {"id": node_id, "properties": {"entity_id": entity_id}}
+
+
+def _edge(edge_id, start_id, end_id, weight="1"):
+    return {
+        "id": edge_id,
+        "label": "DIRECTED",
+        "start_id": start_id,
+        "end_id": end_id,
+        "properties": {"weight": weight},
+    }
+
+
+# ---------------------------------------------------------------------------
+# degree node selection SQL
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_degree_selection_sql_is_undirected_and_preserves_isolated():
+    """The degree query must count start_id + end_id and keep degree-0 nodes."""
+    capture = _QueryCapture(
+        total_nodes=2,
+        degree_rows=[{"node_id": 1, "degree": 3}, {"node_id": 2, "degree": 0}],
+        subgraph_rows=[{"a": _node(1, "Alice"), "r": None, "b": None}],
+    )
+    storage = make_graph_storage()
+
+    with patch.object(storage, "_query", side_effect=capture.as_side_effect()):
+        await storage.get_knowledge_graph("*", max_nodes=50)
+
+    sql = capture.degree_sql
+    assert sql is not None, "degree selection query was never issued"
+    # Undirected: both edge endpoints are counted.
+    assert "start_id" in sql
+    assert "end_id" in sql
+    assert "UNION ALL" in sql
+    # Isolated nodes preserved.
+    assert "LEFT JOIN" in sql
+    assert "COALESCE" in sql
+    # Stable ordering with id tie-break.
+    assert "ORDER BY degree DESC" in sql
+    assert "v.id ASC" in sql
+    # The old outgoing-only Cypher must be gone.
+    assert "-[r]->()" not in sql
+    assert "OPTIONAL MATCH (n)-[r]->()" not in sql
+
+
+@pytest.mark.asyncio
+async def test_degree_selection_limit_is_parameterized():
+    """max_nodes must be passed via params, not interpolated into the SQL."""
+    capture = _QueryCapture(
+        total_nodes=2,
+        degree_rows=[{"node_id": 1, "degree": 3}],
+        subgraph_rows=[{"a": _node(1, "Alice"), "r": None, "b": None}],
+    )
+    storage = make_graph_storage()
+
+    with patch.object(storage, "_query", side_effect=capture.as_side_effect()):
+        await storage.get_knowledge_graph("*", max_nodes=37)
+
+    assert "LIMIT $1" in capture.degree_sql
+    assert "37" not in capture.degree_sql
+    assert capture.degree_params == {"limit": 37}
+
+
+@pytest.mark.asyncio
+async def test_isolated_node_survives_end_to_end():
+    """A degree-0 node selected by the degree query reaches the final KnowledgeGraph."""
+    capture = _QueryCapture(
+        total_nodes=2,
+        degree_rows=[{"node_id": 1, "degree": 2}, {"node_id": 2, "degree": 0}],
+        subgraph_rows=[
+            {"a": _node(1, "Alice"), "r": None, "b": None},
+            # MATCH (a) still returns the isolated node with null r/b.
+            {"a": _node(2, "Bob"), "r": None, "b": None},
+        ],
+    )
+    storage = make_graph_storage()
+
+    with patch.object(storage, "_query", side_effect=capture.as_side_effect()):
+        kg = await storage.get_knowledge_graph("*", max_nodes=50)
+
+    # The isolated node id is formatted into the subgraph id list...
+    assert "2" in capture.subgraph_sql
+    # ...and present in the returned graph.
+    assert {node.id for node in kg.nodes} == {"1", "2"}
+    assert kg.is_truncated is False
+
+
+# ---------------------------------------------------------------------------
+# subgraph edge read — regression guard
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_subgraph_read_stays_directed_and_dedupes_edges():
+    """Subgraph read keeps the directed (a)-[r]->(b) match and dedupes by edge id."""
+    duplicate_edge = _edge(10, 1, 2)
+    capture = _QueryCapture(
+        total_nodes=2,
+        degree_rows=[{"node_id": 1, "degree": 1}, {"node_id": 2, "degree": 1}],
+        subgraph_rows=[
+            {"a": _node(1, "Alice"), "r": duplicate_edge, "b": _node(2, "Bob")},
+            # Same AGE edge id appearing twice must collapse to one edge.
+            {"a": _node(1, "Alice"), "r": duplicate_edge, "b": _node(2, "Bob")},
+        ],
+    )
+    storage = make_graph_storage()
+
+    with patch.object(storage, "_query", side_effect=capture.as_side_effect()):
+        kg = await storage.get_knowledge_graph("*", max_nodes=50)
+
+    assert "(a)-[r]->(b)" in capture.subgraph_sql
+    assert len(kg.edges) == 1
+    edge = kg.edges[0]
+    assert edge.source == "1"
+    assert edge.target == "2"

tests/kg/postgres_impl/test_postgres_graph_batch.py

@@ -0,0 +1,236 @@
+"""Unit tests for PGGraphStorage chunk-level transaction batch paths.
+
+Covers the PR-2 chunk-level fallback: ``upsert_nodes_batch`` /
+``upsert_edges_batch`` group per-row Cypher into payload/record-bounded chunks
+run in a single transaction; ``remove_nodes`` runs all chunks in ONE
+transaction (all-or-nothing); ``remove_edges`` runs one transaction per chunk.
+"""
+
+import pytest
+from unittest.mock import AsyncMock
+
+from lightrag.kg.postgres_impl import PGGraphStorage
+
+
+# ---------------------------------------------------------------------------
+# Capture harness
+# ---------------------------------------------------------------------------
+
+
+class _Capture:
+    def __init__(self) -> None:
+        self.calls: list[dict] = []  # every connection.execute(sql, *args)
+        self.tx_count = 0  # connection.transaction() opens
+        self.run_count = 0  # _run_with_retry invocations (= chunks issued)
+
+
+class _FakeTransaction:
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+
+class _FakeConnection:
+    def __init__(self, capture: _Capture):
+        self._capture = capture
+
+    def transaction(self):
+        self._capture.tx_count += 1
+        return _FakeTransaction()
+
+    async def execute(self, sql, *args):
+        self._capture.calls.append({"sql": sql, "args": args})
+        return ""
+
+
+def make_graph_storage(
+    *,
+    max_upsert_records: int | None = None,
+    max_upsert_payload: int | None = None,
+    max_delete_records: int | None = None,
+) -> tuple[PGGraphStorage, _Capture]:
+    storage = PGGraphStorage.__new__(PGGraphStorage)
+    storage.workspace = "test_ws"
+    storage.namespace = "test_graph"
+    storage.graph_name = "test_graph"
+    storage.__post_init__()  # resolves the chunk-level batch limits
+    if max_upsert_records is not None:
+        storage._max_upsert_records_per_batch = max_upsert_records
+    if max_upsert_payload is not None:
+        storage._max_upsert_payload_bytes = max_upsert_payload
+    if max_delete_records is not None:
+        storage._max_delete_records_per_batch = max_delete_records
+
+    capture = _Capture()
+    conn = _FakeConnection(capture)
+
+    async def fake_run_with_retry(operation, **_kwargs):
+        capture.run_count += 1
+        return await operation(conn)
+
+    storage.db = AsyncMock()
+    storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
+    return storage, capture
+
+
+def _sql_count(capture: _Capture, needle: str) -> int:
+    return sum(1 for c in capture.calls if needle in c["sql"])
+
+
+# ---------------------------------------------------------------------------
+# upsert_nodes_batch
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_upsert_nodes_batch_splits_into_chunk_transactions():
+    storage, cap = make_graph_storage(max_upsert_records=2)
+    nodes = [(f"n{i}", {"entity_id": f"n{i}", "v": str(i)}) for i in range(5)]
+
+    await storage.upsert_nodes_batch(nodes)
+
+    # 5 nodes / cap 2 => 3 chunks, each its own _run_with_retry + transaction.
+    assert cap.run_count == 3
+    assert cap.tx_count == 3
+    # One MERGE per node.
+    assert _sql_count(cap, "MERGE (n:base") == 5
+
+
+@pytest.mark.asyncio
+async def test_upsert_nodes_batch_dedupes_last_write_wins():
+    storage, cap = make_graph_storage()
+    await storage.upsert_nodes_batch(
+        [
+            ("A", {"entity_id": "A", "weight": "first"}),
+            ("A", {"entity_id": "A", "weight": "second"}),
+        ]
+    )
+
+    merge_calls = [c for c in cap.calls if "MERGE (n:base" in c["sql"]]
+    assert len(merge_calls) == 1
+    assert '"second"' in merge_calls[0]["sql"]
+    assert '"first"' not in merge_calls[0]["sql"]
+
+
+@pytest.mark.asyncio
+async def test_upsert_nodes_batch_splits_by_payload_bytes():
+    # Disable the record cap, force a small byte budget so each big node lands
+    # in its own chunk.
+    storage, cap = make_graph_storage(max_upsert_records=0, max_upsert_payload=200)
+    nodes = [(f"n{i}", {"entity_id": f"n{i}", "blob": "X" * 150}) for i in range(3)]
+
+    await storage.upsert_nodes_batch(nodes)
+
+    assert cap.run_count == 3
+
+
+@pytest.mark.asyncio
+async def test_upsert_nodes_batch_missing_entity_id_raises_before_tx():
+    storage, cap = make_graph_storage()
+    with pytest.raises(ValueError, match="entity_id"):
+        await storage.upsert_nodes_batch([("A", {"foo": "bar"})])
+    # The builder runs before the transaction opens, so nothing was issued.
+    assert cap.run_count == 0
+
+
+@pytest.mark.asyncio
+async def test_upsert_nodes_batch_empty_noop():
+    storage, cap = make_graph_storage()
+    await storage.upsert_nodes_batch([])
+    assert cap.run_count == 0
+
+
+# ---------------------------------------------------------------------------
+# upsert_edges_batch
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_upsert_edges_batch_splits_into_chunk_transactions():
+    storage, cap = make_graph_storage(max_upsert_records=2)
+    edges = [(f"s{i}", f"t{i}", {"w": str(i)}) for i in range(5)]
+
+    await storage.upsert_edges_batch(edges)
+
+    # 5 edges / cap 2 => 3 chunks; each chunk = one transaction.
+    assert cap.run_count == 3
+    assert cap.tx_count == 3
+    # One CREATE cypher per edge, and exactly one graph-wide lock per chunk.
+    assert _sql_count(cap, "CREATE (source)-[r:DIRECTED") == 5
+    assert _sql_count(cap, "pg_advisory_xact_lock") == 3
+
+
+@pytest.mark.asyncio
+async def test_upsert_edges_batch_takes_one_graph_lock_per_chunk():
+    """The batch chunk takes a single graph-wide advisory lock (keyed on
+    graph_name only), not one lock per edge -- bounded regardless of edge count."""
+    storage, cap = make_graph_storage()
+    await storage.upsert_edges_batch(
+        [(f"s{i}", f"t{i}", {"w": str(i)}) for i in range(4)]
+    )
+
+    lock_calls = [c for c in cap.calls if "pg_advisory_xact_lock" in c["sql"]]
+    assert len(lock_calls) == 1  # one chunk, one lock (not 4)
+    assert lock_calls[0]["args"] == ("test_graph",)  # graph-keyed, no endpoints
+    assert _sql_count(cap, "CREATE (source)-[r:DIRECTED") == 4
+    # The lock is acquired before any edge cypher in the chunk.
+    assert "pg_advisory_xact_lock" in cap.calls[0]["sql"]
+
+
+@pytest.mark.asyncio
+async def test_upsert_edges_batch_empty_noop():
+    storage, cap = make_graph_storage()
+    await storage.upsert_edges_batch([])
+    assert cap.run_count == 0
+
+
+# ---------------------------------------------------------------------------
+# remove_nodes — all chunks in ONE transaction (all-or-nothing)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_remove_nodes_chunks_share_one_transaction():
+    storage, cap = make_graph_storage(max_delete_records=2)
+    await storage.remove_nodes([f"n{i}" for i in range(5)])
+
+    # One transaction wraps all chunks (preserves single-statement atomicity).
+    assert cap.run_count == 1
+    assert cap.tx_count == 1
+    # 5 ids / cap 2 => 3 bounded IN [...] DETACH DELETE statements.
+    assert _sql_count(cap, "DETACH DELETE n") == 3
+
+
+@pytest.mark.asyncio
+async def test_remove_nodes_empty_noop():
+    storage, cap = make_graph_storage()
+    await storage.remove_nodes([])
+    assert cap.run_count == 0
+
+
+# ---------------------------------------------------------------------------
+# remove_edges — one transaction per chunk
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_remove_edges_one_transaction_per_chunk():
+    storage, cap = make_graph_storage(max_delete_records=2)
+    edges = [(f"s{i}", f"t{i}") for i in range(5)]
+
+    await storage.remove_edges(edges)
+
+    # 5 edges / cap 2 => 3 chunks, each its own transaction.
+    assert cap.run_count == 3
+    assert cap.tx_count == 3
+    # One DELETE r statement per edge.
+    assert _sql_count(cap, "DELETE r") == 5
+
+
+@pytest.mark.asyncio
+async def test_remove_edges_empty_noop():
+    storage, cap = make_graph_storage()
+    await storage.remove_edges([])
+    assert cap.run_count == 0

tests/kg/postgres_impl/test_postgres_halfvec.py

@@ -0,0 +1,374 @@
+import pytest
+import numpy as np
+from unittest.mock import patch, AsyncMock
+from lightrag.utils import EmbeddingFunc
+from lightrag.kg.postgres_impl import (
+    PGVectorStorage,
+    PostgreSQLDB,
+    _safe_index_name,
+)
+from lightrag.exceptions import DataMigrationError
+from lightrag.namespace import NameSpace
+
+
+# Mock PostgreSQLDB
+@pytest.fixture
+def mock_pg_db():
+    """Mock PostgreSQL database connection"""
+    db = AsyncMock()
+    db.workspace = "test_workspace"
+    db.vector_index_type = None
+
+    # Mock query responses: list for search queries (multirows=True), dict for DDL checks
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        if multirows:
+            return []
+        return {"exists": False, "count": 0}
+
+    # Mock for execute that mimics PostgreSQLDB.execute() behavior
+    async def mock_execute(sql, data=None, **kwargs):
+        return None
+
+    db.query = AsyncMock(side_effect=mock_query)
+    db.execute = AsyncMock(side_effect=mock_execute)
+
+    return db
+
+
+# Mock get_data_init_lock to avoid async lock issues in tests
+@pytest.fixture(autouse=True)
+def mock_data_init_lock():
+    with patch("lightrag.kg.postgres_impl.get_data_init_lock") as mock_lock:
+        mock_lock_ctx = AsyncMock()
+        mock_lock.return_value = mock_lock_ctx
+        yield mock_lock
+
+
+# Mock ClientManager
+@pytest.fixture
+def mock_client_manager(mock_pg_db):
+    with patch("lightrag.kg.postgres_impl.ClientManager") as mock_manager:
+        mock_manager.get_client = AsyncMock(return_value=mock_pg_db)
+        mock_manager.release_client = AsyncMock()
+        yield mock_manager
+
+
+# Mock Embedding function
+@pytest.fixture
+def mock_embedding_func():
+    async def embed_func(texts, **kwargs):
+        return np.array([[0.1] * 768 for _ in texts])
+
+    # Note: EmbeddingFunc in this version of lightrag supports model_name
+    func = EmbeddingFunc(embedding_dim=768, func=embed_func, model_name="test_model")
+    return func
+
+
+@pytest.mark.asyncio
+async def test_postgres_halfvec_table_creation(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """Test if table is created with HALFVEC type when HNSW_HALFVEC is selected"""
+    # Set index type to HNSW_HALFVEC
+    mock_pg_db.vector_index_type = "HNSW_HALFVEC"
+
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=mock_embedding_func,
+        workspace="test_ws",
+    )
+
+    # Mock table doesn't exist
+    mock_pg_db.check_table_exists = AsyncMock(return_value=False)
+
+    # Initialize storage (should trigger table creation)
+    await storage.initialize()
+
+    # Verify table creation SQL contains HALFVEC(768)
+    create_table_calls = [
+        call
+        for call in mock_pg_db.execute.call_args_list
+        if "CREATE TABLE" in call[0][0]
+    ]
+
+    assert len(create_table_calls) > 0
+    create_sql = create_table_calls[0][0][0]
+    assert "HALFVEC(768)" in create_sql
+    assert "VECTOR(768)" not in create_sql
+
+
+@pytest.mark.asyncio
+async def test_postgres_vector_table_creation_default(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """Test if table is created with default VECTOR type when other index type is selected"""
+    # Set index type to HNSW (default)
+    mock_pg_db.vector_index_type = "HNSW"
+
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=mock_embedding_func,
+        workspace="test_ws",
+    )
+
+    # Mock table doesn't exist
+    mock_pg_db.check_table_exists = AsyncMock(return_value=False)
+
+    # Initialize storage (should trigger table creation)
+    await storage.initialize()
+
+    # Verify table creation SQL contains VECTOR(768)
+    create_table_calls = [
+        call
+        for call in mock_pg_db.execute.call_args_list
+        if "CREATE TABLE" in call[0][0]
+    ]
+
+    assert len(create_table_calls) > 0
+    create_sql = create_table_calls[0][0][0]
+    assert "VECTOR(768)" in create_sql
+    assert "HALFVEC(768)" not in create_sql
+
+
+# Namespaces that use vector search SQL templates (query path)
+QUERY_NAMESPACES = [
+    NameSpace.VECTOR_STORE_CHUNKS,
+    NameSpace.VECTOR_STORE_ENTITIES,
+    NameSpace.VECTOR_STORE_RELATIONSHIPS,
+]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("namespace", QUERY_NAMESPACES)
+async def test_query_uses_halfvec_cast_when_hnsw_halfvec(
+    mock_client_manager, mock_pg_db, mock_embedding_func, namespace
+):
+    """When HNSW_HALFVEC is set, generated search SQL uses ::halfvec (not ::vector)."""
+    mock_pg_db.vector_index_type = "HNSW_HALFVEC"
+    mock_pg_db.check_table_exists = AsyncMock(return_value=True)
+
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+    storage = PGVectorStorage(
+        namespace=namespace,
+        global_config=config,
+        embedding_func=mock_embedding_func,
+        workspace="test_ws",
+    )
+    await storage.initialize()
+
+    query_embedding = [0.1] * 768
+    await storage.query("test query", top_k=5, query_embedding=query_embedding)
+
+    assert mock_pg_db.query.called
+    call_args = mock_pg_db.query.call_args
+    sql = call_args[0][0]
+    assert "::halfvec" in sql
+    assert "::vector" not in sql
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("namespace", QUERY_NAMESPACES)
+async def test_query_uses_vector_cast_when_hnsw_default(
+    mock_client_manager, mock_pg_db, mock_embedding_func, namespace
+):
+    """When HNSW (default) is set, generated search SQL uses ::vector (not ::halfvec)."""
+    mock_pg_db.vector_index_type = "HNSW"
+    mock_pg_db.check_table_exists = AsyncMock(return_value=True)
+
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+    storage = PGVectorStorage(
+        namespace=namespace,
+        global_config=config,
+        embedding_func=mock_embedding_func,
+        workspace="test_ws",
+    )
+    await storage.initialize()
+
+    query_embedding = [0.1] * 768
+    await storage.query("test query", top_k=5, query_embedding=query_embedding)
+
+    assert mock_pg_db.query.called
+    call_args = mock_pg_db.query.call_args
+    sql = call_args[0][0]
+    assert "::vector" in sql
+    assert "::halfvec" not in sql
+
+
+# ---------------------------------------------------------------------------
+# Index switching: old conflicting indexes are dropped
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_create_vector_index_drops_old_indexes_when_switching(mock_pg_db):
+    """Switching from HNSW to HNSW_HALFVEC drops the old hnsw_cosine index."""
+    mock_pg_db.vector_index_type = "HNSW_HALFVEC"
+    mock_pg_db.hnsw_m = 16
+    mock_pg_db.hnsw_ef = 64
+    mock_pg_db.ivfflat_lists = 100
+    mock_pg_db.vchordrq_build_options = ""
+
+    table_name = "lightrag_vdb_chunks_test"
+
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        if "pg_indexes" in sql:
+            return None
+        return None
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query)
+    mock_pg_db.execute = AsyncMock()
+
+    # Call the real method with mock_pg_db as self
+    await PostgreSQLDB._create_vector_index(mock_pg_db, table_name, 3072)
+
+    execute_calls = [call[0][0] for call in mock_pg_db.execute.call_args_list]
+
+    old_hnsw_name = _safe_index_name(table_name, "hnsw_cosine")
+    old_ivfflat_name = _safe_index_name(table_name, "ivfflat_cosine")
+    old_vchordrq_name = _safe_index_name(table_name, "vchordrq_cosine")
+
+    drop_calls = [c for c in execute_calls if "DROP INDEX IF EXISTS" in c]
+    dropped_names = {c.split("DROP INDEX IF EXISTS ")[1].strip() for c in drop_calls}
+    assert old_hnsw_name in dropped_names
+    assert old_ivfflat_name in dropped_names
+    assert old_vchordrq_name in dropped_names
+
+    new_index_name = _safe_index_name(table_name, "hnsw_halfvec_cosine")
+    assert new_index_name not in dropped_names
+
+    alter_calls = [c for c in execute_calls if "ALTER TABLE" in c]
+    assert any("HALFVEC(3072)" in c for c in alter_calls)
+
+    create_calls = [c for c in execute_calls if "CREATE INDEX" in c]
+    assert any("halfvec_cosine_ops" in c for c in create_calls)
+
+
+@pytest.mark.asyncio
+async def test_create_vector_index_no_drop_when_index_exists(mock_pg_db):
+    """If the target index already exists, no DROP or CREATE is issued."""
+    mock_pg_db.vector_index_type = "HNSW_HALFVEC"
+    mock_pg_db.hnsw_m = 16
+    mock_pg_db.hnsw_ef = 64
+    mock_pg_db.ivfflat_lists = 100
+    mock_pg_db.vchordrq_build_options = ""
+
+    table_name = "lightrag_vdb_chunks_test"
+
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        if "pg_indexes" in sql:
+            return {"?column?": 1}
+        return None
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query)
+    mock_pg_db.execute = AsyncMock()
+
+    await PostgreSQLDB._create_vector_index(mock_pg_db, table_name, 3072)
+
+    execute_calls = [call[0][0] for call in mock_pg_db.execute.call_args_list]
+    assert not any("DROP INDEX" in c for c in execute_calls)
+    assert not any("CREATE INDEX" in c for c in execute_calls)
+
+
+# ---------------------------------------------------------------------------
+# HalfVector dimension detection in setup_table
+# ---------------------------------------------------------------------------
+
+
+class _MockHalfVector:
+    """Mimics pgvector.halfvec.HalfVector for testing dimension detection."""
+
+    def __init__(self, dim: int):
+        self._dim = dim
+
+    def dimensions(self) -> int:
+        return self._dim
+
+    def to_list(self):
+        return [0.0] * self._dim
+
+
+@pytest.mark.asyncio
+async def test_setup_table_detects_halfvector_dimension_mismatch(mock_pg_db):
+    """DataMigrationError is raised when a HalfVector column has a different dimension."""
+    table_name = "lightrag_vdb_chunks_new"
+    legacy_table = "lightrag_vdb_chunks"
+
+    mock_pg_db.check_table_exists = AsyncMock(
+        side_effect=lambda t: t.lower() == legacy_table.lower()
+    )
+
+    call_count = 0
+
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if "COUNT(*)" in sql:
+            return {"count": 5}
+        if "content_vector" in sql:
+            return {"content_vector": _MockHalfVector(1024)}
+        return None
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query)
+    mock_pg_db.execute = AsyncMock()
+
+    with pytest.raises(DataMigrationError, match="Dimension mismatch"):
+        await PGVectorStorage.setup_table(
+            db=mock_pg_db,
+            table_name=table_name,
+            workspace="test_ws",
+            embedding_dim=768,
+            legacy_table_name=legacy_table,
+            base_table=legacy_table,
+        )
+
+
+@pytest.mark.asyncio
+async def test_setup_table_accepts_matching_halfvector_dimension(mock_pg_db):
+    """No error when HalfVector dimension matches the expected embedding_dim."""
+    table_name = "lightrag_vdb_chunks_new"
+    legacy_table = "lightrag_vdb_chunks"
+
+    mock_pg_db.check_table_exists = AsyncMock(
+        side_effect=lambda t: t.lower() == legacy_table.lower()
+    )
+    mock_pg_db.vector_index_type = "HNSW_HALFVEC"
+
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        if "COUNT(*)" in sql:
+            return {"count": 5}
+        if "content_vector" in sql:
+            return {"content_vector": _MockHalfVector(768)}
+        if multirows:
+            return []
+        return None
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query)
+    mock_pg_db.execute = AsyncMock()
+
+    with patch.object(PGVectorStorage, "_pg_create_table", new_callable=AsyncMock):
+        await PGVectorStorage.setup_table(
+            db=mock_pg_db,
+            table_name=table_name,
+            workspace="test_ws",
+            embedding_dim=768,
+            legacy_table_name=legacy_table,
+            base_table=legacy_table,
+        )

tests/kg/postgres_impl/test_postgres_index_name.py

@@ -0,0 +1,210 @@
+"""
+Unit tests for PostgreSQL safe index name generation.
+
+This module tests the _safe_index_name helper function which prevents
+PostgreSQL's silent 63-byte identifier truncation from causing index
+lookup failures.
+"""
+
+import pytest
+
+# Mark all tests as offline (no external dependencies)
+pytestmark = pytest.mark.offline
+
+
+class TestSafeIndexName:
+    """Test suite for _safe_index_name function."""
+
+    def test_short_name_unchanged(self):
+        """Short index names should remain unchanged."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        # Short table name - should return unchanged
+        result = _safe_index_name("lightrag_vdb_entity", "hnsw_cosine")
+        assert result == "idx_lightrag_vdb_entity_hnsw_cosine"
+        assert len(result.encode("utf-8")) <= 63
+
+    def test_long_name_gets_hashed(self):
+        """Long table names exceeding 63 bytes should get hashed."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        # Long table name that would exceed 63 bytes
+        long_table_name = "LIGHTRAG_VDB_ENTITY_text_embedding_3_large_3072d"
+        result = _safe_index_name(long_table_name, "hnsw_cosine")
+
+        # Should be within 63 bytes
+        assert len(result.encode("utf-8")) <= 63
+
+        # Should start with idx_ prefix
+        assert result.startswith("idx_")
+
+        # Should contain the suffix
+        assert result.endswith("_hnsw_cosine")
+
+        # Should NOT be the naive concatenation (which would be truncated)
+        naive_name = f"idx_{long_table_name.lower()}_hnsw_cosine"
+        assert result != naive_name
+
+    def test_deterministic_output(self):
+        """Same input should always produce same output (deterministic)."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        table_name = "LIGHTRAG_VDB_CHUNKS_text_embedding_3_large_3072d"
+        suffix = "hnsw_cosine"
+
+        result1 = _safe_index_name(table_name, suffix)
+        result2 = _safe_index_name(table_name, suffix)
+
+        assert result1 == result2
+
+    def test_different_suffixes_different_results(self):
+        """Different suffixes should produce different index names."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        table_name = "LIGHTRAG_VDB_ENTITY_text_embedding_3_large_3072d"
+
+        result1 = _safe_index_name(table_name, "hnsw_cosine")
+        result2 = _safe_index_name(table_name, "ivfflat_cosine")
+
+        assert result1 != result2
+
+    def test_case_insensitive(self):
+        """Table names should be normalized to lowercase."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        result_upper = _safe_index_name("LIGHTRAG_VDB_ENTITY", "hnsw_cosine")
+        result_lower = _safe_index_name("lightrag_vdb_entity", "hnsw_cosine")
+
+        assert result_upper == result_lower
+
+    def test_boundary_case_exactly_63_bytes(self):
+        """Test boundary case where name is exactly at 63-byte limit."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        # Create a table name that results in exactly 63 bytes
+        # idx_ (4) + table_name + _ (1) + suffix = 63
+        # So table_name + suffix = 58
+
+        # Test a name that's just under the limit (should remain unchanged)
+        short_suffix = "id"
+        # idx_ (4) + 56 chars + _ (1) + id (2) = 63
+        table_56 = "a" * 56
+        result = _safe_index_name(table_56, short_suffix)
+        expected = f"idx_{table_56}_{short_suffix}"
+        assert result == expected
+        assert len(result.encode("utf-8")) == 63
+
+    def test_unicode_handling(self):
+        """Unicode characters should be properly handled (bytes, not chars)."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        # Unicode characters can take more bytes than visible chars
+        # Chinese characters are 3 bytes each in UTF-8
+        table_name = "lightrag_测试_table"  # Contains Chinese chars
+        result = _safe_index_name(table_name, "hnsw_cosine")
+
+        # Should always be within 63 bytes
+        assert len(result.encode("utf-8")) <= 63
+
+    def test_real_world_model_names(self):
+        """Test with real-world embedding model names that cause issues."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        # These are actual model names that have caused issues
+        test_cases = [
+            ("LIGHTRAG_VDB_CHUNKS_text_embedding_3_large_3072d", "hnsw_cosine"),
+            ("LIGHTRAG_VDB_ENTITY_text_embedding_3_large_3072d", "hnsw_cosine"),
+            ("LIGHTRAG_VDB_RELATION_text_embedding_3_large_3072d", "hnsw_cosine"),
+            (
+                "LIGHTRAG_VDB_ENTITY_bge_m3_1024d",
+                "hnsw_cosine",
+            ),  # Shorter model name
+            (
+                "LIGHTRAG_VDB_CHUNKS_nomic_embed_text_v1_768d",
+                "ivfflat_cosine",
+            ),  # Different index type
+        ]
+
+        for table_name, suffix in test_cases:
+            result = _safe_index_name(table_name, suffix)
+
+            # Critical: must be within PostgreSQL's 63-byte limit
+            assert len(result.encode("utf-8")) <= 63, (
+                f"Index name too long: {result} for table {table_name}"
+            )
+
+            # Must have consistent format
+            assert result.startswith("idx_"), f"Missing idx_ prefix: {result}"
+            assert result.endswith(f"_{suffix}"), f"Missing suffix {suffix}: {result}"
+
+    def test_hash_uniqueness_for_similar_tables(self):
+        """Similar but different table names should produce different hashes."""
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        # These tables have similar names but should have different hashes
+        tables = [
+            "LIGHTRAG_VDB_CHUNKS_model_a_1024d",
+            "LIGHTRAG_VDB_CHUNKS_model_b_1024d",
+            "LIGHTRAG_VDB_ENTITY_model_a_1024d",
+        ]
+
+        results = [_safe_index_name(t, "hnsw_cosine") for t in tables]
+
+        # All results should be unique
+        assert len(set(results)) == len(results), "Hash collision detected!"
+
+
+class TestIndexNameIntegration:
+    """Integration-style tests for index name usage patterns."""
+
+    def test_pg_indexes_lookup_compatibility(self):
+        """
+        Test that the generated index name will work with pg_indexes lookup.
+
+        This is the core problem: PostgreSQL stores the truncated name,
+        but we were looking up the untruncated name. Our fix ensures we
+        always use a name that fits within 63 bytes.
+        """
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        table_name = "LIGHTRAG_VDB_CHUNKS_text_embedding_3_large_3072d"
+        suffix = "hnsw_cosine"
+
+        # Generate the index name
+        index_name = _safe_index_name(table_name, suffix)
+
+        # Simulate what PostgreSQL would store (truncate at 63 bytes)
+        stored_name = index_name.encode("utf-8")[:63].decode("utf-8", errors="ignore")
+
+        # The key fix: our generated name should equal the stored name
+        # because it's already within the 63-byte limit
+        assert index_name == stored_name, (
+            "Index name would be truncated by PostgreSQL, causing lookup failures!"
+        )
+
+    def test_backward_compatibility_short_names(self):
+        """
+        Ensure backward compatibility with existing short index names.
+
+        For tables that have existing indexes with short names (pre-model-suffix era),
+        the function should not change their names.
+        """
+        from lightrag.kg.postgres_impl import _safe_index_name
+
+        # Legacy table names without model suffix
+        legacy_tables = [
+            "LIGHTRAG_VDB_ENTITY",
+            "LIGHTRAG_VDB_RELATION",
+            "LIGHTRAG_VDB_CHUNKS",
+        ]
+
+        for table in legacy_tables:
+            for suffix in ["hnsw_cosine", "ivfflat_cosine", "id"]:
+                result = _safe_index_name(table, suffix)
+                expected = f"idx_{table.lower()}_{suffix}"
+
+                # Short names should remain unchanged for backward compatibility
+                if len(expected.encode("utf-8")) <= 63:
+                    assert result == expected, (
+                        f"Short name changed unexpectedly: {result} != {expected}"
+                    )

tests/kg/postgres_impl/test_postgres_migration.py

@@ -0,0 +1,856 @@
+import pytest
+from unittest.mock import patch, AsyncMock
+import numpy as np
+from lightrag.utils import EmbeddingFunc
+from lightrag.kg.postgres_impl import (
+    PGVectorStorage,
+)
+from lightrag.namespace import NameSpace
+
+
+# Mock PostgreSQLDB
+@pytest.fixture
+def mock_pg_db():
+    """Mock PostgreSQL database connection"""
+    db = AsyncMock()
+    db.workspace = "test_workspace"
+
+    # Mock query responses with multirows support
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        # Default return value
+        if multirows:
+            return []  # Return empty list for multirows
+        return {"exists": False, "count": 0}
+
+    # Mock for execute that mimics PostgreSQLDB.execute() behavior
+    async def mock_execute(sql, data=None, **kwargs):
+        """
+        Mock that mimics PostgreSQLDB.execute() behavior:
+        - Accepts data as dict[str, Any] | None (second parameter)
+        - Internally converts dict.values() to tuple for AsyncPG
+        """
+        # Mimic real execute() which accepts dict and converts to tuple
+        if data is not None and not isinstance(data, dict):
+            raise TypeError(
+                f"PostgreSQLDB.execute() expects data as dict, got {type(data).__name__}"
+            )
+        return None
+
+    db.query = AsyncMock(side_effect=mock_query)
+    db.execute = AsyncMock(side_effect=mock_execute)
+
+    return db
+
+
+# Mock get_data_init_lock to avoid async lock issues in tests
+@pytest.fixture(autouse=True)
+def mock_data_init_lock():
+    with patch("lightrag.kg.postgres_impl.get_data_init_lock") as mock_lock:
+        mock_lock_ctx = AsyncMock()
+        mock_lock.return_value = mock_lock_ctx
+        yield mock_lock
+
+
+# Mock ClientManager
+@pytest.fixture
+def mock_client_manager(mock_pg_db):
+    with patch("lightrag.kg.postgres_impl.ClientManager") as mock_manager:
+        mock_manager.get_client = AsyncMock(return_value=mock_pg_db)
+        mock_manager.release_client = AsyncMock()
+        yield mock_manager
+
+
+# Mock Embedding function
+@pytest.fixture
+def mock_embedding_func():
+    async def embed_func(texts, **kwargs):
+        return np.array([[0.1] * 768 for _ in texts])
+
+    func = EmbeddingFunc(embedding_dim=768, func=embed_func, model_name="test_model")
+    return func
+
+
+async def test_postgres_table_naming(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """Test if table name is correctly generated with model suffix"""
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=mock_embedding_func,
+        workspace="test_ws",
+    )
+
+    # Verify table name contains model suffix
+    expected_suffix = "test_model_768d"
+    assert expected_suffix in storage.table_name
+    assert storage.table_name == f"LIGHTRAG_VDB_CHUNKS_{expected_suffix}"
+
+    # Verify legacy table name
+    assert storage.legacy_table_name == "LIGHTRAG_VDB_CHUNKS"
+
+
+async def test_postgres_migration_trigger(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """Test if migration logic is triggered correctly"""
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=mock_embedding_func,
+        workspace="test_ws",
+    )
+
+    # Setup mocks for migration scenario
+    # 1. New table does not exist, legacy table exists
+    async def mock_check_table_exists(table_name):
+        return table_name == storage.legacy_table_name
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
+
+    # 2. Legacy table has 100 records
+    mock_rows = [
+        {"id": f"test_id_{i}", "content": f"content_{i}", "workspace": "test_ws"}
+        for i in range(100)
+    ]
+    migration_state = {"new_table_count": 0}
+
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        if "COUNT(*)" in sql:
+            sql_upper = sql.upper()
+            legacy_table = storage.legacy_table_name.upper()
+            new_table = storage.table_name.upper()
+            is_new_table = new_table in sql_upper
+            is_legacy_table = legacy_table in sql_upper and not is_new_table
+
+            if is_new_table:
+                return {"count": migration_state["new_table_count"]}
+            if is_legacy_table:
+                return {"count": 100}
+            return {"count": 0}
+        elif multirows and "SELECT *" in sql:
+            # Mock batch fetch for migration using keyset pagination
+            # New pattern: WHERE workspace = $1 AND id > $2 ORDER BY id LIMIT $3
+            # or first batch: WHERE workspace = $1 ORDER BY id LIMIT $2
+            if "WHERE workspace" in sql:
+                if "id >" in sql:
+                    # Keyset pagination: params = [workspace, last_id, limit]
+                    last_id = params[1] if len(params) > 1 else None
+                    # Find rows after last_id
+                    start_idx = 0
+                    for i, row in enumerate(mock_rows):
+                        if row["id"] == last_id:
+                            start_idx = i + 1
+                            break
+                    limit = params[2] if len(params) > 2 else 500
+                else:
+                    # First batch (no last_id): params = [workspace, limit]
+                    start_idx = 0
+                    limit = params[1] if len(params) > 1 else 500
+            else:
+                # No workspace filter with keyset
+                if "id >" in sql:
+                    last_id = params[0] if params else None
+                    start_idx = 0
+                    for i, row in enumerate(mock_rows):
+                        if row["id"] == last_id:
+                            start_idx = i + 1
+                            break
+                    limit = params[1] if len(params) > 1 else 500
+                else:
+                    start_idx = 0
+                    limit = params[0] if params else 500
+            end = min(start_idx + limit, len(mock_rows))
+            return mock_rows[start_idx:end]
+        return {}
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query)
+
+    # Track migration through _run_with_retry calls
+    migration_executed = []
+
+    async def mock_run_with_retry(operation, **kwargs):
+        # Track that migration batch operation was called
+        migration_executed.append(True)
+        migration_state["new_table_count"] = 100
+        return None
+
+    mock_pg_db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)
+
+    with patch(
+        "lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
+    ):
+        # Initialize storage (should trigger migration)
+        await storage.initialize()
+
+        # Verify migration was executed by checking _run_with_retry was called
+        # (batch migration uses _run_with_retry with executemany)
+        assert len(migration_executed) > 0, "Migration should have been executed"
+
+
+async def test_postgres_no_migration_needed(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """Test scenario where new table already exists (no migration needed)"""
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=mock_embedding_func,
+        workspace="test_ws",
+    )
+
+    # Mock: new table already exists
+    async def mock_check_table_exists(table_name):
+        return table_name == storage.table_name
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
+
+    with patch(
+        "lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
+    ) as mock_create:
+        await storage.initialize()
+
+        # Verify no table creation was attempted
+        mock_create.assert_not_called()
+
+
+async def test_scenario_1_new_workspace_creation(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """
+    Scenario 1: New workspace creation
+
+    Expected behavior:
+    - No legacy table exists
+    - Directly create new table with model suffix
+    - No migration needed
+    """
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    embedding_func = EmbeddingFunc(
+        embedding_dim=3072,
+        func=mock_embedding_func.func,
+        model_name="text-embedding-3-large",
+    )
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=embedding_func,
+        workspace="new_workspace",
+    )
+
+    # Mock: neither table exists
+    async def mock_check_table_exists(table_name):
+        return False
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
+
+    with patch(
+        "lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
+    ) as mock_create:
+        await storage.initialize()
+
+        # Verify table name format
+        assert "text_embedding_3_large_3072d" in storage.table_name
+
+        # Verify new table creation was called
+        mock_create.assert_called_once()
+        call_args = mock_create.call_args
+        assert (
+            call_args[0][1] == storage.table_name
+        )  # table_name is second positional arg
+
+
+async def test_scenario_2_legacy_upgrade_migration(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """
+    Scenario 2: Upgrade from legacy version
+
+    Expected behavior:
+    - Legacy table exists (without model suffix)
+    - New table doesn't exist
+    - Automatically migrate data to new table with suffix
+    """
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    embedding_func = EmbeddingFunc(
+        embedding_dim=1536,
+        func=mock_embedding_func.func,
+        model_name="text-embedding-ada-002",
+    )
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=embedding_func,
+        workspace="legacy_workspace",
+    )
+
+    # Mock: only legacy table exists
+    async def mock_check_table_exists(table_name):
+        return table_name == storage.legacy_table_name
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
+
+    # Mock: legacy table has 50 records
+    mock_rows = [
+        {
+            "id": f"legacy_id_{i}",
+            "content": f"legacy_content_{i}",
+            "workspace": "legacy_workspace",
+        }
+        for i in range(50)
+    ]
+
+    # Track which queries have been made for proper response
+    query_history = []
+    migration_state = {"new_table_count": 0}
+
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        query_history.append(sql)
+
+        if "COUNT(*)" in sql:
+            # Determine table type:
+            # - Legacy: contains base name but NOT model suffix
+            # - New: contains model suffix (e.g., text_embedding_ada_002_1536d)
+            sql_upper = sql.upper()
+            base_name = storage.legacy_table_name.upper()
+
+            # Check if this is querying the new table (has model suffix)
+            has_model_suffix = storage.table_name.upper() in sql_upper
+
+            is_legacy_table = base_name in sql_upper and not has_model_suffix
+            has_workspace_filter = "WHERE workspace" in sql
+
+            if is_legacy_table and has_workspace_filter:
+                # Count for legacy table with workspace filter (before migration)
+                return {"count": 50}
+            elif is_legacy_table and not has_workspace_filter:
+                # Total count for legacy table
+                return {"count": 50}
+            else:
+                # New table count (before/after migration)
+                return {"count": migration_state["new_table_count"]}
+        elif multirows and "SELECT *" in sql:
+            # Mock batch fetch for migration using keyset pagination
+            # New pattern: WHERE workspace = $1 AND id > $2 ORDER BY id LIMIT $3
+            # or first batch: WHERE workspace = $1 ORDER BY id LIMIT $2
+            if "WHERE workspace" in sql:
+                if "id >" in sql:
+                    # Keyset pagination: params = [workspace, last_id, limit]
+                    last_id = params[1] if len(params) > 1 else None
+                    # Find rows after last_id
+                    start_idx = 0
+                    for i, row in enumerate(mock_rows):
+                        if row["id"] == last_id:
+                            start_idx = i + 1
+                            break
+                    limit = params[2] if len(params) > 2 else 500
+                else:
+                    # First batch (no last_id): params = [workspace, limit]
+                    start_idx = 0
+                    limit = params[1] if len(params) > 1 else 500
+            else:
+                # No workspace filter with keyset
+                if "id >" in sql:
+                    last_id = params[0] if params else None
+                    start_idx = 0
+                    for i, row in enumerate(mock_rows):
+                        if row["id"] == last_id:
+                            start_idx = i + 1
+                            break
+                    limit = params[1] if len(params) > 1 else 500
+                else:
+                    start_idx = 0
+                    limit = params[0] if params else 500
+            end = min(start_idx + limit, len(mock_rows))
+            return mock_rows[start_idx:end]
+        return {}
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query)
+
+    # Track migration through _run_with_retry calls
+    migration_executed = []
+
+    async def mock_run_with_retry(operation, **kwargs):
+        # Track that migration batch operation was called
+        migration_executed.append(True)
+        migration_state["new_table_count"] = 50
+        return None
+
+    mock_pg_db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)
+
+    with patch(
+        "lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
+    ) as mock_create:
+        await storage.initialize()
+
+        # Verify table name contains ada-002
+        assert "text_embedding_ada_002_1536d" in storage.table_name
+
+        # Verify migration was executed (batch migration uses _run_with_retry)
+        assert len(migration_executed) > 0, "Migration should have been executed"
+        mock_create.assert_called_once()
+
+
+async def test_scenario_3_multi_model_coexistence(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """
+    Scenario 3: Multiple embedding models coexist
+
+    Expected behavior:
+    - Different embedding models create separate tables
+    - Tables are isolated by model suffix
+    - No interference between different models
+    """
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    # Workspace A: uses bge-small (768d)
+    embedding_func_a = EmbeddingFunc(
+        embedding_dim=768, func=mock_embedding_func.func, model_name="bge-small"
+    )
+
+    storage_a = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=embedding_func_a,
+        workspace="workspace_a",
+    )
+
+    # Workspace B: uses bge-large (1024d)
+    async def embed_func_b(texts, **kwargs):
+        return np.array([[0.1] * 1024 for _ in texts])
+
+    embedding_func_b = EmbeddingFunc(
+        embedding_dim=1024, func=embed_func_b, model_name="bge-large"
+    )
+
+    storage_b = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=embedding_func_b,
+        workspace="workspace_b",
+    )
+
+    # Verify different table names
+    assert storage_a.table_name != storage_b.table_name
+    assert "bge_small_768d" in storage_a.table_name
+    assert "bge_large_1024d" in storage_b.table_name
+
+    # Mock: both tables don't exist yet
+    async def mock_check_table_exists(table_name):
+        return False
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
+
+    with patch(
+        "lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
+    ) as mock_create:
+        # Initialize both storages
+        await storage_a.initialize()
+        await storage_b.initialize()
+
+        # Verify two separate tables were created
+        assert mock_create.call_count == 2
+
+        # Verify table names are different
+        call_args_list = mock_create.call_args_list
+        table_names = [call[0][1] for call in call_args_list]  # Second positional arg
+        assert len(set(table_names)) == 2  # Two unique table names
+        assert storage_a.table_name in table_names
+        assert storage_b.table_name in table_names
+
+
+async def test_case1_empty_legacy_auto_cleanup(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """
+    Case 1a: Both new and legacy tables exist, but legacy is EMPTY
+    Expected: Automatically delete empty legacy table (safe cleanup)
+    """
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    embedding_func = EmbeddingFunc(
+        embedding_dim=1536,
+        func=mock_embedding_func.func,
+        model_name="test-model",
+    )
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=embedding_func,
+        workspace="test_ws",
+    )
+
+    # Mock: Both tables exist
+    async def mock_check_table_exists(table_name):
+        return True  # Both new and legacy exist
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
+
+    # Mock: Legacy table is empty (0 records)
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        if "COUNT(*)" in sql:
+            if storage.legacy_table_name in sql:
+                return {"count": 0}  # Empty legacy table
+            else:
+                return {"count": 100}  # New table has data
+        return {}
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query)
+
+    with patch("lightrag.kg.postgres_impl.logger"):
+        await storage.initialize()
+
+        # Verify: Empty legacy table should be automatically cleaned up
+        # Empty tables are safe to delete without data loss risk
+        delete_calls = [
+            call
+            for call in mock_pg_db.execute.call_args_list
+            if call[0][0] and "DROP TABLE" in call[0][0]
+        ]
+        assert len(delete_calls) >= 1, "Empty legacy table should be auto-deleted"
+        # Check if legacy table was dropped
+        dropped_table = storage.legacy_table_name
+        assert any(dropped_table in str(call) for call in delete_calls), (
+            f"Expected to drop empty legacy table '{dropped_table}'"
+        )
+
+        print(
+            f"✅ Case 1a: Empty legacy table '{dropped_table}' auto-deleted successfully"
+        )
+
+
+async def test_case1_nonempty_legacy_warning(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """
+    Case 1b: Both new and legacy tables exist, and legacy HAS DATA
+    Expected: Log warning, do not delete legacy (preserve data)
+    """
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    embedding_func = EmbeddingFunc(
+        embedding_dim=1536,
+        func=mock_embedding_func.func,
+        model_name="test-model",
+    )
+
+    storage = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=embedding_func,
+        workspace="test_ws",
+    )
+
+    # Mock: Both tables exist
+    async def mock_check_table_exists(table_name):
+        return True  # Both new and legacy exist
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
+
+    # Mock: Legacy table has data (50 records)
+    async def mock_query(sql, params=None, multirows=False, **kwargs):
+        if "COUNT(*)" in sql:
+            if storage.legacy_table_name in sql:
+                return {"count": 50}  # Legacy has data
+            else:
+                return {"count": 100}  # New table has data
+        return {}
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query)
+
+    with patch("lightrag.kg.postgres_impl.logger"):
+        await storage.initialize()
+
+        # Verify: Legacy table with data should be preserved
+        # We never auto-delete tables that contain data to prevent accidental data loss
+        delete_calls = [
+            call
+            for call in mock_pg_db.execute.call_args_list
+            if call[0][0] and "DROP TABLE" in call[0][0]
+        ]
+        # Check if legacy table was deleted (it should not be)
+        dropped_table = storage.legacy_table_name
+        legacy_deleted = any(dropped_table in str(call) for call in delete_calls)
+        assert not legacy_deleted, "Legacy table with data should NOT be auto-deleted"
+
+        print(
+            f"✅ Case 1b: Legacy table '{dropped_table}' with data preserved (warning only)"
+        )
+
+
+async def test_case1_sequential_workspace_migration(
+    mock_client_manager, mock_pg_db, mock_embedding_func
+):
+    """
+    Case 1c: Sequential workspace migration (Multi-tenant scenario)
+
+    Critical bug fix verification:
+    Timeline:
+    1. Legacy table has workspace_a (3 records) + workspace_b (3 records)
+    2. Workspace A initializes first → Case 3 (only legacy exists) → migrates A's data
+    3. Workspace B initializes later → Case 3 (both tables exist, legacy has B's data) → should migrate B's data
+    4. Verify workspace B's data is correctly migrated to new table
+
+    This test verifies the migration logic correctly handles multi-tenant scenarios
+    where different workspaces migrate sequentially.
+    """
+    config = {
+        "embedding_batch_num": 10,
+        "vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
+    }
+
+    embedding_func = EmbeddingFunc(
+        embedding_dim=1536,
+        func=mock_embedding_func.func,
+        model_name="test-model",
+    )
+
+    # Mock data: Legacy table has 6 records total (3 from workspace_a, 3 from workspace_b)
+    mock_rows_a = [
+        {"id": f"a_{i}", "content": f"A content {i}", "workspace": "workspace_a"}
+        for i in range(3)
+    ]
+    mock_rows_b = [
+        {"id": f"b_{i}", "content": f"B content {i}", "workspace": "workspace_b"}
+        for i in range(3)
+    ]
+
+    # Track migration state
+    migration_state = {
+        "new_table_exists": False,
+        "workspace_a_migrated": False,
+        "workspace_a_migration_count": 0,
+        "workspace_b_migration_count": 0,
+    }
+
+    # Step 1: Simulate workspace_a initialization (Case 3 - only legacy exists)
+    # CRITICAL: Set db.workspace to workspace_a
+    mock_pg_db.workspace = "workspace_a"
+
+    storage_a = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=embedding_func,
+        workspace="workspace_a",
+    )
+
+    # Mock table_exists for workspace_a
+    async def mock_check_table_exists_a(table_name):
+        if table_name == storage_a.legacy_table_name:
+            return True
+        if table_name == storage_a.table_name:
+            return migration_state["new_table_exists"]
+        return False
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists_a)
+
+    # Mock query for workspace_a (Case 3)
+    async def mock_query_a(sql, params=None, multirows=False, **kwargs):
+        sql_upper = sql.upper()
+        base_name = storage_a.legacy_table_name.upper()
+
+        if "COUNT(*)" in sql:
+            has_model_suffix = "TEST_MODEL_1536D" in sql_upper
+            is_legacy = base_name in sql_upper and not has_model_suffix
+            has_workspace_filter = "WHERE workspace" in sql
+
+            if is_legacy and has_workspace_filter:
+                workspace = params[0] if params and len(params) > 0 else None
+                if workspace == "workspace_a":
+                    return {"count": 3}
+                elif workspace == "workspace_b":
+                    return {"count": 3}
+            elif is_legacy and not has_workspace_filter:
+                # Global count in legacy table
+                return {"count": 6}
+            elif has_model_suffix:
+                if has_workspace_filter:
+                    workspace = params[0] if params and len(params) > 0 else None
+                    if workspace == "workspace_a":
+                        return {"count": migration_state["workspace_a_migration_count"]}
+                    if workspace == "workspace_b":
+                        return {"count": migration_state["workspace_b_migration_count"]}
+                return {
+                    "count": migration_state["workspace_a_migration_count"]
+                    + migration_state["workspace_b_migration_count"]
+                }
+        elif multirows and "SELECT *" in sql:
+            if "WHERE workspace" in sql:
+                workspace = params[0] if params and len(params) > 0 else None
+                if workspace == "workspace_a":
+                    # Handle keyset pagination
+                    if "id >" in sql:
+                        # params = [workspace, last_id, limit]
+                        last_id = params[1] if len(params) > 1 else None
+                        start_idx = 0
+                        for i, row in enumerate(mock_rows_a):
+                            if row["id"] == last_id:
+                                start_idx = i + 1
+                                break
+                        limit = params[2] if len(params) > 2 else 500
+                    else:
+                        # First batch: params = [workspace, limit]
+                        start_idx = 0
+                        limit = params[1] if len(params) > 1 else 500
+                    end = min(start_idx + limit, len(mock_rows_a))
+                    return mock_rows_a[start_idx:end]
+        return {}
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query_a)
+
+    # Track migration via _run_with_retry (batch migration uses this)
+    migration_a_executed = []
+
+    async def mock_run_with_retry_a(operation, **kwargs):
+        migration_a_executed.append(True)
+        migration_state["workspace_a_migration_count"] = len(mock_rows_a)
+        return None
+
+    mock_pg_db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry_a)
+
+    # Initialize workspace_a (Case 3)
+    with patch("lightrag.kg.postgres_impl.logger"):
+        await storage_a.initialize()
+        migration_state["new_table_exists"] = True
+        migration_state["workspace_a_migrated"] = True
+
+    print("✅ Step 1: Workspace A initialized")
+    # Verify migration was executed via _run_with_retry (batch migration uses executemany)
+    assert len(migration_a_executed) > 0, (
+        "Migration should have been executed for workspace_a"
+    )
+    print(f"✅ Step 1: Migration executed {len(migration_a_executed)} batch(es)")
+
+    # Step 2: Simulate workspace_b initialization (Case 3 - both exist, but legacy has B's data)
+    # CRITICAL: Set db.workspace to workspace_b
+    mock_pg_db.workspace = "workspace_b"
+
+    storage_b = PGVectorStorage(
+        namespace=NameSpace.VECTOR_STORE_CHUNKS,
+        global_config=config,
+        embedding_func=embedding_func,
+        workspace="workspace_b",
+    )
+
+    mock_pg_db.reset_mock()
+
+    # Mock table_exists for workspace_b (both exist)
+    async def mock_check_table_exists_b(table_name):
+        return True  # Both tables exist
+
+    mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists_b)
+
+    # Mock query for workspace_b (Case 3)
+    async def mock_query_b(sql, params=None, multirows=False, **kwargs):
+        sql_upper = sql.upper()
+        base_name = storage_b.legacy_table_name.upper()
+
+        if "COUNT(*)" in sql:
+            has_model_suffix = "TEST_MODEL_1536D" in sql_upper
+            is_legacy = base_name in sql_upper and not has_model_suffix
+            has_workspace_filter = "WHERE workspace" in sql
+
+            if is_legacy and has_workspace_filter:
+                workspace = params[0] if params and len(params) > 0 else None
+                if workspace == "workspace_b":
+                    return {"count": 3}  # workspace_b still has data in legacy
+                elif workspace == "workspace_a":
+                    return {"count": 0}  # workspace_a already migrated
+            elif is_legacy and not has_workspace_filter:
+                # Global count: only workspace_b data remains
+                return {"count": 3}
+            elif has_model_suffix:
+                if has_workspace_filter:
+                    workspace = params[0] if params and len(params) > 0 else None
+                    if workspace == "workspace_b":
+                        return {"count": migration_state["workspace_b_migration_count"]}
+                    elif workspace == "workspace_a":
+                        return {"count": 3}
+                else:
+                    return {"count": 3 + migration_state["workspace_b_migration_count"]}
+        elif multirows and "SELECT *" in sql:
+            if "WHERE workspace" in sql:
+                workspace = params[0] if params and len(params) > 0 else None
+                if workspace == "workspace_b":
+                    # Handle keyset pagination
+                    if "id >" in sql:
+                        # params = [workspace, last_id, limit]
+                        last_id = params[1] if len(params) > 1 else None
+                        start_idx = 0
+                        for i, row in enumerate(mock_rows_b):
+                            if row["id"] == last_id:
+                                start_idx = i + 1
+                                break
+                        limit = params[2] if len(params) > 2 else 500
+                    else:
+                        # First batch: params = [workspace, limit]
+                        start_idx = 0
+                        limit = params[1] if len(params) > 1 else 500
+                    end = min(start_idx + limit, len(mock_rows_b))
+                    return mock_rows_b[start_idx:end]
+        return {}
+
+    mock_pg_db.query = AsyncMock(side_effect=mock_query_b)
+
+    # Track migration via _run_with_retry for workspace_b
+    migration_b_executed = []
+
+    async def mock_run_with_retry_b(operation, **kwargs):
+        migration_b_executed.append(True)
+        migration_state["workspace_b_migration_count"] = len(mock_rows_b)
+        return None
+
+    mock_pg_db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry_b)
+
+    # Initialize workspace_b (Case 3 - both tables exist)
+    with patch("lightrag.kg.postgres_impl.logger"):
+        await storage_b.initialize()
+
+    print("✅ Step 2: Workspace B initialized")
+
+    # Verify workspace_b migration happens when new table has no workspace_b data
+    # but legacy table still has workspace_b data.
+    assert len(migration_b_executed) > 0, (
+        "Migration should have been executed for workspace_b"
+    )
+    print("✅ Step 2: Migration executed for workspace_b")
+
+    print("\n🎉 Case 1c: Sequential workspace migration verification complete!")
+    print("   - Workspace A: Migrated successfully (only legacy existed)")
+    print("   - Workspace B: Migrated successfully (new table empty for workspace_b)")

tests/kg/postgres_impl/test_postgres_performance_timing.py

@@ -0,0 +1,138 @@
+import importlib
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+import lightrag.utils as utils_module
+from lightrag.kg.postgres_impl import PGGraphStorage, PostgreSQLDB
+from lightrag.namespace import NameSpace
+
+
+def make_db() -> PostgreSQLDB:
+    return PostgreSQLDB(
+        {
+            "host": "localhost",
+            "port": 5432,
+            "user": "postgres",
+            "password": "postgres",
+            "database": "postgres",
+            "workspace": "test_ws",
+            "max_connections": 10,
+            "connection_retry_attempts": 3,
+            "connection_retry_backoff": 0,
+            "connection_retry_backoff_max": 0,
+            "pool_close_timeout": 5.0,
+        }
+    )
+
+
+@pytest.mark.asyncio
+async def test_execute_timing_logs_success():
+    db = make_db()
+
+    async def fake_run_with_retry(operation, **kwargs):
+        conn = AsyncMock()
+        conn.execute = AsyncMock(return_value="INSERT 0 1")
+        await operation(conn)
+
+    db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
+
+    with patch("lightrag.kg.postgres_impl.performance_timing_log") as timing_log:
+        await db.execute("SELECT 1", timing_label="test label")
+
+    assert any(
+        "connection.execute completed" in call.args[0]
+        for call in timing_log.call_args_list
+    )
+
+
+@pytest.mark.asyncio
+async def test_execute_timing_logs_failure():
+    db = make_db()
+
+    async def fake_run_with_retry(operation, **kwargs):
+        conn = AsyncMock()
+        conn.execute = AsyncMock(side_effect=RuntimeError("boom"))
+        await operation(conn)
+
+    db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
+
+    with patch("lightrag.kg.postgres_impl.performance_timing_log") as timing_log:
+        with pytest.raises(RuntimeError, match="boom"):
+            await db.execute("SELECT 1", timing_label="test label")
+
+    assert any(
+        "connection.execute failed" in call.args[0]
+        for call in timing_log.call_args_list
+    )
+
+
+@pytest.mark.asyncio
+async def test_graph_upsert_node_passes_timing_label():
+    storage = PGGraphStorage(
+        namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
+        workspace="test_ws",
+        global_config={},
+        embedding_func=AsyncMock(),
+    )
+    storage.graph_name = "test_graph"
+    storage._query = AsyncMock(return_value=[])
+
+    await storage.upsert_node(
+        "node-1",
+        {
+            "entity_id": "node-1",
+            "description": "desc",
+        },
+    )
+
+    assert storage._query.await_args.kwargs["timing_label"] == (
+        "test_ws PGGraphStorage.upsert_node"
+    )
+
+
+@pytest.mark.asyncio
+async def test_graph_upsert_edge_passes_timing_label():
+    storage = PGGraphStorage(
+        namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
+        workspace="test_ws",
+        global_config={},
+        embedding_func=AsyncMock(),
+    )
+    storage.graph_name = "test_graph"
+    # upsert_edge drives the lock + cypher via db._run_with_retry, not _query.
+    storage.db = MagicMock()
+    storage.db._run_with_retry = AsyncMock(return_value=None)
+
+    await storage.upsert_edge(
+        "node-1",
+        "node-2",
+        {
+            "weight": 1.0,
+            "description": "desc",
+        },
+    )
+
+    assert storage.db._run_with_retry.await_args.kwargs["timing_label"] == (
+        "test_ws PGGraphStorage.upsert_edge"
+    )
+
+
+def test_performance_timing_logs_reads_new_env_only(monkeypatch):
+    with monkeypatch.context() as m:
+        m.setenv("LIGHTRAG_DOC_QUERY_TIMING_LOGS", "false")
+        m.setenv("LIGHTRAG_PERFORMANCE_TIMING_LOGS", "true")
+        reloaded = importlib.reload(utils_module)
+        assert reloaded.PERFORMANCE_TIMING_LOGS is True
+
+    importlib.reload(utils_module)
+
+
+def test_performance_timing_logs_ignores_old_env(monkeypatch):
+    with monkeypatch.context() as m:
+        m.setenv("LIGHTRAG_DOC_QUERY_TIMING_LOGS", "true")
+        m.setenv("LIGHTRAG_PERFORMANCE_TIMING_LOGS", "false")
+        reloaded = importlib.reload(utils_module)
+        assert reloaded.PERFORMANCE_TIMING_LOGS is False
+
+    importlib.reload(utils_module)