import asyncio
import json
import os
import time
from typing import Any, final, Optional, Dict
from dataclasses import dataclass, fields
import numpy as np
from lightrag.utils import (
    logger,
    compute_mdhash_id,
    _cooperative_yield,
    validate_workspace,
)
from ..base import BaseVectorStorage
from ..constants import (
    DEFAULT_MAX_FILE_PATH_LENGTH,
    DEFAULT_QUERY_PRIORITY,
    GRAPH_FIELD_SEP,
)
from ..kg.shared_storage import get_data_init_lock, get_namespace_lock
import pipmaster as pm

if not pm.is_installed("pymilvus"):
    pm.install("pymilvus>=2.6.2")

import configparser
import grpc  # type: ignore
from pymilvus import (  # type: ignore
    MilvusClient,
    MilvusException,
    DataType,
    CollectionSchema,
    FieldSchema,
)
from packaging import version

config = configparser.ConfigParser()
config.read("config.ini", "utf-8")


@dataclass
class _PendingVectorDoc:
    """Buffered vector upsert waiting for embedding and/or bulk flush."""

    source: dict[str, Any]
    content: str
    vector: list[float] | None = None


# Flush-time batching limits. Milvus' server-side proxy rejects any single
# gRPC message larger than ~64MB (grpc.serverMaxRecvSize); the client library
# cannot raise that ceiling, so large flushes must be split client-side.
# The payload-byte budget is the primary limiter; the record-count caps are a
# secondary guard that only binds when individual records are small.
# Upsert and delete have separate count caps on purpose: upsert records each
# carry a full embedding vector and are far heavier than delete pks, so the
# upsert batch count is kept much smaller than the delete one.
DEFAULT_MILVUS_UPSERT_MAX_PAYLOAD_BYTES = (
    32 * 1024 * 1024
)  # 32MB, well below the 64MB gRPC ceiling
DEFAULT_MILVUS_UPSERT_MAX_RECORDS_PER_BATCH = 128
DEFAULT_MILVUS_DELETE_MAX_RECORDS_PER_BATCH = 1000

# Schema-migration resilience. A transient Milvus outage during the long
# iterator-based migration must not kill worker startup: when pymilvus'
# internal reconnect fails it closes the gRPC channel for good, so every later
# call on the same client raises "Cannot invoke RPC on closed channel!". On a
# connection-class failure the whole migration attempt is therefore retried
# from scratch with a rebuilt MilvusClient (the source collection is untouched
# until the final rename and each attempt drops the leftover _temp collection
# first, so a full re-run is always safe). The max backoff is kept above
# pymilvus' connection-pool idle health-check threshold (IDLE_THRESHOLD_SECONDS
# = 30s in pymilvus 3.x) so a rebuilt client is guaranteed to get a
# health-checked/recovered channel rather than the same dead pooled handler.
DEFAULT_MILVUS_MIGRATION_MAX_RETRIES = 5
DEFAULT_MILVUS_MIGRATION_RETRY_BACKOFF_SECONDS = 5.0
DEFAULT_MILVUS_MIGRATION_RETRY_MAX_BACKOFF_SECONDS = 60.0
MILVUS_MIGRATION_RETRY_BACKOFF_MULTIPLIER = 3.0
DEFAULT_MILVUS_MIGRATION_ITERATOR_BATCH_SIZE = 2000

# Schema-migration memory back-pressure. The bulk copy inserts the whole source
# collection into the temp collection with no client-side throttle, so the
# Milvus data node accumulates growing insert-buffer segments until its own
# auto-flush catches up; under a large migration this can exhaust server memory
# (the temp collection is not loaded, so query-node memory is already bounded —
# this is the data-node write buffer). Flushing every N migrated rows seals
# those segments to object storage and blocks until the flush returns, giving
# the server natural back-pressure. 0 disables periodic flush (rely on Milvus
# auto-flush). The interval is kept coarse so it does not spawn many tiny
# segments (which would burden later compaction). An optional per-batch sleep
# lets a small server breathe between batches; 0 disables it.
DEFAULT_MILVUS_MIGRATION_FLUSH_INTERVAL_ROWS = 50000
DEFAULT_MILVUS_MIGRATION_BATCH_SLEEP_SECONDS = 0.0

# Substrings that mark an exception as a transient connection failure (worth a
# retry with a rebuilt client) rather than a schema/parameter error.
MILVUS_RETRYABLE_CONNECTION_ERROR_MARKERS = (
    "unavailable",  # grpc UNAVAILABLE / "server unavailable"
    "ping timeout",
    "deadline exceeded",
    "connection refused",
    "connection reset",
    "broken pipe",
    "closed channel",  # ValueError: Cannot invoke RPC on closed channel!
    "fail connecting to server",  # pymilvus _wait_for_channel_ready
    "failed to connect",
)

MILVUS_MAX_VARCHAR_BYTES = 65535
# The Milvus primary key. Truncating it would let two distinct ids collapse to
# the same key (silent overwrite) and make the row unreachable by its real id
# via get_by_id/delete, so it must never be truncated under any circumstance.
MILVUS_PRIMARY_KEY_FIELDS = frozenset({"id"})
# Non-primary identity fields. They are not the Milvus primary key, so the row
# stays uniquely keyed by `id` even if these collide after truncation (no
# storage-level overwrite). On the live upsert path we still reject oversize
# values so callers fix their input; during migration of pre-existing data we
# truncate-and-warn instead, so a single pathological legacy value cannot abort
# the whole collection migration.
MILVUS_IDENTITY_VARCHAR_FIELDS = frozenset(
    {"id", "entity_name", "full_doc_id", "src_id", "tgt_id"}
)
# Fields whose value is a GRAPH_FIELD_SEP-joined list of ids (chunk ids / file
# paths). When such a value overflows we truncate on the last separator that
# fits rather than mid-id, so we drop whole ids instead of leaving a dangling
# partial id that resolves to nothing.
MILVUS_SEPARATOR_JOINED_FIELDS = frozenset({"source_id", "file_path"})

# Supported index types
SUPPORTED_INDEX_TYPES = {
    "AUTOINDEX",
    "HNSW",
    "HNSW_SQ",
    "HNSW_PQ",
    "HNSW_PRQ",
    "IVF_FLAT",
    "IVF_SQ8",
    "IVF_PQ",
    "DISKANN",
    "SCANN",
}

# Supported metric types
SUPPORTED_METRIC_TYPES = {"COSINE", "L2", "IP"}

# HNSW_SQ quantization types
SUPPORTED_SQ_TYPES = {"SQ4U", "SQ6", "SQ8", "BF16", "FP16"}
SUPPORTED_REFINE_TYPES = {"SQ6", "SQ8", "BF16", "FP16", "FP32"}

# Index type version requirements
# Important: HNSW_SQ was first introduced in Milvus 2.6.8 (not 2.5)
INDEX_VERSION_REQUIREMENTS = {
    "HNSW_SQ": "2.6.8",  # HNSW_SQ requires Milvus 2.6.8+ (supports sq_types such as SQ4U, SQ6, SQ8, BF16, FP16)
}


def _get_env_bool(key: str, default: bool = False) -> bool:
    """Parse environment variable as boolean"""
    val = os.environ.get(key, "").lower()
    if val in ("true", "1", "yes", "on"):
        return True
    elif val in ("false", "0", "no", "off"):
        return False
    return default


def _get_env_int(key: str, default: int) -> int:
    """Parse environment variable as integer"""
    val = os.environ.get(key, "")
    if val:
        try:
            return int(val)
        except ValueError:
            logger.warning(
                f"Invalid integer value for {key}: {val}, using default {default}"
            )
    return default


@dataclass
class MilvusIndexConfig:
    """
    Milvus vector index configuration class

    Supports configuration via environment variables or initialization parameters.
    Initialization parameters take precedence over environment variables.
    """

    # Base configuration
    index_type: Optional[str] = None
    metric_type: Optional[str] = None

    # HNSW series parameters
    hnsw_m: Optional[int] = None
    hnsw_ef_construction: Optional[int] = None
    hnsw_ef: Optional[int] = None

    # HNSW_SQ specific parameters
    sq_type: Optional[str] = None
    sq_refine: Optional[bool] = None
    sq_refine_type: Optional[str] = None
    sq_refine_k: Optional[int] = None

    # IVF series parameters
    ivf_nlist: Optional[int] = None
    ivf_nprobe: Optional[int] = None

    def __post_init__(self):
        """Load configuration from environment variables (init parameters take precedence)"""
        # Index type
        self.index_type = (
            self.index_type or os.environ.get("MILVUS_INDEX_TYPE", "AUTOINDEX")
        ).upper()

        # Metric type
        self.metric_type = (
            self.metric_type or os.environ.get("MILVUS_METRIC_TYPE", "COSINE")
        ).upper()

        # HNSW parameters
        # Defaults aligned with Milvus 2.4+ official documentation
        if self.hnsw_m is None:
            self.hnsw_m = _get_env_int("MILVUS_HNSW_M", 16)
        if self.hnsw_ef_construction is None:
            self.hnsw_ef_construction = _get_env_int("MILVUS_HNSW_EF_CONSTRUCTION", 360)
        if self.hnsw_ef is None:
            self.hnsw_ef = _get_env_int("MILVUS_HNSW_EF", 200)

        # HNSW_SQ parameters
        if self.sq_type is None:
            self.sq_type = os.environ.get("MILVUS_HNSW_SQ_TYPE", "SQ8").upper()
        if self.sq_refine is None:
            self.sq_refine = _get_env_bool("MILVUS_HNSW_SQ_REFINE", False)
        if self.sq_refine_type is None:
            self.sq_refine_type = os.environ.get(
                "MILVUS_HNSW_SQ_REFINE_TYPE", "FP32"
            ).upper()
        if self.sq_refine_k is None:
            self.sq_refine_k = _get_env_int("MILVUS_HNSW_SQ_REFINE_K", 10)

        # IVF parameters
        if self.ivf_nlist is None:
            self.ivf_nlist = _get_env_int("MILVUS_IVF_NLIST", 1024)
        if self.ivf_nprobe is None:
            self.ivf_nprobe = _get_env_int("MILVUS_IVF_NPROBE", 16)

        # Validate configuration
        self._validate()

    def _validate(self):
        """Validate configuration validity"""
        if self.index_type not in SUPPORTED_INDEX_TYPES:
            raise ValueError(
                f"Unsupported index type: {self.index_type}. "
                f"Supported: {SUPPORTED_INDEX_TYPES}"
            )

        if self.metric_type not in SUPPORTED_METRIC_TYPES:
            raise ValueError(
                f"Unsupported metric type: {self.metric_type}. "
                f"Supported: {SUPPORTED_METRIC_TYPES}"
            )

        if self.index_type == "HNSW_SQ":
            if self.sq_type not in SUPPORTED_SQ_TYPES:
                raise ValueError(
                    f"Unsupported sq_type: {self.sq_type}. "
                    f"Supported: {SUPPORTED_SQ_TYPES}"
                )
            if self.sq_refine and self.sq_refine_type not in SUPPORTED_REFINE_TYPES:
                raise ValueError(
                    f"Unsupported refine_type: {self.sq_refine_type}. "
                    f"Supported: {SUPPORTED_REFINE_TYPES}"
                )

        # Parameter range validation
        if not (2 <= self.hnsw_m <= 2048):
            raise ValueError(f"hnsw_m must be in [2, 2048], got {self.hnsw_m}")
        if self.hnsw_ef_construction < 1:
            raise ValueError(
                f"hnsw_ef_construction must be >= 1, got {self.hnsw_ef_construction}"
            )
        if self.ivf_nlist < 1 or self.ivf_nlist > 65536:
            raise ValueError(f"ivf_nlist must be in [1, 65536], got {self.ivf_nlist}")

    def validate_milvus_version(self, server_version: str) -> None:
        """
        Validate Milvus server version supports the configured index type

        Args:
            server_version: Milvus server version string (e.g., "2.6.9")

        Raises:
            ValueError: Version does not meet index type requirements
        """
        current_ver = version.parse(
            server_version.split("-")[0]
        )  # Handle "2.6.9-dev" format

        # Check HNSW_SQ index type version requirements (requires 2.6.8+)
        if self.index_type == "HNSW_SQ":
            required = INDEX_VERSION_REQUIREMENTS["HNSW_SQ"]
            if current_ver < version.parse(required):
                raise ValueError(
                    f"HNSW_SQ requires Milvus {required}+, "
                    f"current version: {server_version}"
                )

        logger.info(
            f"Milvus version {server_version} validated for index type "
            f"{self.index_type}"
            + (f" with sq_type {self.sq_type}" if self.index_type == "HNSW_SQ" else "")
        )

    def build_index_params(self, index_params, field_name: str = "vector"):
        """
        Build pymilvus index parameters

        Args:
            index_params: IndexParams instance (from compatibility helper or client.prepare_index_params())
            field_name: Vector field name

        Returns:
            IndexParams object, or a dict fallback when direct API creation is needed.
        """
        if index_params is None:
            if self.index_type == "AUTOINDEX":
                logger.info(
                    "Using AUTOINDEX with direct API fallback because IndexParams is unavailable"
                )
                return {
                    "field_name": field_name,
                    "index_type": self.index_type,
                    "metric_type": self.metric_type,
                    "params": {},
                }
            raise RuntimeError(
                f"IndexParams not available but required for index type "
                f"'{self.index_type}'. Ensure pymilvus is installed correctly."
            )

        params: Dict[str, Any] = {}

        # HNSW series indexes
        if self.index_type in ("HNSW", "HNSW_SQ", "HNSW_PQ", "HNSW_PRQ"):
            params["M"] = self.hnsw_m
            params["efConstruction"] = self.hnsw_ef_construction

            # HNSW_SQ specific parameters
            if self.index_type == "HNSW_SQ":
                params["sq_type"] = self.sq_type
                if self.sq_refine:
                    params["refine"] = True
                    params["refine_type"] = self.sq_refine_type

        # IVF series indexes
        elif self.index_type in ("IVF_FLAT", "IVF_SQ8", "IVF_PQ"):
            params["nlist"] = self.ivf_nlist

        # DISKANN / SCANN have no additional params

        index_params.add_index(
            field_name=field_name,
            index_type=self.index_type,
            metric_type=self.metric_type,
            params=params,
        )

        logger.info(
            f"Milvus index configured: type={self.index_type}, "
            f"metric={self.metric_type}, params={params}"
        )

        return index_params

    def build_search_params(self) -> Dict[str, Any]:
        """
        Build search parameters

        Returns:
            Search parameters dictionary
        """
        search_params: Dict[str, Any] = {}

        if self.index_type in ("HNSW", "HNSW_SQ", "HNSW_PQ", "HNSW_PRQ"):
            search_params["ef"] = self.hnsw_ef
            if self.index_type == "HNSW_SQ" and self.sq_refine:
                search_params["refine_k"] = self.sq_refine_k

        elif self.index_type in ("IVF_FLAT", "IVF_SQ8", "IVF_PQ"):
            search_params["nprobe"] = self.ivf_nprobe

        return {"params": search_params} if search_params else {}

    @classmethod
    def get_config_field_names(cls) -> set:
        """Get all configuration field names from the dataclass.

        This method provides a single source of truth for configuration parameter names,
        eliminating the need to maintain duplicate hardcoded lists elsewhere.

        Returns:
            Set of field names that can be used to extract configuration from kwargs
        """
        return {f.name for f in fields(cls)}

    def to_dict(self) -> Dict[str, Any]:
        """Export configuration as dictionary (for logging/debugging)"""
        return {
            "index_type": self.index_type,
            "metric_type": self.metric_type,
            "hnsw_m": self.hnsw_m,
            "hnsw_ef_construction": self.hnsw_ef_construction,
            "hnsw_ef": self.hnsw_ef,
            "sq_type": self.sq_type if self.index_type == "HNSW_SQ" else None,
            "sq_refine": self.sq_refine if self.index_type == "HNSW_SQ" else None,
            "sq_refine_type": (
                self.sq_refine_type
                if self.index_type == "HNSW_SQ" and self.sq_refine
                else None
            ),
            "sq_refine_k": (
                self.sq_refine_k
                if self.index_type == "HNSW_SQ" and self.sq_refine
                else None
            ),
            "ivf_nlist": (
                self.ivf_nlist if self.index_type.startswith("IVF") else None
            ),
            "ivf_nprobe": (
                self.ivf_nprobe if self.index_type.startswith("IVF") else None
            ),
        }


@final
@dataclass
class MilvusVectorDBStorage(BaseVectorStorage):
    def _get_milvus_connection_kwargs(self, include_db_name: bool = True) -> dict:
        """Build Milvus connection kwargs from env/config."""
        connection_kwargs = {
            "uri": os.environ.get(
                "MILVUS_URI",
                config.get(
                    "milvus",
                    "uri",
                    fallback=os.path.join(
                        self.global_config["working_dir"], "milvus_lite.db"
                    ),
                ),
            ),
            "user": os.environ.get(
                "MILVUS_USER", config.get("milvus", "user", fallback=None)
            ),
            "password": os.environ.get(
                "MILVUS_PASSWORD",
                config.get("milvus", "password", fallback=None),
            ),
            "token": os.environ.get(
                "MILVUS_TOKEN", config.get("milvus", "token", fallback=None)
            ),
        }

        db_name = os.environ.get(
            "MILVUS_DB_NAME",
            config.get("milvus", "db_name", fallback=None),
        )
        if include_db_name and db_name:
            connection_kwargs["db_name"] = db_name

        return connection_kwargs

    def _get_milvus_db_name(self) -> Optional[str]:
        """Return the configured Milvus database name, if any."""
        db_name = self._get_milvus_connection_kwargs(include_db_name=True).get(
            "db_name"
        )
        if db_name is None:
            return None

        normalized_name = str(db_name).strip()
        return normalized_name or None

    def _create_milvus_client(self) -> MilvusClient:
        """Create a Milvus client and ensure the configured database exists."""
        client = MilvusClient(
            **self._get_milvus_connection_kwargs(include_db_name=False)
        )
        db_name = self._get_milvus_db_name()

        if not db_name:
            return client

        existing_databases = set(client.list_databases())
        if db_name not in existing_databases:
            logger.warning(
                f"[{self.workspace}] Milvus database '{db_name}' not found, creating it"
            )
            client.create_database(db_name)

        use_database = getattr(client, "use_database", None) or getattr(
            client, "using_database", None
        )
        if callable(use_database):
            use_database(db_name)
            logger.debug(
                f"[{self.workspace}] Using Milvus database '{db_name}' for namespace '{self.namespace}'"
            )
            return client

        return MilvusClient(**self._get_milvus_connection_kwargs(include_db_name=True))

    def _rebuild_milvus_client(self) -> None:
        """Replace the (possibly dead) client with a freshly created one.

        Once pymilvus' internal reconnect has failed, the gRPC channel is
        closed permanently and every later RPC on the same client raises
        "Cannot invoke RPC on closed channel!" — the client cannot heal
        itself, so migration retries must rebuild it. Safe during migration:
        it runs inside get_data_init_lock(), so this instance owns
        self._client exclusively. close() is best-effort — on a dead channel
        it is a no-op release. In pymilvus 3.x the pooled handler is
        health-checked and recovered in place, which also heals other clients
        sharing the same address.
        """
        old_client, self._client = self._client, None
        if old_client is not None:
            try:
                old_client.close()
            except Exception as close_error:
                logger.warning(
                    f"[{self.workspace}] Failed to close stale Milvus client: {close_error}"
                )
        self._client = self._create_milvus_client()

    @staticmethod
    def _is_retryable_connection_error(error: BaseException) -> bool:
        """Return True when the error chain indicates a transient connection failure.

        Walks __cause__/__context__ because the migration wraps low-level
        errors in RuntimeError and pymilvus wraps grpc errors in
        MilvusException. Schema, dimension and parameter errors fall through
        to False so they keep failing fast.
        """
        seen: set[int] = set()
        current: BaseException | None = error
        while current is not None and id(current) not in seen:
            seen.add(id(current))
            if isinstance(current, grpc.RpcError):
                code_getter = getattr(current, "code", None)
                code = code_getter() if callable(code_getter) else None
                if code in (
                    grpc.StatusCode.UNAVAILABLE,
                    grpc.StatusCode.DEADLINE_EXCEEDED,
                ):
                    return True
            elif isinstance(current, MilvusException):
                # Status.CONNECT_FAILED == 2 (client-side connect failure)
                message = str(current).lower()
                if current.code == 2 or any(
                    marker in message
                    for marker in MILVUS_RETRYABLE_CONNECTION_ERROR_MARKERS
                ):
                    return True
            elif isinstance(current, ValueError):
                # grpc raises ValueError("Cannot invoke RPC on closed channel!")
                if "closed channel" in str(current).lower():
                    return True
            current = current.__cause__ or current.__context__
        return False

    def _create_schema_for_namespace(self) -> CollectionSchema:
        """Create schema based on the current instance's namespace"""

        # Get vector dimension from embedding_func
        dimension = self.embedding_func.embedding_dim
        varchar_limits = self._get_varchar_field_limits_for_namespace()

        # Base fields (common to all collections)
        base_fields = [
            FieldSchema(
                name="id", dtype=DataType.VARCHAR, max_length=64, is_primary=True
            ),
            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension),
            FieldSchema(name="created_at", dtype=DataType.INT64),
        ]

        # Determine specific fields based on namespace
        if self.namespace.endswith("entities"):
            specific_fields = [
                FieldSchema(
                    name="entity_name",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["entity_name"],
                    nullable=True,
                ),
                FieldSchema(
                    name="content",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["content"],
                    nullable=True,
                ),
                FieldSchema(
                    name="source_id",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["source_id"],
                    nullable=True,
                ),
                FieldSchema(
                    name="file_path",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["file_path"],
                    nullable=True,
                ),
            ]
            description = "LightRAG entities vector storage"

        elif self.namespace.endswith("relationships"):
            specific_fields = [
                FieldSchema(
                    name="src_id",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["src_id"],
                    nullable=True,
                ),
                FieldSchema(
                    name="tgt_id",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["tgt_id"],
                    nullable=True,
                ),
                FieldSchema(
                    name="content",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["content"],
                    nullable=True,
                ),
                FieldSchema(
                    name="source_id",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["source_id"],
                    nullable=True,
                ),
                FieldSchema(
                    name="file_path",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["file_path"],
                    nullable=True,
                ),
            ]
            description = "LightRAG relationships vector storage"

        elif self.namespace.endswith("chunks"):
            specific_fields = [
                FieldSchema(
                    name="full_doc_id",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["full_doc_id"],
                    nullable=True,
                ),
                FieldSchema(
                    name="content",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["content"],
                    nullable=True,
                ),
                FieldSchema(
                    name="file_path",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["file_path"],
                    nullable=True,
                ),
            ]
            description = "LightRAG chunks vector storage"

        else:
            # Default generic schema (backward compatibility)
            specific_fields = [
                FieldSchema(
                    name="file_path",
                    dtype=DataType.VARCHAR,
                    max_length=varchar_limits["file_path"],
                    nullable=True,
                ),
            ]
            description = "LightRAG generic vector storage"

        # Merge all fields
        all_fields = base_fields + specific_fields

        return CollectionSchema(
            fields=all_fields,
            description=description,
            enable_dynamic_field=True,  # Support dynamic fields
        )

    def _get_varchar_field_limits_for_namespace(self) -> dict[str, int]:
        base_fields = {
            "id": 64,
            "content": MILVUS_MAX_VARCHAR_BYTES,
            "file_path": DEFAULT_MAX_FILE_PATH_LENGTH,
        }
        if self.namespace.endswith("entities"):
            return {
                **base_fields,
                "entity_name": 512,
                "source_id": MILVUS_MAX_VARCHAR_BYTES,
            }
        if self.namespace.endswith("relationships"):
            return {
                **base_fields,
                "src_id": 512,
                "tgt_id": 512,
                "source_id": MILVUS_MAX_VARCHAR_BYTES,
            }
        if self.namespace.endswith("chunks"):
            return {**base_fields, "full_doc_id": 64}
        return base_fields

    def _get_migrated_metadata_field_limits(self) -> dict[str, int]:
        if self.namespace.endswith("entities"):
            return {
                "content": MILVUS_MAX_VARCHAR_BYTES,
                "source_id": MILVUS_MAX_VARCHAR_BYTES,
            }
        if self.namespace.endswith("relationships"):
            return {
                "content": MILVUS_MAX_VARCHAR_BYTES,
                "source_id": MILVUS_MAX_VARCHAR_BYTES,
            }
        if self.namespace.endswith("chunks"):
            return {"content": MILVUS_MAX_VARCHAR_BYTES}
        return {}

    @staticmethod
    def _field_max_length(field: dict) -> int | None:
        max_length = field.get("params", {}).get("max_length")
        if max_length is None:
            return None
        try:
            return int(max_length)
        except (TypeError, ValueError):
            return None

    def _truncate_varchar_value(
        self,
        field_name: str,
        value: Any,
        record_id: str | None = None,
        allow_identity_truncation: bool = False,
    ) -> Any:
        limit = self._varchar_field_limits.get(field_name)
        if limit is None or not isinstance(value, str):
            return value

        encoded = value.encode("utf-8")
        if len(encoded) <= limit:
            return value

        # The primary key is never truncated: collapsing two ids into one would
        # silently overwrite a row and orphan it from get_by_id/delete.
        if field_name in MILVUS_PRIMARY_KEY_FIELDS:
            raise ValueError(
                f"[{self.workspace}] Milvus primary key '{field_name}' for record "
                f"'{record_id or '<unknown>'}' exceeds {limit} bytes "
                f"({len(encoded)} bytes); primary keys cannot be truncated"
            )

        # Other identity fields: reject on the live upsert path, but allow
        # truncate-and-warn during migration so legacy data can be carried over
        # without aborting the whole collection.
        if (
            field_name in MILVUS_IDENTITY_VARCHAR_FIELDS
            and not allow_identity_truncation
        ):
            raise ValueError(
                f"[{self.workspace}] Milvus field '{field_name}' for record "
                f"'{record_id or '<unknown>'}' exceeds {limit} bytes "
                f"({len(encoded)} bytes); identity fields cannot be truncated"
            )

        # Cut to the byte budget on a valid UTF-8 boundary first.
        truncated = encoded[:limit].decode("utf-8", errors="ignore")
        # For separator-joined id lists, back off to the last separator that
        # fits so we never persist a half id. Fall back to the raw byte cut when
        # no separator fits (e.g. a single id longer than the limit).
        if field_name in MILVUS_SEPARATOR_JOINED_FIELDS:
            boundary = truncated.rfind(GRAPH_FIELD_SEP)
            if boundary > 0:
                truncated = truncated[:boundary]
        logger.warning(
            "[%s] Milvus field '%s' for record '%s' truncated from %d to %d bytes",
            self.workspace,
            field_name,
            record_id or "<unknown>",
            len(encoded),
            len(truncated.encode("utf-8")),
        )
        return truncated

    def _sanitize_varchar_fields(
        self, row: dict[str, Any], allow_identity_truncation: bool = False
    ) -> dict[str, Any]:
        record_id = str(row.get("id", "")) or None
        return {
            field_name: self._truncate_varchar_value(
                field_name,
                value,
                record_id,
                allow_identity_truncation=allow_identity_truncation,
            )
            for field_name, value in row.items()
        }

    def _normalize_migration_row(self, row: dict[str, Any]) -> dict[str, Any]:
        normalized = dict(row)
        metadata = normalized.pop("$meta", None)
        if isinstance(metadata, dict):
            for field_name, value in metadata.items():
                # Explicit nullable fields can hold None while the real value
                # still lives in $meta (schema-drift rows), so backfill on None
                # rather than mere key presence.
                if normalized.get(field_name) is None:
                    normalized[field_name] = value
        # Migration carries pre-existing rows: non-primary identity fields are
        # truncated-and-warned rather than rejected (see _truncate_varchar_value).
        return self._sanitize_varchar_fields(normalized, allow_identity_truncation=True)

    def _get_index_params(self):
        """Get IndexParams in a version-compatible way"""
        try:
            # Try to use client's prepare_index_params method (most common)
            if hasattr(self._client, "prepare_index_params"):
                return self._client.prepare_index_params()
        except Exception:
            pass

        try:
            # Try to import IndexParams from different possible locations
            from pymilvus.client.prepare import IndexParams  # type: ignore

            return IndexParams()
        except ImportError:
            pass

        try:
            from pymilvus.client.types import IndexParams  # type: ignore

            return IndexParams()
        except ImportError:
            pass

        try:
            from pymilvus import IndexParams  # type: ignore

            return IndexParams()
        except ImportError:
            pass

        # If all else fails, return None to use fallback method
        return None

    def _create_scalar_index_fallback(self, field_name: str, index_type: str):
        """Fallback method to create scalar index using direct API"""
        # Skip unsupported index types
        if index_type == "SORTED":
            logger.info(
                f"[{self.workspace}] Skipping SORTED index for {field_name} (not supported in this Milvus version)"
            )
            return

        try:
            self._client.create_index(
                collection_name=self.final_namespace,
                field_name=field_name,
                index_params={"index_type": index_type},
            )
            logger.debug(
                f"[{self.workspace}] Created {field_name} index using fallback method"
            )
        except Exception as e:
            logger.info(
                f"[{self.workspace}] Could not create {field_name} index using fallback method: {e}"
            )

    def _create_indexes_after_collection(self):
        """Create indexes after collection is created"""
        # Build vector index using index configuration
        # Use compatibility helper to get IndexParams
        index_params_for_vector = self._get_index_params()

        vector_index_params = self.index_config.build_index_params(
            index_params_for_vector, field_name="vector"
        )

        # Re-raise exceptions to surface vector index creation failures
        if isinstance(vector_index_params, dict):
            self._client.create_index(
                collection_name=self.final_namespace,
                field_name=vector_index_params["field_name"],
                index_params={
                    "index_type": vector_index_params["index_type"],
                    "metric_type": vector_index_params["metric_type"],
                    "params": vector_index_params["params"],
                },
            )
        else:
            self._client.create_index(
                collection_name=self.final_namespace,
                index_params=vector_index_params,
            )

        logger.debug(
            f"[{self.workspace}] Created vector index with config: {self.index_config.to_dict()}"
        )

        # Create scalar indexes based on namespace
        # Wrap scalar index creation in try-except to allow graceful degradation
        try:
            # Try to get IndexParams in a version-compatible way
            scalar_index_params = self._get_index_params()

            if scalar_index_params is not None:
                # Create scalar indexes based on namespace
                if self.namespace.endswith("entities"):
                    # Create indexes for entity fields
                    try:
                        entity_name_index = self._get_index_params()
                        entity_name_index.add_index(
                            field_name="entity_name", index_type="INVERTED"
                        )
                        self._client.create_index(
                            collection_name=self.final_namespace,
                            index_params=entity_name_index,
                        )
                    except Exception as e:
                        logger.debug(
                            f"[{self.workspace}] IndexParams method failed for entity_name: {e}"
                        )
                        self._create_scalar_index_fallback("entity_name", "INVERTED")

                elif self.namespace.endswith("relationships"):
                    # Create indexes for relationship fields
                    try:
                        src_id_index = self._get_index_params()
                        src_id_index.add_index(
                            field_name="src_id", index_type="INVERTED"
                        )
                        self._client.create_index(
                            collection_name=self.final_namespace,
                            index_params=src_id_index,
                        )
                    except Exception as e:
                        logger.debug(
                            f"[{self.workspace}] IndexParams method failed for src_id: {e}"
                        )
                        self._create_scalar_index_fallback("src_id", "INVERTED")

                    try:
                        tgt_id_index = self._get_index_params()
                        tgt_id_index.add_index(
                            field_name="tgt_id", index_type="INVERTED"
                        )
                        self._client.create_index(
                            collection_name=self.final_namespace,
                            index_params=tgt_id_index,
                        )
                    except Exception as e:
                        logger.debug(
                            f"[{self.workspace}] IndexParams method failed for tgt_id: {e}"
                        )
                        self._create_scalar_index_fallback("tgt_id", "INVERTED")

                elif self.namespace.endswith("chunks"):
                    # Create indexes for chunk fields
                    try:
                        doc_id_index = self._get_index_params()
                        doc_id_index.add_index(
                            field_name="full_doc_id", index_type="INVERTED"
                        )
                        self._client.create_index(
                            collection_name=self.final_namespace,
                            index_params=doc_id_index,
                        )
                    except Exception as e:
                        logger.debug(
                            f"[{self.workspace}] IndexParams method failed for full_doc_id: {e}"
                        )
                        self._create_scalar_index_fallback("full_doc_id", "INVERTED")

            else:
                # Fallback to direct API calls if IndexParams is not available
                logger.info(
                    f"[{self.workspace}] IndexParams not available, using fallback methods for {self.namespace}"
                )

                # Create scalar indexes using fallback
                if self.namespace.endswith("entities"):
                    self._create_scalar_index_fallback("entity_name", "INVERTED")
                elif self.namespace.endswith("relationships"):
                    self._create_scalar_index_fallback("src_id", "INVERTED")
                    self._create_scalar_index_fallback("tgt_id", "INVERTED")
                elif self.namespace.endswith("chunks"):
                    self._create_scalar_index_fallback("full_doc_id", "INVERTED")

            logger.info(
                f"[{self.workspace}] Created indexes for collection: {self.namespace}"
            )

        except Exception as e:
            # Scalar index failures are logged as warnings (not critical)
            logger.warning(
                f"[{self.workspace}] Failed to create some scalar indexes for {self.namespace}: {e}"
            )

    def _get_required_fields_for_namespace(self) -> dict:
        """Get required core field definitions for current namespace"""

        # Base fields (common to all types)
        base_fields = {
            "id": {"type": "VarChar", "is_primary": True},
            "vector": {"type": "FloatVector"},
            "created_at": {"type": "Int64"},
        }

        # Add specific fields based on namespace
        if self.namespace.endswith("entities"):
            specific_fields = {
                "entity_name": {"type": "VarChar"},
                "content": {"type": "VarChar"},
                "source_id": {"type": "VarChar"},
                "file_path": {"type": "VarChar"},
            }
        elif self.namespace.endswith("relationships"):
            specific_fields = {
                "src_id": {"type": "VarChar"},
                "tgt_id": {"type": "VarChar"},
                "content": {"type": "VarChar"},
                "source_id": {"type": "VarChar"},
                "file_path": {"type": "VarChar"},
            }
        elif self.namespace.endswith("chunks"):
            specific_fields = {
                "full_doc_id": {"type": "VarChar"},
                "content": {"type": "VarChar"},
                "file_path": {"type": "VarChar"},
            }
        else:
            specific_fields = {
                "file_path": {"type": "VarChar"},
            }

        return {**base_fields, **specific_fields}

    def _is_field_compatible(self, existing_field: dict, expected_config: dict) -> bool:
        """Check compatibility of a single field"""
        field_name = existing_field.get("name", "unknown")
        existing_type = existing_field.get("type")
        expected_type = expected_config.get("type")

        logger.debug(
            f"[{self.workspace}] Checking field '{field_name}': existing_type={existing_type} (type={type(existing_type)}), expected_type={expected_type}"
        )

        # Convert DataType enum values to string names if needed
        original_existing_type = existing_type
        if hasattr(existing_type, "name"):
            existing_type = existing_type.name
            logger.debug(
                f"[{self.workspace}] Converted enum to name: {original_existing_type} -> {existing_type}"
            )
        elif isinstance(existing_type, int):
            # Map common Milvus internal type codes to type names for backward compatibility
            type_mapping = {
                21: "VarChar",
                101: "FloatVector",
                5: "Int64",
                9: "Double",
            }
            mapped_type = type_mapping.get(existing_type, str(existing_type))
            logger.debug(
                f"[{self.workspace}] Mapped numeric type: {existing_type} -> {mapped_type}"
            )
            existing_type = mapped_type

        # Normalize type names for comparison
        type_aliases = {
            "VARCHAR": "VarChar",
            "String": "VarChar",
            "FLOAT_VECTOR": "FloatVector",
            "INT64": "Int64",
            "BigInt": "Int64",
            "DOUBLE": "Double",
            "Float": "Double",
        }

        original_existing = existing_type
        original_expected = expected_type
        existing_type = type_aliases.get(existing_type, existing_type)
        expected_type = type_aliases.get(expected_type, expected_type)

        if original_existing != existing_type or original_expected != expected_type:
            logger.debug(
                f"[{self.workspace}] Applied aliases: {original_existing} -> {existing_type}, {original_expected} -> {expected_type}"
            )

        # Basic type compatibility check
        type_compatible = existing_type == expected_type
        logger.debug(
            f"[{self.workspace}] Type compatibility for '{field_name}': {existing_type} == {expected_type} -> {type_compatible}"
        )

        if not type_compatible:
            logger.warning(
                f"[{self.workspace}] Type mismatch for field '{field_name}': expected {expected_type}, got {existing_type}"
            )
            return False

        # Primary key check - be more flexible about primary key detection
        if expected_config.get("is_primary"):
            # Check multiple possible field names for primary key status
            is_primary = (
                existing_field.get("is_primary_key", False)
                or existing_field.get("is_primary", False)
                or existing_field.get("primary_key", False)
            )
            logger.debug(
                f"[{self.workspace}] Primary key check for '{field_name}': expected=True, actual={is_primary}"
            )
            logger.debug(
                f"[{self.workspace}] Raw field data for '{field_name}': {existing_field}"
            )

            # For ID field, be more lenient - if it's the ID field, assume it should be primary
            if field_name == "id" and not is_primary:
                logger.info(
                    f"[{self.workspace}] ID field '{field_name}' not marked as primary in existing collection, but treating as compatible"
                )
                # Don't fail for ID field primary key mismatch
            elif not is_primary:
                logger.warning(
                    f"[{self.workspace}] Primary key mismatch for field '{field_name}': expected primary key, but field is not primary"
                )
                return False

        logger.debug(f"[{self.workspace}] Field '{field_name}' is compatible")
        return True

    def _check_vector_dimension(self, collection_info: dict):
        """Check vector dimension compatibility"""
        current_dimension = self.embedding_func.embedding_dim

        # Find vector field dimension
        for field in collection_info.get("fields", []):
            if field.get("name") == "vector":
                field_type = field.get("type")

                # Extract type name from DataType enum or string
                type_name = None
                if hasattr(field_type, "name"):
                    type_name = field_type.name
                elif isinstance(field_type, str):
                    type_name = field_type
                else:
                    type_name = str(field_type)

                # Check if it's a vector type (supports multiple formats)
                if type_name in ["FloatVector", "FLOAT_VECTOR"]:
                    existing_dimension = field.get("params", {}).get("dim")

                    # Convert both to int for comparison to handle type mismatches
                    # (Milvus API may return string "1024" vs int 1024)
                    try:
                        existing_dim_int = (
                            int(existing_dimension)
                            if existing_dimension is not None
                            else None
                        )
                        current_dim_int = (
                            int(current_dimension)
                            if current_dimension is not None
                            else None
                        )
                    except (TypeError, ValueError) as e:
                        logger.error(
                            f"[{self.workspace}] Failed to parse dimensions: existing={existing_dimension} (type={type(existing_dimension)}), "
                            f"current={current_dimension} (type={type(current_dimension)}), error={e}"
                        )
                        raise ValueError(
                            f"Invalid dimension values for collection '{self.final_namespace}': "
                            f"existing={existing_dimension}, current={current_dimension}"
                        ) from e

                    if existing_dim_int != current_dim_int:
                        raise ValueError(
                            f"Vector dimension mismatch for collection '{self.final_namespace}': "
                            f"existing={existing_dim_int}, current={current_dim_int}"
                        )

                    logger.debug(
                        f"[{self.workspace}] Vector dimension check passed: {current_dim_int}"
                    )
                    return

        # If no vector field found, this might be an old collection created with simple schema
        logger.warning(
            f"[{self.workspace}] Vector field not found in collection '{self.namespace}'. This might be an old collection created with simple schema."
        )
        logger.warning(
            f"[{self.workspace}] Consider recreating the collection for optimal performance."
        )
        return

    @staticmethod
    def _has_vector_field(collection_info: dict) -> bool:
        """Return True when the collection exposes a 'vector' field.

        Old simple-schema collections may lack a vector field entirely. Their
        rows therefore carry no vector data, and copying them into the new
        schema (whose vector field is required) would fail at insert time, so
        callers use this to skip migration for such collections.
        """
        return any(
            field.get("name") == "vector" for field in collection_info.get("fields", [])
        )

    def _check_file_path_length_restriction(self, collection_info: dict) -> bool:
        """Check if collection has file_path length restrictions that need migration

        Returns:
            bool: True if migration is needed, False otherwise
        """
        existing_fields = {
            field["name"]: field for field in collection_info.get("fields", [])
        }

        # Check if file_path field exists and has length restrictions
        if "file_path" in existing_fields:
            file_path_field = existing_fields["file_path"]
            # Get max_length from field params
            max_length = file_path_field.get("params", {}).get("max_length")

            if max_length and max_length < DEFAULT_MAX_FILE_PATH_LENGTH:
                logger.info(
                    f"[{self.workspace}] Collection {self.namespace} has file_path max_length={max_length}, "
                    f"needs migration to {DEFAULT_MAX_FILE_PATH_LENGTH}"
                )
                return True

        return False

    def _check_metadata_schema_migration_needed(self, collection_info: dict) -> bool:
        existing_fields = {
            field["name"]: field for field in collection_info.get("fields", [])
        }

        for (
            field_name,
            expected_max_length,
        ) in self._get_migrated_metadata_field_limits().items():
            existing_field = existing_fields.get(field_name)
            if existing_field is None:
                logger.info(
                    f"[{self.workspace}] Collection {self.namespace} missing explicit Milvus field '{field_name}', needs migration"
                )
                return True

            if not self._is_field_compatible(existing_field, {"type": "VarChar"}):
                logger.info(
                    f"[{self.workspace}] Collection {self.namespace} has incompatible Milvus field '{field_name}', needs migration"
                )
                return True

            max_length = self._field_max_length(existing_field)
            if max_length is not None and max_length < expected_max_length:
                logger.info(
                    f"[{self.workspace}] Collection {self.namespace} has {field_name} max_length={max_length}, "
                    f"needs migration to {expected_max_length}"
                )
                return True

        return False

    def _check_schema_compatibility(self, collection_info: dict):
        """Check schema field compatibility and detect migration needs"""
        existing_fields = {
            field["name"]: field for field in collection_info.get("fields", [])
        }

        # Check if this is an old collection created with simple schema
        has_vector_field = self._has_vector_field(collection_info)

        if not has_vector_field:
            logger.warning(
                f"[{self.workspace}] Collection {self.namespace} appears to be created with old simple schema (no vector field)"
            )
            logger.warning(
                f"[{self.workspace}] This collection will work but may have suboptimal performance"
            )
            logger.warning(
                f"[{self.workspace}] Consider recreating the collection for optimal performance"
            )
            return

        if self._check_file_path_length_restriction(
            collection_info
        ) or self._check_metadata_schema_migration_needed(collection_info):
            logger.info(
                f"[{self.workspace}] Starting automatic migration for collection {self.namespace}"
            )
            self._migrate_collection_schema()
            return

        # For collections with vector field, check basic compatibility
        # Only check for critical incompatibilities, not missing optional fields
        critical_fields = {"id": {"type": "VarChar", "is_primary": True}}

        incompatible_fields = []

        for field_name, expected_config in critical_fields.items():
            if field_name in existing_fields:
                existing_field = existing_fields[field_name]
                if not self._is_field_compatible(existing_field, expected_config):
                    incompatible_fields.append(
                        f"{field_name}: expected {expected_config['type']}, "
                        f"got {existing_field.get('type')}"
                    )

        if incompatible_fields:
            raise ValueError(
                f"Critical schema incompatibility in collection '{self.final_namespace}': {incompatible_fields}"
            )

        # Get all expected fields for informational purposes
        expected_fields = self._get_required_fields_for_namespace()
        missing_fields = [
            field for field in expected_fields if field not in existing_fields
        ]

        if missing_fields:
            logger.info(
                f"[{self.workspace}] Collection {self.namespace} missing optional fields: {missing_fields}"
            )
            logger.info(
                "These fields would be available in a newly created collection for better performance"
            )

        logger.debug(
            f"[{self.workspace}] Schema compatibility check passed for {self.namespace}"
        )

    def _create_collection_with_schema(
        self, collection_name: str, ignore_index_errors: bool = False
    ) -> None:
        original_final_namespace = self.final_namespace
        try:
            self.final_namespace = collection_name
            schema = self._create_schema_for_namespace()
            self._client.create_collection(
                collection_name=collection_name, schema=schema
            )
            try:
                self._create_indexes_after_collection()
            except Exception as index_error:
                if not ignore_index_errors:
                    raise
                logger.warning(
                    f"[{self.workspace}] Failed to create indexes for new collection: {index_error}"
                )
        finally:
            self.final_namespace = original_final_namespace

    def _recover_interrupted_inplace_migration(
        self, target_collection_name: str
    ) -> str:
        """Recover from a crash/disconnect inside the in-place migration commit window.

        An in-place migration vacates the source collection (rename to _old,
        or the drop-source fallback) in Step 3 *before* promoting the temp
        collection to the target name in Step 4. A failure between those steps
        leaves the target name free but the migrated rows stranded in the temp
        collection (and the pre-migration rows, if any, in _old). Both are
        recoverable state, never scratch: dropping the temp collection here —
        as the normal Step 1 reset and the failed-attempt cleanup would —
        could destroy the only surviving copy.

        Only call this for the in-place case (source == target); a suffix
        migration never vacates its source. It is a no-op unless the target
        name is actually missing.

        Returns:
            "promoted" - temp was renamed to the target; the migration is
                         complete and the caller should stop.
            "restored" - the _old backup was renamed back to the target; the
                         source is restored and a fresh migration can re-run.
            "none"     - no interrupted commit detected; proceed normally.
        """
        if self._client.has_collection(target_collection_name):
            return "none"

        temp_collection_name = f"{target_collection_name}_temp"
        old_backup_name = f"{target_collection_name}_old"

        if self._client.has_collection(temp_collection_name):
            logger.warning(
                f"[{self.workspace}] Resuming interrupted migration: promoting "
                f"{temp_collection_name} -> {target_collection_name}"
            )
            self._client.rename_collection(temp_collection_name, target_collection_name)
            # The migrated data is now safe under the target name; the _old
            # backup (if any) is no longer the active copy, so release it.
            if self._client.has_collection(old_backup_name):
                try:
                    self._client.release_collection(old_backup_name)
                except Exception as release_error:
                    logger.warning(
                        f"[{self.workspace}] Failed to release backup collection "
                        f"{old_backup_name}: {release_error}"
                    )
            return "promoted"

        if self._client.has_collection(old_backup_name):
            logger.warning(
                f"[{self.workspace}] Resuming interrupted migration: restoring "
                f"{old_backup_name} -> {target_collection_name} (no migrated copy survived)"
            )
            self._client.rename_collection(old_backup_name, target_collection_name)
            return "restored"

        return "none"

    def _migrate_collection_schema(
        self,
        source_collection_name: str | None = None,
        target_collection_name: str | None = None,
    ):
        """Run the iterator-based migration, retrying transient connection failures.

        Each attempt is idempotent: it starts by dropping the leftover _temp
        collection and the source collection is never touched until the final
        rename, so a failed attempt can always be re-run from scratch. A
        connection-class failure leaves the current MilvusClient permanently
        dead (see _rebuild_milvus_client), so every retry rebuilds the client
        before re-running the attempt.
        """
        attempt = 0
        backoff = self._migration_retry_backoff
        needs_client_rebuild = False

        while True:
            try:
                if needs_client_rebuild:
                    self._rebuild_milvus_client()
                    needs_client_rebuild = False
                return self._migrate_collection_schema_attempt(
                    source_collection_name=source_collection_name,
                    target_collection_name=target_collection_name,
                )
            except Exception as e:
                if not self._is_retryable_connection_error(e):
                    raise
                attempt += 1
                if attempt > self._migration_max_retries:
                    logger.error(
                        f"[{self.workspace}] Migration of {self.namespace} failed after "
                        f"{attempt} attempt(s) due to connection errors"
                    )
                    raise
                needs_client_rebuild = True
                logger.warning(
                    f"[{self.workspace}] Migration attempt "
                    f"{attempt}/{self._migration_max_retries + 1} for {self.namespace} "
                    f"failed with a connection error: {e}. "
                    f"Rebuilding Milvus client and retrying in {backoff:.0f}s"
                )
                time.sleep(backoff)
                backoff = min(
                    backoff * MILVUS_MIGRATION_RETRY_BACKOFF_MULTIPLIER,
                    self._migration_retry_max_backoff,
                )

    def _migrate_collection_schema_attempt(
        self,
        source_collection_name: str | None = None,
        target_collection_name: str | None = None,
    ):
        source_collection_name = source_collection_name or self.final_namespace
        target_collection_name = target_collection_name or self.final_namespace
        temp_collection_name = f"{target_collection_name}_temp"
        original_final_namespace = self.final_namespace
        is_inplace = source_collection_name == target_collection_name
        iterator = None
        # Once an in-place migration has vacated its source (Step 3), the temp
        # collection holds the only migrated copy and must not be dropped by
        # the failure cleanup below — a later attempt or startup recovers it.
        commit_phase = False
        # True once we explicitly loaded a suffix migration's legacy source, so
        # the failure cleanup can release it again (in-place sources are the
        # active collection and are left as-is).
        source_loaded = False

        try:
            logger.info(
                f"[{self.workspace}] Starting iterator-based schema migration for {self.namespace}: "
                f"{source_collection_name} -> {target_collection_name}"
            )

            # A previous attempt may have failed inside the in-place commit
            # window (after the source was vacated). Finish that commit instead
            # of restarting the copy, which would drop the recovery copy.
            if is_inplace:
                recovery = self._recover_interrupted_inplace_migration(
                    target_collection_name
                )
                if recovery == "promoted":
                    self.final_namespace = target_collection_name
                    return
                # "restored": the source is back, fall through to a clean copy.
                # "none": nothing to recover, proceed normally.

            logger.info(
                f"[{self.workspace}] Step 1: Creating temporary collection: {temp_collection_name}"
            )
            if self._client.has_collection(temp_collection_name):
                self._client.drop_collection(temp_collection_name)
            self._create_collection_with_schema(
                temp_collection_name, ignore_index_errors=True
            )

            # The temp collection is deliberately NOT loaded here: insert does
            # not require a loaded collection, and loading it would keep every
            # migrated row's growing segments in query-node memory for the
            # whole bulk copy (nearly doubling the collection's footprint on
            # the server). The final collection is loaded after the rename via
            # _ensure_collection_loaded.

            logger.info(
                f"[{self.workspace}] Step 2: Copying data using query_iterator from: {source_collection_name}"
            )

            # query_iterator issues a server-side query, which requires the
            # source collection to be loaded. An in-place source is the active
            # collection and is already loaded, but a legacy/suffix source is
            # typically NOT loaded (it is an old backup), so the iterator setup
            # fails with "collection not loaded" (code 101). Load it explicitly
            # (idempotent); on a successful migration the source becomes the
            # backup and is released again from memory at the end.
            self._client.load_collection(source_collection_name)
            # Only track suffix loads: an in-place source is the active
            # collection and must not be released on failure.
            source_loaded = not is_inplace

            try:
                iterator = self._client.query_iterator(
                    collection_name=source_collection_name,
                    batch_size=self._migration_iterator_batch_size,
                    output_fields=["*"],
                )
                logger.debug(f"[{self.workspace}] Query iterator created successfully")
            except Exception as iterator_error:
                logger.error(
                    f"[{self.workspace}] Failed to create query iterator: {iterator_error}"
                )
                raise

            total_migrated = 0
            batch_number = 1
            rows_since_flush = 0

            while True:
                try:
                    batch_data = iterator.next()
                    if not batch_data:
                        # No more data available
                        break

                    sanitized_batch_data = [
                        self._normalize_migration_row(row) for row in batch_data
                    ]
                    insert_batches = self._build_upsert_batches(
                        sanitized_batch_data,
                        max_payload_bytes=self._max_upsert_payload_bytes,
                        max_records_per_batch=self._max_upsert_records_per_batch,
                    )
                    try:
                        for insert_batch_number, (
                            records_batch,
                            estimated_bytes,
                        ) in enumerate(insert_batches, 1):
                            logger.debug(
                                f"[{self.workspace}] Milvus migration insert batch "
                                f"{batch_number}.{insert_batch_number}/{len(insert_batches)}: "
                                f"records={len(records_batch)}, estimated_payload_bytes={estimated_bytes}"
                            )
                            self._client.insert(
                                collection_name=temp_collection_name,
                                data=records_batch,
                            )
                        total_migrated += len(batch_data)
                        rows_since_flush += len(batch_data)

                        logger.info(
                            f"[{self.workspace}] Iterator batch {batch_number}: "
                            f"processed {len(batch_data)} records, total migrated: {total_migrated}"
                        )
                        batch_number += 1

                        # Seal the accumulated insert buffer to bound the
                        # data-node memory and block until the flush returns,
                        # giving the server back-pressure during a large copy.
                        if (
                            self._migration_flush_interval_rows > 0
                            and rows_since_flush >= self._migration_flush_interval_rows
                        ):
                            logger.info(
                                f"[{self.workspace}] Flushing temp collection after "
                                f"{rows_since_flush} rows (total migrated: {total_migrated})"
                            )
                            self._client.flush(temp_collection_name)
                            rows_since_flush = 0

                        # Optional throttle to let a small server breathe.
                        if self._migration_batch_sleep > 0:
                            time.sleep(self._migration_batch_sleep)

                    except Exception as batch_error:
                        logger.error(
                            f"[{self.workspace}] Failed to insert iterator batch {batch_number}: {batch_error}"
                        )
                        raise

                except Exception as next_error:
                    logger.error(
                        f"[{self.workspace}] Iterator next() failed at batch {batch_number}: {next_error}"
                    )
                    raise

            # Final flush so the last unsealed insert buffer is persisted and
            # released before the collection is promoted and loaded.
            if self._migration_flush_interval_rows > 0 and total_migrated > 0:
                logger.info(
                    f"[{self.workspace}] Final flush of temp collection ({total_migrated} rows migrated)"
                )
                self._client.flush(temp_collection_name)

            if total_migrated > 0:
                logger.info(
                    f"[{self.workspace}] Successfully migrated {total_migrated} records using iterator"
                )
            else:
                logger.info(
                    f"[{self.workspace}] No data found in original collection, migration completed"
                )

            backup_collection_name: str | None = None
            if is_inplace:
                logger.info(
                    f"[{self.workspace}] Step 3: Rename origin collection to {source_collection_name}_old"
                )
                # Entering the commit window: from here the source is vacated,
                # so the temp collection becomes the recovery copy and must
                # survive a failure rather than being cleaned up as scratch.
                commit_phase = True
                old_backup_name = f"{source_collection_name}_old"
                try:
                    # Drop a stale backup from a previous migration first:
                    # rename_collection cannot overwrite an existing name, and
                    # a failed rename here falls back to dropping the source
                    # collection outright (losing the backup entirely).
                    if self._client.has_collection(old_backup_name):
                        logger.info(
                            f"[{self.workspace}] Dropping stale backup collection {old_backup_name}"
                        )
                        self._client.drop_collection(old_backup_name)
                    self._client.rename_collection(
                        source_collection_name, old_backup_name
                    )
                    backup_collection_name = old_backup_name
                except Exception as rename_error:
                    try:
                        logger.warning(
                            f"[{self.workspace}] Try to drop origin collection instead"
                        )
                        self._client.drop_collection(source_collection_name)
                    except Exception as e:
                        logger.error(
                            f"[{self.workspace}] Rename operation failed: {rename_error}"
                        )
                        raise e
            elif self._client.has_collection(target_collection_name):
                raise RuntimeError(
                    f"Target collection already exists: {target_collection_name}"
                )
            else:
                # Suffix migration: the legacy source collection stays in
                # place as the backup.
                backup_collection_name = source_collection_name

            logger.info(
                f"[{self.workspace}] Step 4: Renaming collection {temp_collection_name} -> {target_collection_name}"
            )
            try:
                self._client.rename_collection(
                    temp_collection_name, target_collection_name
                )
                logger.info(f"[{self.workspace}] Rename operation completed")
            except Exception as rename_error:
                if source_collection_name == target_collection_name:
                    logger.error(
                        f"[{self.workspace}] Rename operation failed: {rename_error}"
                    )
                else:
                    logger.error(
                        f"[{self.workspace}] Target rename operation failed: {rename_error}"
                    )
                raise RuntimeError(
                    f"Failed to rename collection: {rename_error}"
                ) from rename_error

            self.final_namespace = target_collection_name

            # The backup collection (legacy source or renamed _old) inherits
            # the loaded state of the pre-migration active collection, which
            # would keep a full second copy of the data in query-node memory
            # forever. Release it so the backup only occupies disk.
            if backup_collection_name is not None:
                try:
                    self._client.release_collection(backup_collection_name)
                    logger.info(
                        f"[{self.workspace}] Released backup collection {backup_collection_name} from memory"
                    )
                except Exception as release_error:
                    logger.warning(
                        f"[{self.workspace}] Failed to release backup collection "
                        f"{backup_collection_name}: {release_error}"
                    )

        except Exception as e:
            self.final_namespace = original_final_namespace
            logger.error(
                f"[{self.workspace}] Iterator-based migration failed for {self.namespace}: {e}"
            )

            if commit_phase:
                # The source has already been vacated, so the temp collection
                # is the only migrated copy: keep it as recovery state. The
                # next attempt (or startup) finishes the interrupted commit via
                # _recover_interrupted_inplace_migration.
                logger.warning(
                    f"[{self.workspace}] Migration failed after the source was vacated; "
                    f"keeping {temp_collection_name} as recovery state for the next attempt or startup"
                )
            else:
                try:
                    if self._client and self._client.has_collection(
                        temp_collection_name
                    ):
                        logger.info(
                            f"[{self.workspace}] Cleaning up failed migration temporary collection"
                        )
                        self._client.drop_collection(temp_collection_name)
                except Exception as cleanup_error:
                    logger.warning(
                        f"[{self.workspace}] Failed to cleanup temporary collection: {cleanup_error}. "
                        f"The leftover temp collection will be dropped on the next migration attempt or startup"
                    )

            # Release the suffix source we loaded above so a failed startup or
            # retry does not leave a full backup resident in query-node memory.
            # In-place sources are the active collection and are left as-is.
            if source_loaded:
                try:
                    self._client.release_collection(source_collection_name)
                    logger.info(
                        f"[{self.workspace}] Released source collection "
                        f"{source_collection_name} from memory after failed migration"
                    )
                except Exception as release_error:
                    logger.warning(
                        f"[{self.workspace}] Failed to release source collection "
                        f"{source_collection_name}: {release_error}"
                    )

            raise RuntimeError(
                f"Iterator-based migration failed for collection {self.namespace}: {e}"
            ) from e

        finally:
            if iterator:
                try:
                    iterator.close()
                    logger.debug(
                        f"[{self.workspace}] Query iterator closed successfully"
                    )
                except Exception as close_error:
                    logger.warning(
                        f"[{self.workspace}] Failed to close query iterator: {close_error}"
                    )

    def _validate_collection_compatibility(self):
        """Validate existing collection's dimension and schema compatibility"""
        try:
            collection_info = self._client.describe_collection(self.final_namespace)

            # 1. Check vector dimension
            self._check_vector_dimension(collection_info)

            # 2. Check schema compatibility
            self._check_schema_compatibility(collection_info)

            logger.info(
                f"[{self.workspace}] VectorDB Collection '{self.namespace}' compatibility validation passed"
            )

        except Exception as e:
            logger.error(
                f"[{self.workspace}] Collection compatibility validation failed for {self.namespace}: {e}"
            )
            raise

    def _validate_collection_and_load(self) -> None:
        try:
            self._client.describe_collection(self.final_namespace)
            self._validate_collection_compatibility()
        except Exception as validation_error:
            logger.error(
                f"[{self.workspace}] CRITICAL ERROR: Collection '{self.namespace}' exists but validation failed!"
            )
            logger.error(
                f"[{self.workspace}] This indicates potential data migration failure or schema incompatibility."
            )
            logger.error(f"[{self.workspace}] Validation error: {validation_error}")
            logger.error(f"[{self.workspace}] MANUAL INTERVENTION REQUIRED:")
            logger.error(
                f"[{self.workspace}] 1. Check the existing collection schema and data integrity"
            )
            logger.error(f"[{self.workspace}] 2. Backup existing data if needed")
            logger.error(
                f"[{self.workspace}] 3. Manually resolve schema compatibility issues"
            )
            logger.error(
                f"[{self.workspace}] 4. Consider recreating the collection using: lightrag-rebuild-vdb "
            )
            logger.error(
                f"[{self.workspace}] Program execution stopped to prevent potential data loss."
            )
            raise RuntimeError(
                f"Collection validation failed for '{self.final_namespace}'. "
                f"Data migration failure detected. Manual intervention required to prevent data loss. "
                f"Original error: {validation_error}"
            )

        try:
            self._ensure_collection_loaded()
        except Exception as load_error:
            if not self._is_missing_vector_index_error(load_error):
                raise

            try:
                self._repair_missing_vector_index()
                self._ensure_collection_loaded()
                logger.info(
                    f"[{self.workspace}] Repaired missing vector index for existing collection '{self.namespace}'"
                )
            except Exception as repair_error:
                raise RuntimeError(
                    f"Index repair failed for collection '{self.final_namespace}'. "
                    f"Original error: {repair_error}"
                ) from repair_error

    @staticmethod
    def _is_missing_vector_index_error(error: Exception) -> bool:
        """Return True when the error indicates the collection lacks a vector index."""
        error_message = str(error).lower()
        return (
            "no vector index" in error_message
            or "please create index firstly" in error_message
        )

    def _repair_missing_vector_index(self):
        """Create indexes for an existing collection that is missing its vector index."""
        logger.warning(
            f"[{self.workspace}] Collection '{self.namespace}' is missing a vector index, attempting repair"
        )
        self._create_indexes_after_collection()

    def _ensure_collection_loaded(self):
        """Ensure the collection is loaded into memory for search operations"""
        try:
            # Check if collection exists first
            if not self._client.has_collection(self.final_namespace):
                logger.error(
                    f"[{self.workspace}] Collection {self.namespace} does not exist"
                )
                raise ValueError(f"Collection {self.final_namespace} does not exist")

            # Load the collection if it's not already loaded
            # In Milvus, collections need to be loaded before they can be searched
            self._client.load_collection(self.final_namespace)
            # logger.debug(f"[{self.workspace}] Collection {self.namespace} loaded successfully")

        except Exception as e:
            logger.error(
                f"[{self.workspace}] Failed to load collection {self.namespace}: {e}"
            )
            raise

    def _create_collection_if_not_exist(self):
        """Create collection if not exists and check existing collection compatibility"""

        try:
            collection_exists = self._client.has_collection(self.final_namespace)
            logger.info(
                f"[{self.workspace}] VectorDB collection '{self.namespace}' exists check: {collection_exists}"
            )

            if collection_exists:
                self._validate_collection_and_load()
                return

            legacy_collection_exists = (
                self.legacy_namespace != self.final_namespace
                and self._client.has_collection(self.legacy_namespace)
            )

            # Interrupted in-place commit recovery — MUST run before the legacy
            # migration below. An in-place migration vacates its source (Step 3)
            # before promoting the temp collection to the target name (Step 4);
            # a crash in that window leaves the target name free with the
            # completed copy stranded in {final}_temp and the pre-migration data
            # in {final}_old. If a legacy migration ran first it would migrate
            # the stale legacy collection over the target and drop {final}_temp
            # as scratch (Step 1), silently losing every write made to the
            # suffixed collection since the legacy split.
            #
            # {final}_old is the reliable marker that an in-place copy COMPLETED
            # (it only ever exists once Step 3 renamed the source): when it is
            # present, recover unconditionally. A lone {final}_temp next to a
            # live legacy source is instead an aborted *partial* suffix copy and
            # must be treated as scratch by the legacy migration — so only
            # recover a lone temp when there is no legacy source to fall back to.
            temp_collection_name = f"{self.final_namespace}_temp"
            old_backup_name = f"{self.final_namespace}_old"
            has_old_backup = self._client.has_collection(old_backup_name)
            has_temp = self._client.has_collection(temp_collection_name)
            if has_old_backup or (has_temp and not legacy_collection_exists):
                recovery = self._recover_interrupted_inplace_migration(
                    self.final_namespace
                )
                if recovery in ("promoted", "restored"):
                    self._ensure_collection_loaded()
                    return

            if legacy_collection_exists:
                legacy_collection_info = self._client.describe_collection(
                    self.legacy_namespace
                )
                if not self._has_vector_field(legacy_collection_info):
                    # Old simple-schema collection with no vector field: its rows
                    # carry no vectors, so migrating them into the required-vector
                    # schema would fail at insert and block startup. Skip the
                    # migration and create a fresh suffixed collection instead.
                    logger.warning(
                        f"[{self.workspace}] Legacy collection '{self.legacy_namespace}' "
                        f"has no vector field (old simple schema); cannot migrate its rows "
                        f"into '{self.final_namespace}'. Creating a new collection instead."
                    )
                else:
                    try:
                        self._check_vector_dimension(legacy_collection_info)
                    except ValueError as legacy_error:
                        logger.warning(
                            f"[{self.workspace}] Legacy collection '{self.legacy_namespace}' "
                            f"is not compatible with '{self.final_namespace}': {legacy_error}. "
                            f"Creating a new collection without migrating legacy vectors."
                        )
                    else:
                        self._migrate_collection_schema(
                            source_collection_name=self.legacy_namespace,
                            target_collection_name=self.final_namespace,
                        )
                        self._ensure_collection_loaded()
                        return

            # Collection doesn't exist, create new collection
            logger.info(f"[{self.workspace}] Creating new collection: {self.namespace}")
            self._create_collection_with_schema(self.final_namespace)
            self._ensure_collection_loaded()

            logger.info(
                f"[{self.workspace}] Successfully created Milvus collection: {self.namespace}"
            )

        except RuntimeError:
            # Re-raise RuntimeError (validation failures) without modification
            # These are critical errors that should stop execution
            raise

        except Exception as e:
            # A transient connection failure must never trigger the
            # force-create fallback below: dropping and recreating the
            # collection on a flaky connection would destroy healthy data
            # (or create an empty suffixed collection that permanently
            # shadows an unmigrated legacy collection). Let it propagate so
            # initialize() fails and can be retried.
            if self._is_retryable_connection_error(e):
                logger.error(
                    f"[{self.workspace}] Connection error in _create_collection_if_not_exist "
                    f"for {self.namespace}: {e}"
                )
                raise

            logger.error(
                f"[{self.workspace}] Error in _create_collection_if_not_exist for {self.namespace}: {e}"
            )

            # If there's any error (other than validation failure), try to force create the collection
            logger.info(
                f"[{self.workspace}] Attempting to force create collection {self.namespace}..."
            )
            try:
                # Try to drop the collection first if it exists in a bad state
                try:
                    if self._client.has_collection(self.final_namespace):
                        logger.info(
                            f"[{self.workspace}] Dropping potentially corrupted collection {self.namespace}"
                        )
                        self._client.drop_collection(self.final_namespace)
                except Exception as drop_error:
                    logger.warning(
                        f"[{self.workspace}] Could not drop collection {self.namespace}: {drop_error}"
                    )

                # Create fresh collection
                self._create_collection_with_schema(self.final_namespace)

                # Load the newly created collection
                self._ensure_collection_loaded()

                logger.info(
                    f"[{self.workspace}] Successfully force-created collection {self.namespace}"
                )

            except Exception as create_error:
                logger.error(
                    f"[{self.workspace}] Failed to force-create collection {self.namespace}: {create_error}"
                )
                raise

    def __post_init__(self):
        validate_workspace(self.workspace)
        self._validate_embedding_func()

        # Extract MilvusIndexConfig parameters from vector_db_storage_cls_kwargs
        #
        # IMPORTANT: This approach allows Milvus index configuration via vector_db_storage_cls_kwargs,
        # which is the RECOMMENDED method for framework integration (e.g., RAGAnything).
        #
        # All 11 index configuration parameters can be passed through vector_db_storage_cls_kwargs:
        #   - 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
        #
        # Example:
        #   LightRAG(
        #       vector_storage="MilvusVectorDBStorage",
        #       vector_db_storage_cls_kwargs={
        #           "cosine_better_than_threshold": 0.2,
        #           "index_type": "HNSW",
        #           "metric_type": "COSINE",
        #           "hnsw_m": 32,
        #           "hnsw_ef_construction": 256,
        #       }
        #   )
        #
        # Use MilvusIndexConfig.get_config_field_names() to dynamically extract valid parameters.
        # This ensures we always stay in sync with the MilvusIndexConfig dataclass definition.
        kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {})
        index_config_keys = MilvusIndexConfig.get_config_field_names()
        index_config_params = {
            k: v for k, v in kwargs.items() if k in index_config_keys
        }

        # Initialize index configuration (if not already set)
        # Configuration priority: init params from kwargs > environment variables > defaults
        if not hasattr(self, "index_config") or self.index_config is None:
            self.index_config = MilvusIndexConfig(**index_config_params)

        # Check for MILVUS_WORKSPACE environment variable first (higher priority)
        # This allows administrators to force a specific workspace for all Milvus storage instances
        milvus_workspace = os.environ.get("MILVUS_WORKSPACE")
        if milvus_workspace and milvus_workspace.strip():
            # Use environment variable value, overriding the passed workspace parameter
            effective_workspace = milvus_workspace.strip()
            logger.info(
                f"Using MILVUS_WORKSPACE environment variable: '{effective_workspace}' (overriding '{self.workspace}/{self.namespace}')"
            )
        else:
            # Use the workspace parameter passed during initialization
            effective_workspace = self.workspace
            if effective_workspace:
                logger.debug(
                    f"Using passed workspace parameter: '{effective_workspace}'"
                )

        self.workspace = effective_workspace or ""
        self.model_suffix = self._generate_collection_suffix()
        if self.workspace:
            self.legacy_namespace = f"{self.workspace}_{self.namespace}"
            logger.debug(
                f"Legacy namespace with workspace prefix: '{self.legacy_namespace}'"
            )
        else:
            self.legacy_namespace = self.namespace
            logger.debug(f"Legacy namespace (no workspace): '{self.legacy_namespace}'")
        if self.model_suffix:
            self.final_namespace = f"{self.legacy_namespace}_{self.model_suffix}"
            logger.info(f"Milvus collection: {self.final_namespace}")
        else:
            self.final_namespace = self.legacy_namespace
            logger.warning(
                f"Milvus collection: {self.final_namespace} missing suffix. Please add model_name to embedding_func for proper model-based data isolation."
            )
        cosine_threshold = kwargs.get("cosine_better_than_threshold")
        if cosine_threshold is None:
            raise ValueError(
                "cosine_better_than_threshold must be specified in vector_db_storage_cls_kwargs"
            )
        self.cosine_better_than_threshold = cosine_threshold

        # Ensure created_at is in meta_fields
        if "created_at" not in self.meta_fields:
            self.meta_fields.add("created_at")
        self._varchar_field_limits = self._get_varchar_field_limits_for_namespace()

        # Initialize client as None - will be created in initialize() method
        self._client = None
        self._max_batch_size = self.global_config["embedding_batch_num"]

        # Flush-time batching limits (see module-level DEFAULT_MILVUS_* constants).
        # A non-positive value disables that splitting dimension.
        self._max_upsert_payload_bytes = int(
            os.getenv(
                "MILVUS_UPSERT_MAX_PAYLOAD_BYTES",
                str(DEFAULT_MILVUS_UPSERT_MAX_PAYLOAD_BYTES),
            )
        )
        self._max_upsert_records_per_batch = int(
            os.getenv(
                "MILVUS_UPSERT_MAX_RECORDS_PER_BATCH",
                str(DEFAULT_MILVUS_UPSERT_MAX_RECORDS_PER_BATCH),
            )
        )
        self._max_delete_records_per_batch = int(
            os.getenv(
                "MILVUS_DELETE_MAX_RECORDS_PER_BATCH",
                str(DEFAULT_MILVUS_DELETE_MAX_RECORDS_PER_BATCH),
            )
        )
        if self._max_upsert_payload_bytes <= 0:
            logger.warning(
                f"MILVUS_UPSERT_MAX_PAYLOAD_BYTES={self._max_upsert_payload_bytes} is non-positive, disable payload-size splitting"
            )
        if self._max_upsert_records_per_batch <= 0:
            logger.warning(
                f"MILVUS_UPSERT_MAX_RECORDS_PER_BATCH={self._max_upsert_records_per_batch} is non-positive, disable upsert record-count splitting"
            )
        if self._max_delete_records_per_batch <= 0:
            logger.warning(
                f"MILVUS_DELETE_MAX_RECORDS_PER_BATCH={self._max_delete_records_per_batch} is non-positive, disable delete record-count splitting"
            )

        # Schema-migration retry knobs (see DEFAULT_MILVUS_MIGRATION_*
        # constants). MILVUS_MIGRATION_MAX_RETRIES=0 restores fail-fast.
        self._migration_max_retries = max(
            0,
            _get_env_int(
                "MILVUS_MIGRATION_MAX_RETRIES", DEFAULT_MILVUS_MIGRATION_MAX_RETRIES
            ),
        )
        self._migration_retry_backoff = float(
            os.getenv(
                "MILVUS_MIGRATION_RETRY_BACKOFF",
                str(DEFAULT_MILVUS_MIGRATION_RETRY_BACKOFF_SECONDS),
            )
        )
        self._migration_retry_max_backoff = float(
            os.getenv(
                "MILVUS_MIGRATION_RETRY_MAX_BACKOFF",
                str(DEFAULT_MILVUS_MIGRATION_RETRY_MAX_BACKOFF_SECONDS),
            )
        )
        self._migration_iterator_batch_size = _get_env_int(
            "MILVUS_MIGRATION_ITERATOR_BATCH_SIZE",
            DEFAULT_MILVUS_MIGRATION_ITERATOR_BATCH_SIZE,
        )

        # Schema-migration memory back-pressure knobs (see the
        # DEFAULT_MILVUS_MIGRATION_FLUSH_*/_BATCH_SLEEP_* constants).
        self._migration_flush_interval_rows = max(
            0,
            _get_env_int(
                "MILVUS_MIGRATION_FLUSH_INTERVAL_ROWS",
                DEFAULT_MILVUS_MIGRATION_FLUSH_INTERVAL_ROWS,
            ),
        )
        self._migration_batch_sleep = max(
            0.0,
            float(
                os.getenv(
                    "MILVUS_MIGRATION_BATCH_SLEEP",
                    str(DEFAULT_MILVUS_MIGRATION_BATCH_SLEEP_SECONDS),
                )
            ),
        )
        self._initialized = False

        # Deferred-embedding buffers and the per-namespace flush lock.
        # The lock keys on final_namespace so two instances pointing at the
        # same Milvus collection (e.g. when MILVUS_WORKSPACE env override is
        # used) share a single writer lock. We construct it here in
        # __post_init__ — not in initialize() — so any code path that
        # touches the buffer before initialize() still has a valid lock.
        self._pending_vector_docs: dict[str, _PendingVectorDoc] = {}
        self._pending_vector_deletes: set[str] = set()
        self._flush_lock = get_namespace_lock(
            namespace=self.final_namespace, workspace=""
        )

    async def initialize(self):
        """Initialize Milvus collection"""
        async with get_data_init_lock():
            if self._initialized:
                return

            try:
                # Create MilvusClient if not already created
                if self._client is None:
                    self._client = self._create_milvus_client()
                    logger.debug(
                        f"[{self.workspace}] MilvusClient created successfully"
                    )

                # Validate Milvus version compatibility with configured index
                if self.index_config.index_type in INDEX_VERSION_REQUIREMENTS:
                    try:
                        server_version = self._client.get_server_version()
                        self.index_config.validate_milvus_version(server_version)
                    except Exception as version_error:
                        logger.error(
                            f"[{self.workspace}] Milvus version validation failed: {version_error}"
                        )
                        raise

                # Create collection and check compatibility
                self._create_collection_if_not_exist()
                self._initialized = True
                logger.info(
                    f"[{self.workspace}] Milvus collection '{self.namespace}' initialized successfully"
                )
            except Exception as e:
                logger.error(
                    f"[{self.workspace}] Failed to initialize Milvus collection '{self.namespace}': {e}"
                )
                raise

    async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
        """Buffer vector docs for embedding and batched flush.

        Embedding deliberately does NOT happen here: repeated upserts of the
        same id, or many small batches, collapse into a single flush-time
        embedding pass. Reads (`get_by_id`/`get_by_ids`/`get_vectors_by_ids`)
        observe pending docs via the same lock for read-your-writes.
        """
        if not data:
            return

        current_time = int(time.time())

        pending_docs: list[tuple[str, _PendingVectorDoc]] = []
        for i, (k, v) in enumerate(data.items(), start=1):
            # _sanitize_varchar_fields already byte-truncates the stored
            # `content` when it is a meta field; the pending doc keeps the full
            # untruncated text so the embedding sees the complete chunk.
            source = self._sanitize_varchar_fields(
                {
                    "id": k,
                    "created_at": current_time,
                    **{k1: v1 for k1, v1 in v.items() if k1 in self.meta_fields},
                }
            )
            pending_docs.append(
                (
                    k,
                    _PendingVectorDoc(source=source, content=v["content"]),
                )
            )
            await _cooperative_yield(i)

        # An upsert overrides any pending delete on the same id; installing
        # a fresh _PendingVectorDoc instance invalidates any vector cached
        # by a prior get_vectors_by_ids() call on a stale revision.
        async with self._flush_lock:
            for doc_id, pdoc in pending_docs:
                self._pending_vector_deletes.discard(doc_id)
                self._pending_vector_docs[doc_id] = pdoc

    async def query(
        self, query: str, top_k: int, query_embedding: list[float] = None
    ) -> list[dict[str, Any]]:
        """Similarity search against the persisted Milvus collection.

        Note: buffered-but-unflushed upserts are NOT visible to this method —
        they exist only in `_pending_vector_docs` until `index_done_callback()`
        embeds and writes them. Callers that need read-after-write visibility
        for similarity search must run an explicit flush first.
        """
        # Ensure collection is loaded before querying
        self._ensure_collection_loaded()

        # Use provided embedding or compute it
        if query_embedding is not None:
            embedding = [query_embedding]  # Milvus expects a list of embeddings
        else:
            embedding = await self.embedding_func(
                [query], context="query", _priority=DEFAULT_QUERY_PRIORITY
            )  # higher priority for query

        # Include all meta_fields (created_at is now always included)
        output_fields = list(self.meta_fields)

        # Build search params from index config
        search_params_base = self.index_config.build_search_params()

        # Merge with metric type and radius threshold
        search_params = {
            "metric_type": self.index_config.metric_type,
            "params": {
                **search_params_base.get("params", {}),
                "radius": self.cosine_better_than_threshold,
            },
        }

        results = self._client.search(
            collection_name=self.final_namespace,
            data=embedding,
            limit=top_k,
            output_fields=output_fields,
            search_params=search_params,
        )
        return [
            {
                **dp["entity"],
                "id": dp["id"],
                "distance": dp["distance"],
                "created_at": dp.get("created_at"),
            }
            for dp in results[0]
        ]

    @staticmethod
    def _build_upsert_batches(
        records: list[dict[str, Any]],
        max_payload_bytes: int,
        max_records_per_batch: int,
    ) -> list[tuple[list[dict[str, Any]], int]]:
        """Split upsert records into batches by estimated payload size and count.

        The byte budget is the primary limiter: records accumulate until adding
        the next one would exceed ``max_payload_bytes``, then a new batch starts.
        Size is estimated by JSON-serializing each record; this overestimates the
        actual gRPC protobuf size (a JSON float string is far longer than the 4
        protobuf bytes it encodes), so the split stays conservatively below the
        server limit and never underestimates.

        A single record larger than the byte budget is emitted as its own batch
        rather than raising: JSON overestimation means such a record's real
        protobuf size is often still under Milvus' 64MB ceiling, so we let the
        server be the final arbiter instead of failing client-side. Returns a
        list of ``(batch, estimated_bytes)`` tuples (estimate used for logging).
        """
        if not records:
            return []

        payload_limit = max_payload_bytes if max_payload_bytes > 0 else float("inf")
        records_limit = (
            max_records_per_batch if max_records_per_batch > 0 else float("inf")
        )

        batches: list[tuple[list[dict[str, Any]], int]] = []
        current_batch: list[dict[str, Any]] = []
        # JSON array overhead ("[]")
        current_estimated_bytes = 2

        for record in records:
            record_size = len(
                json.dumps(
                    record,
                    ensure_ascii=False,
                    separators=(",", ":"),
                    default=str,
                ).encode("utf-8")
            )

            # If current batch not empty, a comma is needed before next element.
            separator_overhead = 1 if current_batch else 0
            next_batch_size = current_estimated_bytes + separator_overhead + record_size

            if current_batch and (
                len(current_batch) >= records_limit or next_batch_size > payload_limit
            ):
                batches.append((current_batch, current_estimated_bytes))
                current_batch = []
                current_estimated_bytes = 2
                next_batch_size = current_estimated_bytes + record_size

            current_batch.append(record)
            current_estimated_bytes = next_batch_size

        if current_batch:
            batches.append((current_batch, current_estimated_bytes))

        return batches

    async def index_done_callback(self) -> None:
        """Flush all buffered vector ops to Milvus before returning.

        Contract: on a successful return, every previously buffered upsert
        has been embedded and committed to the collection, and every buffered
        delete has been issued — i.e. all pending vectors are durable in
        Milvus (which persists automatically once written). On any embed-
        or server-side failure this method raises and leaves both buffers
        intact for the next callback to retry; the caller MUST NOT assume
        clean persistence in that case.
        """
        await self._flush_pending_vector_ops()

    async def drop_pending_index_ops(self) -> None:
        """Discard buffered upserts/deletes (pipeline aborting on error)."""
        async with self._flush_lock:
            self._pending_vector_docs.clear()
            self._pending_vector_deletes.clear()

    async def _flush_pending_vector_ops(self) -> None:
        """Flush buffered vector upserts and deletes to Milvus.

        Embedding runs *inside* this lock (not in `upsert` or lock-free):
        it makes deferred embedding and bulk indexing atomic against
        concurrent upserts and destructive mutations. Any failure (embed
        or server write) raises and leaves both buffers intact; the next
        `index_done_callback` retries automatically.
        """
        async with self._flush_lock:
            if not self._pending_vector_docs and not self._pending_vector_deletes:
                return
            if self._client is None:
                return

            # Milvus requires the collection to be loaded before upsert/delete.
            self._ensure_collection_loaded()

            pending_docs = self._pending_vector_docs
            pending_deletes = self._pending_vector_deletes

            docs_to_embed: list[tuple[str, _PendingVectorDoc]] = [
                (doc_id, pdoc)
                for doc_id, pdoc in pending_docs.items()
                if pdoc.vector is None
            ]

            if docs_to_embed:
                contents = [pdoc.content for _, pdoc in docs_to_embed]
                batches = [
                    contents[i : i + self._max_batch_size]
                    for i in range(0, len(contents), self._max_batch_size)
                ]
                logger.info(
                    f"[{self.workspace}] {self.namespace} flush: embedding "
                    f"{len(docs_to_embed)} vectors in {len(batches)} batch(es) "
                    f"(batch_num={self._max_batch_size})"
                )
                try:
                    embeddings_list = await asyncio.gather(
                        *[
                            self.embedding_func(batch, context="document")
                            for batch in batches
                        ]
                    )
                except Exception as e:
                    logger.error(
                        f"[{self.workspace}] Error embedding pending vector ops "
                        f"(upserts={len(docs_to_embed)}): {e}"
                    )
                    raise

                embeddings = np.concatenate(embeddings_list)
                if len(embeddings) != len(docs_to_embed):
                    raise RuntimeError(
                        f"[{self.workspace}] Embedding count mismatch: expected "
                        f"{len(docs_to_embed)}, got {len(embeddings)}"
                    )
                for i, ((_, pdoc), embedding) in enumerate(
                    zip(docs_to_embed, embeddings), start=1
                ):
                    # Cache as float32 so a second flush after a server-side
                    # error doesn't re-embed, and so the upsert JSON payload
                    # stays compact (float32 serializes to a shorter string
                    # than float64, and Milvus stores FLOAT_VECTOR as float32
                    # anyway, so the cast is lossless).
                    pdoc.vector = np.array(embedding, dtype=np.float32).tolist()
                    await _cooperative_yield(i)

            # Assemble final upsert payload. After the embed loop above every
            # pending doc has a non-None vector (count-mismatch was checked),
            # so we can iterate without re-guarding.
            committed_ids: list[str] = list(pending_docs.keys())
            # source was already byte-truncated in upsert(); no need to
            # re-sanitize here (vector is not a VarChar field).
            list_data: list[dict[str, Any]] = [
                {
                    **pending_docs[doc_id].source,
                    "vector": pending_docs[doc_id].vector,
                }
                for doc_id in committed_ids
            ]

            try:
                if list_data:
                    # Split the upsert into batches that stay under the server-side
                    # 64MB gRPC message limit. Fail-fast: any batch failure raises
                    # immediately and the full buffer is retained for the next flush.
                    upsert_batches = self._build_upsert_batches(
                        list_data,
                        max_payload_bytes=self._max_upsert_payload_bytes,
                        max_records_per_batch=self._max_upsert_records_per_batch,
                    )
                    if len(upsert_batches) > 1:
                        logger.info(
                            f"[{self.workspace}] {self.namespace} flush: upsert split into "
                            f"{len(upsert_batches)} batches for {len(list_data)} records "
                            f"(max_payload={self._max_upsert_payload_bytes} batch={self._max_upsert_records_per_batch})"
                        )
                    for batch_index, (records_batch, estimated_bytes) in enumerate(
                        upsert_batches, 1
                    ):
                        if (
                            len(records_batch) == 1
                            and self._max_upsert_payload_bytes > 0
                            and estimated_bytes > self._max_upsert_payload_bytes
                        ):
                            logger.warning(
                                f"[{self.workspace}] {self.namespace} flush: single record "
                                f"id={records_batch[0].get('id')} estimated {estimated_bytes} bytes "
                                f"exceeds {self._max_upsert_payload_bytes}"
                            )
                        logger.debug(
                            f"[{self.workspace}] Milvus upsert batch {batch_index}/{len(upsert_batches)}: "
                            f"records={len(records_batch)}, estimated_payload_bytes={estimated_bytes}"
                        )
                        self._client.upsert(
                            collection_name=self.final_namespace, data=records_batch
                        )
                if pending_deletes:
                    # Chunk deletes by record count; pks are short strings so a
                    # count cap is enough to stay under the gRPC message limit.
                    delete_ids = list(pending_deletes)
                    delete_chunk = (
                        self._max_delete_records_per_batch
                        if self._max_delete_records_per_batch > 0
                        else len(delete_ids)
                    )
                    for i in range(0, len(delete_ids), delete_chunk):
                        self._client.delete(
                            collection_name=self.final_namespace,
                            pks=delete_ids[i : i + delete_chunk],
                        )
            except Exception as e:
                logger.error(
                    f"[{self.workspace}] Error flushing vector ops "
                    f"(upserts={len(pending_docs)}, "
                    f"deletes={len(pending_deletes)}): {e}"
                )
                raise

            # On success, clear the buffers in-place so external references
            # (e.g. drop()) see the cleared state.
            for doc_id in committed_ids:
                pending_docs.pop(doc_id, None)
            pending_deletes.clear()

    async def delete_entity(self, entity_name: str) -> None:
        """Buffer an entity vector delete by computing its hash ID."""
        entity_id = compute_mdhash_id(entity_name, prefix="ent-")
        async with self._flush_lock:
            self._pending_vector_docs.pop(entity_id, None)
            self._pending_vector_deletes.add(entity_id)
        logger.debug(
            f"[{self.workspace}] Buffered delete for entity {entity_name} (id={entity_id})"
        )

    async def delete_entity_relation(self, entity_name: str) -> None:
        """Delete all relation vectors where entity appears as src or tgt.

        The whole method runs under ``_flush_lock`` so the server-side query
        + delete cannot interleave with an in-flight bulk upsert.
        Server-side failures are re-raised (no log-and-swallow): the caller
        decides whether to retry.

        Buffer semantics — post-prune with caller short-circuit contract:
            Matching pending upserts in ``_pending_vector_docs`` are
            pruned **only after** the server-side query + delete
            succeeds. On failure the pending buffer stays intact and
            the exception propagates so the caller (``adelete_by_entity``
            in ``utils_graph.py``) can short-circuit before
            ``_persist_graph_updates`` flushes a half-cleaned buffer.

        Semantic note (deferred-buffer ↔ persisted divergence): pruning only
        consults the *current* buffered ``src_id`` / ``tgt_id`` view; we do
        not re-read the persisted row a buffered upsert is about to
        overwrite. So if a pending upsert is rewriting an already-persisted
        ``rel-X-Y`` so that its new ``src_id`` / ``tgt_id`` matches
        ``entity_name`` while the persisted row's do not (or vice versa),
        the persisted row will not be deleted by the server-side filter and
        the pending overwrite is dropped — i.e. the final state can diverge
        from the eager-flush ordering (upsert → flush → delete). Callers
        that require eager-equivalent semantics should call
        ``index_done_callback()`` before ``delete_entity_relation``.
        """

        def _prune_pending() -> None:
            for doc_id in [
                k
                for k, v in self._pending_vector_docs.items()
                if v.source.get("src_id") == entity_name
                or v.source.get("tgt_id") == entity_name
            ]:
                self._pending_vector_docs.pop(doc_id, None)

        async with self._flush_lock:
            if self._client is None:
                # No server state to mutate; buffer prune is the only
                # delete intent we can record.
                _prune_pending()
                return

            self._ensure_collection_loaded()

            expr = f'src_id == "{entity_name}" or tgt_id == "{entity_name}"'
            results = self._client.query(
                collection_name=self.final_namespace,
                filter=expr,
                output_fields=["id"],
            )

            if not results:
                # No server rows to delete — still safe to prune any
                # pending upserts so they can't re-create the relation.
                _prune_pending()
                logger.debug(
                    f"[{self.workspace}] No relations found for entity {entity_name}"
                )
                return

            relation_ids = [item["id"] for item in results]
            self._client.delete(collection_name=self.final_namespace, pks=relation_ids)
            # Server-side delete succeeded — safe to prune the pending
            # buffer so subsequent flushes don't re-upsert the deleted
            # relations.
            _prune_pending()
            logger.debug(
                f"[{self.workspace}] Deleted {len(relation_ids)} relations for {entity_name}"
            )

    async def delete(self, ids: list[str]) -> None:
        """Buffer vector deletes for batched flush."""
        if not ids:
            return
        if isinstance(ids, set):
            ids = list(ids)
        async with self._flush_lock:
            for doc_id in ids:
                self._pending_vector_docs.pop(doc_id, None)
                self._pending_vector_deletes.add(doc_id)
        logger.debug(
            f"[{self.workspace}] Buffered delete for {len(ids)} vectors in {self.namespace}"
        )

    async def get_by_id(self, id: str) -> dict[str, Any] | None:
        """Get vector data by its ID, with read-your-writes against the buffer."""
        async with self._flush_lock:
            if id in self._pending_vector_deletes:
                return None
            pending = self._pending_vector_docs.get(id)
            if pending is not None:
                doc = dict(pending.source)
                doc["id"] = id
                return doc

        try:
            # Ensure collection is loaded before querying
            self._ensure_collection_loaded()

            # Include all meta_fields (created_at is now always included) plus id
            output_fields = list(self.meta_fields) + ["id"]

            result = self._client.query(
                collection_name=self.final_namespace,
                filter=f'id == "{id}"',
                output_fields=output_fields,
            )

            if not result or len(result) == 0:
                return None

            return result[0]
        except Exception as e:
            logger.error(
                f"[{self.workspace}] Error retrieving vector data for ID {id}: {e}"
            )
            return None

    async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
        """Get multiple vector data by their IDs (read-your-writes), preserving order."""
        if not ids:
            return []

        buffered: dict[str, dict[str, Any] | None] = {}
        remaining: list[str] = []
        async with self._flush_lock:
            for doc_id in ids:
                if doc_id in self._pending_vector_deletes:
                    buffered[doc_id] = None
                    continue
                pending = self._pending_vector_docs.get(doc_id)
                if pending is not None:
                    doc = dict(pending.source)
                    doc["id"] = doc_id
                    buffered[doc_id] = doc
                    continue
                remaining.append(doc_id)

        result_map: dict[str, dict[str, Any]] = {}
        if remaining:
            try:
                # Ensure collection is loaded before querying
                self._ensure_collection_loaded()

                # Include all meta_fields (created_at is now always included) plus id
                output_fields = list(self.meta_fields) + ["id"]

                id_list = '", "'.join(remaining)
                filter_expr = f'id in ["{id_list}"]'

                result = self._client.query(
                    collection_name=self.final_namespace,
                    filter=filter_expr,
                    output_fields=output_fields,
                )

                if result:
                    for row in result:
                        if not row:
                            continue
                        row_id = row.get("id")
                        if row_id is not None:
                            result_map[str(row_id)] = row
            except Exception as e:
                logger.error(
                    f"[{self.workspace}] Error retrieving vector data for IDs {remaining}: {e}"
                )
                return []

        return [
            buffered[doc_id] if doc_id in buffered else result_map.get(str(doc_id))
            for doc_id in ids
        ]

    async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]:
        """Get vector embeddings for given IDs, with read-your-writes.

        Pending docs with `vector is None` trigger a lazy embed inside the
        lock; the resulting vector is cached on the buffered `_PendingVectorDoc`
        so the next flush won't re-embed the same content.
        """
        if not ids:
            return {}

        result: dict[str, list[float]] = {}
        remaining: list[str] = []
        async with self._flush_lock:
            docs_to_embed: list[tuple[str, _PendingVectorDoc]] = []
            for doc_id in ids:
                if doc_id in self._pending_vector_deletes:
                    continue
                pending = self._pending_vector_docs.get(doc_id)
                if pending is not None:
                    if pending.vector is None:
                        docs_to_embed.append((doc_id, pending))
                    else:
                        result[doc_id] = pending.vector
                    continue
                remaining.append(doc_id)

            if docs_to_embed:
                contents = [pdoc.content for _, pdoc in docs_to_embed]
                batches = [
                    contents[i : i + self._max_batch_size]
                    for i in range(0, len(contents), self._max_batch_size)
                ]
                try:
                    embeddings_list = await asyncio.gather(
                        *[
                            self.embedding_func(batch, context="document")
                            for batch in batches
                        ]
                    )
                except Exception as e:
                    logger.error(
                        f"[{self.workspace}] Error lazily embedding pending vectors "
                        f"(upserts={len(docs_to_embed)}): {e}"
                    )
                    raise
                embeddings = np.concatenate(embeddings_list)
                if len(embeddings) != len(docs_to_embed):
                    raise RuntimeError(
                        f"[{self.workspace}] Embedding count mismatch: expected "
                        f"{len(docs_to_embed)}, got {len(embeddings)}"
                    )
                for i, ((doc_id, pdoc), embedding) in enumerate(
                    zip(docs_to_embed, embeddings), start=1
                ):
                    # Cache float32 to match the flush path so the buffered
                    # vector dtype is uniform regardless of which path embedded.
                    pdoc.vector = np.array(embedding, dtype=np.float32).tolist()
                    result[doc_id] = pdoc.vector
                    await _cooperative_yield(i)

        if not remaining:
            return result

        try:
            self._ensure_collection_loaded()

            id_list = '", "'.join(remaining)
            filter_expr = f'id in ["{id_list}"]'

            rows = self._client.query(
                collection_name=self.final_namespace,
                filter=filter_expr,
                output_fields=["id", "vector"],
            )

            for item in rows or []:
                if item and "vector" in item and "id" in item:
                    vector_data = item["vector"]
                    if isinstance(vector_data, np.ndarray):
                        vector_data = vector_data.tolist()
                    # Match get_by_ids: stringify the server-returned id so
                    # callers can index the dict by the original requested id.
                    result[str(item["id"])] = vector_data
            return result
        except Exception as e:
            logger.error(
                f"[{self.workspace}] Error retrieving vectors by IDs from {self.namespace}: {e}"
            )
            return result

    async def finalize(self):
        """Flush pending vector ops, then release the Milvus gRPC channel.

        Every other server-backed storage releases its client here (Neo4j
        driver, Postgres pool, Mongo client, OpenSearch ClientManager); the
        MilvusClient owns a gRPC channel that should be closed explicitly
        too — ``close()`` is already used for that purpose by
        ``_rebuild_milvus_client``. We still fail loudly when a transient
        bulk error left writes buffered — the caller must not believe
        storage finalized cleanly.
        """
        flush_error: Exception | None = None
        try:
            await self._flush_pending_vector_ops()
        except Exception as e:
            flush_error = e

        # Release the gRPC channel after the flush so the flush can still use
        # the client. Best-effort: close() on a dead channel is a no-op (see
        # _rebuild_milvus_client). The connection is freed on every exit path,
        # matching the close-on-release pattern of the other server-backed
        # storages (Neo4j / Postgres / Mongo / OpenSearch).
        if self._client is not None:
            try:
                self._client.close()
            except Exception as close_error:
                logger.warning(
                    f"[{self.workspace}] Failed to close Milvus client: {close_error}"
                )

        # Read the residual buffer sizes under the flush lock so the
        # snapshot is consistent with any racing late-arriving mutator
        # (cancellation paths can land an upsert/delete between the flush
        # above and the post-mortem check below).
        async with self._flush_lock:
            pending_docs = len(self._pending_vector_docs)
            pending_deletes = len(self._pending_vector_deletes)

        if flush_error is not None:
            raise RuntimeError(
                f"[{self.workspace}] MilvusVectorDBStorage.finalize() flush raised; "
                f"{pending_docs} pending upserts and {pending_deletes} pending "
                f"deletes were left buffered (data lost)"
            ) from flush_error
        if pending_docs or pending_deletes:
            raise RuntimeError(
                f"[{self.workspace}] MilvusVectorDBStorage.finalize() left "
                f"{pending_docs} pending upserts and {pending_deletes} pending "
                f"deletes buffered after final flush attempt (these writes have been lost)"
            )

    async def drop(self) -> dict[str, str]:
        """Drop all data from the Milvus collection. Destructive.

        MUST only be called when ``pipeline_status`` is idle (see the
        Pipeline concurrency contract in ``AGENTS.md``); the only
        in-tree caller ``clear_documents`` enforces this.

        Caveat — only this instance's buffers are cleared. Other
        ``MilvusVectorDBStorage`` instances aliased onto the same
        ``final_namespace`` (multi-worker processes, or distinct
        workspaces collapsed by ``MILVUS_WORKSPACE``) keep their own
        buffers; a sibling whose prior flush failed and left buffers
        intact will, on its next flush, upsert those stale rows into
        the freshly recreated collection. Direct callers bypassing the
        idle precondition MUST flush every aliased instance first.

        Returns:
            dict[str, str]: ``{"status": "success"|"error", "message": str}``
        """
        try:
            async with self._flush_lock:
                # Discard any buffered writes before the collection is gone;
                # a concurrent flush would otherwise resurrect them.
                self._pending_vector_docs.clear()
                self._pending_vector_deletes.clear()

                # Drop the collection and recreate it empty.
                if self._client.has_collection(self.final_namespace):
                    self._client.drop_collection(self.final_namespace)

                # Recreate an EMPTY collection. Do NOT route through
                # _create_collection_if_not_exist here: with the suffixed
                # collection now gone it would see the intentionally-kept legacy
                # collection and re-run the legacy->suffixed migration, pulling
                # the just-dropped rows back in. That makes drop() non-empty
                # (clear_documents would leave stale legacy data behind) and
                # forces a needless full migration on every rebuild/clear.
                self._create_collection_with_schema(self.final_namespace)
                self._ensure_collection_loaded()

            logger.info(
                f"[{self.workspace}] Process {os.getpid()} drop Milvus collection {self.namespace}"
            )
            return {"status": "success", "message": "data dropped"}
        except Exception as e:
            logger.error(
                f"[{self.workspace}] Error dropping Milvus collection {self.namespace}: {e}"
            )
            return {"status": "error", "message": str(e)}