feat: 支持工作流多知识库向量检索

easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusClientManager.java

@@ -7,10 +7,12 @@
 package com.easyagents.store.milvus;
 
 import com.easyagents.core.util.StringUtil;
+import io.grpc.Context;
 import io.milvus.pool.MilvusClientV2Pool;
 import io.milvus.pool.PoolConfig;
 import io.milvus.v2.client.ConnectConfig;
 import io.milvus.v2.client.MilvusClientV2;
+import io.milvus.v2.client.RetryConfig;
 
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
@@ -20,6 +22,11 @@ import java.time.Duration;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Function;
 
@@ -29,14 +36,27 @@ import java.util.function.Function;
 public class MilvusClientManager implements AutoCloseable {
 
     private static final String POOL_KEY = "default";
+    private static final RetryConfig SINGLE_ATTEMPT_RETRY_CONFIG = RetryConfig.builder()
+        .maxRetryTimes(1)
+        .retryOnRateLimit(false)
+        .maxRetryTimeoutMs(0L)
+        .build();
 
     private final ReentrantReadWriteLock lifecycleLock = new ReentrantReadWriteLock();
     private final Set<String> initializedCollections =
         Collections.synchronizedSet(new HashSet<String>());
     private final Set<String> loadedCollections =
         Collections.synchronizedSet(new HashSet<String>());
+    private final ConcurrentMap<String, CollectionLoadTicket> collectionLoads =
+        new ConcurrentHashMap<String, CollectionLoadTicket>();
+    private final Set<Context.CancellableContext> activeContexts =
+        ConcurrentHashMap.newKeySet();
+    private final ConcurrentMap<Thread, Integer> activeOperations =
+        new ConcurrentHashMap<Thread, Integer>();
     private volatile ManagedMilvusClientV2Pool pool;
     private volatile String poolFingerprint;
+    private volatile long poolGeneration;
+    private volatile boolean acceptingOperations = true;
     private volatile boolean closed;
 
     public MilvusClientManager(MilvusVectorStoreConfig config) {
@@ -67,10 +87,20 @@ public class MilvusClientManager implements AutoCloseable {
     }
 
     public <T> T withClient(Function<MilvusClientV2, T> operation) {
+        return withClient(null, operation);
+    }
+
+    public <T> T withClient(
+        Duration maxWait,
+        Function<MilvusClientV2, T> operation
+    ) {
+        Thread operationThread = registerActiveOperation();
         lifecycleLock.readLock().lock();
         try {
             ManagedMilvusClientV2Pool currentPool = requireOpenPool();
-            MilvusClientV2 client = currentPool.getClient(POOL_KEY);
+            MilvusClientV2 client = maxWait == null
+                ? currentPool.getClient(POOL_KEY)
+                : currentPool.getClient(POOL_KEY, maxWait);
             if (client == null) {
                 throw new IllegalStateException(
                     "Milvus client pool is exhausted or unavailable"
@@ -78,18 +108,27 @@ public class MilvusClientManager implements AutoCloseable {
             }
             Throwable operationFailure = null;
             try {
-                return operation.apply(client);
+                client.retryConfig(SINGLE_ATTEMPT_RETRY_CONFIG);
+                Context.CancellableContext operationContext =
+                    Context.current().withCancellation();
+                try {
+                    return withRequestContext(operationContext,
+                        () -> operation.apply(client));
+                } catch (RuntimeException | Error exception) {
+                    throw exception;
+                } catch (Exception exception) {
+                    throw new IllegalStateException(
+                        "Milvus client operation failed", exception);
+                } finally {
+                    operationContext.cancel(null);
+                }
             } catch (RuntimeException | Error exception) {
                 operationFailure = exception;
                 throw exception;
             } finally {
                 RuntimeException cleanupFailure = null;
                 try {
-                    if (operationFailure == null) {
-                        currentPool.returnClient(POOL_KEY, client);
-                    } else {
-                        discardFailedClient(currentPool, client);
-                    }
+                    releaseClient(currentPool, client);
                 } catch (RuntimeException exception) {
                     cleanupFailure = exception;
                 }
@@ -101,10 +140,74 @@ public class MilvusClientManager implements AutoCloseable {
                 }
             }
         } finally {
+            unregisterActiveOperation(operationThread);
             lifecycleLock.readLock().unlock();
         }
     }
 
+    private Thread registerActiveOperation() {
+        ensureAcceptingOperations();
+        Thread currentThread = Thread.currentThread();
+        activeOperations.merge(currentThread, 1, Integer::sum);
+        if (!acceptingOperations) {
+            unregisterActiveOperation(currentThread);
+            ensureAcceptingOperations();
+        }
+        return currentThread;
+    }
+
+    private void unregisterActiveOperation(Thread operationThread) {
+        activeOperations.computeIfPresent(operationThread,
+            (thread, depth) -> depth <= 1 ? null : depth - 1);
+    }
+
+    private void releaseClient(
+        ManagedMilvusClientV2Pool currentPool,
+        MilvusClientV2 client
+    ) {
+        RuntimeException readinessFailure = null;
+        boolean reusable = false;
+        try {
+            reusable = client.clientIsReady();
+        } catch (RuntimeException exception) {
+            readinessFailure = exception;
+        }
+        try {
+            if (reusable) {
+                currentPool.returnClient(POOL_KEY, client);
+            } else {
+                discardFailedClient(currentPool, client);
+            }
+        } catch (RuntimeException cleanupFailure) {
+            if (readinessFailure == null) {
+                throw cleanupFailure;
+            }
+            readinessFailure.addSuppressed(cleanupFailure);
+        }
+        if (readinessFailure != null) {
+            throw readinessFailure;
+        }
+    }
+
+    <T> T withRequestContext(
+        Context.CancellableContext context,
+        Callable<T> operation
+    ) throws Exception {
+        ensureAcceptingOperations();
+        activeContexts.add(context);
+        if (!acceptingOperations) {
+            activeContexts.remove(context);
+            context.cancel(new CancellationException(
+                "Milvus client pool is unavailable"));
+            ensureAcceptingOperations();
+        }
+        try {
+            return context.call(operation);
+        } finally {
+            activeContexts.remove(context);
+        }
+    }
+
     private void discardFailedClient(
         ManagedMilvusClientV2Pool currentPool,
         MilvusClientV2 client
@@ -115,15 +218,6 @@ public class MilvusClientManager implements AutoCloseable {
         } catch (RuntimeException exception) {
             cleanupFailure = exception;
         }
-        try {
-            currentPool.clear(POOL_KEY);
-        } catch (RuntimeException exception) {
-            if (cleanupFailure == null) {
-                cleanupFailure = exception;
-            } else {
-                cleanupFailure.addSuppressed(exception);
-            }
-        }
         if (cleanupFailure != null) {
             throw cleanupFailure;
         }
@@ -131,16 +225,19 @@ public class MilvusClientManager implements AutoCloseable {
 
     /**
      * Rebuilds the pool when connection or pool settings change.
-     * Active RPCs finish before the old pool is closed.
+     * Active operations are cancelled before the old pool is closed.
      *
      * @return true when a new pool was installed
      */
-    public boolean reconfigureIfNeeded(MilvusVectorStoreConfig config) {
+    public synchronized boolean reconfigureIfNeeded(MilvusVectorStoreConfig config) {
         PoolSettings nextSettings = PoolSettings.from(config);
         String nextFingerprint = fingerprint(nextSettings);
         if (nextFingerprint.equals(poolFingerprint)) {
             return false;
         }
+        acceptingOperations = false;
+        cancelActiveContexts("Milvus client pool is reconfiguring");
+        interruptActiveOperations();
         lifecycleLock.writeLock().lock();
         try {
             if (closed) {
@@ -153,14 +250,19 @@ public class MilvusClientManager implements AutoCloseable {
             ManagedMilvusClientV2Pool previous = pool;
             pool = replacement;
             poolFingerprint = nextFingerprint;
+            poolGeneration++;
             initializedCollections.clear();
             loadedCollections.clear();
+            failCollectionLoads("Milvus client pool was reconfigured");
             if (previous != null) {
                 previous.close();
             }
             return true;
         } finally {
             lifecycleLock.writeLock().unlock();
+            if (!closed) {
+                acceptingOperations = true;
+            }
         }
     }
 
@@ -180,10 +282,6 @@ public class MilvusClientManager implements AutoCloseable {
         return loadedCollections.contains(collectionName);
     }
 
-    Object loadedCollectionsLock() {
-        return loadedCollections;
-    }
-
     void markCollectionLoaded(String collectionName) {
         loadedCollections.add(collectionName);
     }
@@ -192,6 +290,62 @@ public class MilvusClientManager implements AutoCloseable {
         loadedCollections.remove(collectionName);
     }
 
+    CollectionLoadTicket beginCollectionLoad(String collectionName) {
+        ensureAcceptingOperations();
+        lifecycleLock.readLock().lock();
+        try {
+            requireOpenPool();
+            CollectionLoadTicket candidate = new CollectionLoadTicket(
+                collectionName,
+                poolGeneration,
+                new CompletableFuture<Void>(),
+                true
+            );
+            CollectionLoadTicket existing = collectionLoads.putIfAbsent(
+                collectionName, candidate);
+            return existing == null ? candidate : existing.asFollower();
+        } finally {
+            lifecycleLock.readLock().unlock();
+        }
+    }
+
+    void completeCollectionLoad(CollectionLoadTicket ticket) {
+        lifecycleLock.readLock().lock();
+        try {
+            requireOpenPool();
+            if (ticket.generation != poolGeneration) {
+                throw new IllegalStateException(
+                    "Milvus client pool changed while loading collection: "
+                        + ticket.collectionName
+                );
+            }
+            loadedCollections.add(ticket.collectionName);
+            ticket.completion.complete(null);
+        } finally {
+            lifecycleLock.readLock().unlock();
+        }
+    }
+
+    void failCollectionLoad(
+        CollectionLoadTicket ticket,
+        Throwable failure,
+        boolean retryableForFollowers
+    ) {
+        if (retryableForFollowers && ticket.leader) {
+            collectionLoads.remove(ticket.collectionName, ticket);
+        }
+        Throwable sharedFailure = retryableForFollowers
+            ? new RetryableCollectionLoadException(failure)
+            : failure;
+        ticket.completion.completeExceptionally(sharedFailure);
+    }
+
+    void endCollectionLoad(CollectionLoadTicket ticket) {
+        if (ticket.leader) {
+            collectionLoads.remove(ticket.collectionName, ticket);
+        }
+    }
+
     public int getActiveClientCount() {
         lifecycleLock.readLock().lock();
         try {
@@ -211,15 +365,23 @@ public class MilvusClientManager implements AutoCloseable {
     }
 
     @Override
-    public void close() {
+    public synchronized void close() {
+        if (closed) {
+            return;
+        }
+        acceptingOperations = false;
+        closed = true;
+        cancelActiveContexts("Milvus client pool is closing");
+        interruptActiveOperations();
         lifecycleLock.writeLock().lock();
         try {
             initializedCollections.clear();
             loadedCollections.clear();
+            poolGeneration++;
+            failCollectionLoads("Milvus client pool was closed");
             ManagedMilvusClientV2Pool currentPool = pool;
             pool = null;
             poolFingerprint = null;
-            closed = true;
             if (currentPool != null) {
                 currentPool.close();
             }
@@ -236,6 +398,41 @@ public class MilvusClientManager implements AutoCloseable {
         return currentPool;
     }
 
+    boolean isClosed() {
+        return closed;
+    }
+
+    private void ensureAcceptingOperations() {
+        if (!acceptingOperations) {
+            throw new IllegalStateException(closed
+                ? "Milvus client pool is closed"
+                : "Milvus client pool is reconfiguring");
+        }
+    }
+
+    private void failCollectionLoads(String message) {
+        IllegalStateException failure = new IllegalStateException(message);
+        for (CollectionLoadTicket ticket : collectionLoads.values()) {
+            ticket.completion.completeExceptionally(failure);
+        }
+        collectionLoads.clear();
+    }
+
+    private void cancelActiveContexts(String message) {
+        for (Context.CancellableContext context : activeContexts) {
+            context.cancel(new CancellationException(message));
+        }
+    }
+
+    private void interruptActiveOperations() {
+        Thread currentThread = Thread.currentThread();
+        for (Thread operationThread : activeOperations.keySet()) {
+            if (operationThread != currentThread) {
+                operationThread.interrupt();
+            }
+        }
+    }
+
     private static String fingerprint(PoolSettings settings) {
         String value = String.join("\u0000",
             String.valueOf(settings.uri()),
@@ -334,6 +531,64 @@ public class MilvusClientManager implements AutoCloseable {
                 throw new IllegalStateException("Unable to invalidate Milvus client", exception);
             }
         }
+
+        private MilvusClientV2 getClient(String key, Duration maxWait) {
+            if (maxWait == null || maxWait.isZero() || maxWait.isNegative()) {
+                throw new IllegalArgumentException("maxWait must be greater than zero");
+            }
+            try {
+                long waitMillis = Math.max(1L, maxWait.toMillis());
+                return clientPool.borrowObject(key, waitMillis);
+            } catch (InterruptedException exception) {
+                Thread.currentThread().interrupt();
+                throw new IllegalStateException(
+                    "Interrupted while waiting for a Milvus client", exception);
+            } catch (Exception exception) {
+                throw new IllegalStateException(
+                    "Unable to borrow a Milvus client", exception);
+            }
+        }
+    }
+
+    static final class CollectionLoadTicket {
+
+        private final String collectionName;
+        private final long generation;
+        private final CompletableFuture<Void> completion;
+        private final boolean leader;
+
+        private CollectionLoadTicket(
+            String collectionName,
+            long generation,
+            CompletableFuture<Void> completion,
+            boolean leader
+        ) {
+            this.collectionName = collectionName;
+            this.generation = generation;
+            this.completion = completion;
+            this.leader = leader;
+        }
+
+        boolean isLeader() {
+            return leader;
+        }
+
+        CompletableFuture<Void> completion() {
+            return completion;
+        }
+
+        private CollectionLoadTicket asFollower() {
+            return new CollectionLoadTicket(
+                collectionName, generation, completion, false);
+        }
+    }
+
+    static final class RetryableCollectionLoadException
+        extends RuntimeException {
+
+        private RetryableCollectionLoadException(Throwable cause) {
+            super("The collection load leader exhausted its local budget", cause);
+        }
     }
 
     static String normalizeAndValidateUri(String uri) {

easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStore.java

@@ -22,9 +22,12 @@ import com.easyagents.core.store.DocumentStore;
 import com.easyagents.core.store.SearchWrapper;
 import com.easyagents.core.store.StoreOptions;
 import com.easyagents.core.store.StoreResult;
+import com.easyagents.core.store.StoreTimeoutException;
 import com.easyagents.core.util.CollectionUtil;
 import com.easyagents.core.util.Maps;
 import com.easyagents.core.util.StringUtil;
+import io.grpc.Context;
+import io.grpc.Status;
 import io.milvus.v2.client.MilvusClientV2;
 import io.milvus.v2.common.ConsistencyLevel;
 import io.milvus.v2.common.DataType;
@@ -40,7 +43,17 @@ import io.milvus.v2.service.vector.response.SearchResp;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.time.Duration;
 import java.util.*;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Function;
 
 /**
  * Milvus vector store based on Milvus Java SDK v2.
@@ -50,6 +63,16 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
     private static final Logger LOG = LoggerFactory.getLogger(MilvusVectorStore.class);
     private static final long LOAD_TIMEOUT_MS = 30_000L;
     private static final long LOAD_POLL_INTERVAL_MS = 200L;
+    private static final long DEADLINE_SAFETY_MARGIN_MS = 200L;
+    private static final ScheduledExecutorService DEADLINE_SCHEDULER =
+        Executors.newSingleThreadScheduledExecutor(new ThreadFactory() {
+            @Override
+            public Thread newThread(Runnable runnable) {
+                Thread thread = new Thread(runnable, "milvus-deadline");
+                thread.setDaemon(true);
+                return thread;
+            }
+        });
 
     private static final String FIELD_ID = "id";
     private static final String FIELD_CONTENT = "content";
@@ -63,9 +86,10 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
     private final String defaultCollectionName;
     private final boolean ownsClientManager;
     private volatile MilvusClientV2 compatibilityClient;
+    private volatile boolean closed;
 
     public MilvusVectorStore(MilvusVectorStoreConfig config) {
-        this(config, new MilvusClientManager(config), true);
+        this(config, createOwnedClientManager(config), true);
     }
 
     public MilvusVectorStore(
@@ -80,12 +104,35 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
         MilvusClientManager clientManager,
         boolean ownsClientManager
     ) {
+        validateConfig(config);
         this.config = config;
         this.defaultCollectionName = config.getDefaultCollectionName();
         this.clientManager = Objects.requireNonNull(clientManager, "clientManager");
         this.ownsClientManager = ownsClientManager;
     }
 
+    private static MilvusClientManager createOwnedClientManager(
+        MilvusVectorStoreConfig config
+    ) {
+        validateConfig(config);
+        return new MilvusClientManager(config);
+    }
+
+    private static void validateConfig(MilvusVectorStoreConfig config) {
+        Objects.requireNonNull(config, "config");
+        if (config.getSearchTimeoutMillis() <= DEADLINE_SAFETY_MARGIN_MS) {
+            throw new IllegalArgumentException(
+                "Milvus searchTimeoutMillis must be greater than "
+                    + DEADLINE_SAFETY_MARGIN_MS
+            );
+        }
+        if (config.getPoolMaxWaitMillis() <= 0L) {
+            throw new IllegalArgumentException(
+                "Milvus poolMaxWaitMillis must be greater than zero"
+            );
+        }
+    }
+
     @Override
     public StoreResult doStore(List<Document> documents, StoreOptions options) {
         if (CollectionUtil.noItems(documents)) {
@@ -184,78 +231,104 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
         if (StringUtil.noText(collectionName)) {
             throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName.");
         }
-        return clientManager.withClient(client -> {
-            ensureCollectionLoaded(client, collectionName);
-            if (wrapper.getVector() == null || wrapper.getVector().length == 0) {
-                return queryByCondition(client, wrapper, options, collectionName);
+        long timeoutMillis = resolveSearchTimeoutMillis(options);
+        long rpcBudgetMillis = timeoutMillis - DEADLINE_SAFETY_MARGIN_MS;
+        long deadlineNanos = deadlineAfterMillis(rpcBudgetMillis);
+        Context.CancellableContext context = Context.current().withDeadlineAfter(
+            rpcBudgetMillis, TimeUnit.MILLISECONDS, DEADLINE_SCHEDULER);
+        try {
+            return clientManager.withRequestContext(context, () ->
+                searchWithinDeadline(
+                    wrapper, options, collectionName, deadlineNanos));
+        } catch (RuntimeException exception) {
+            if (!(exception instanceof StoreTimeoutException)
+                && deadlineExpired(deadlineNanos, exception)) {
+                throw timeoutException(collectionName, exception);
             }
-            return searchByVector(client, wrapper, options, collectionName);
-        });
+            throw exception;
+        } catch (Exception exception) {
+            throw new IllegalStateException("Milvus search failed", exception);
+        } finally {
+            context.cancel(null);
+        }
+    }
+
+    private List<Document> searchWithinDeadline(
+        SearchWrapper wrapper,
+        StoreOptions options,
+        String collectionName,
+        long deadlineNanos
+    ) {
+        String operation = wrapper.getVector() == null
+            || wrapper.getVector().length == 0
+            ? "query"
+            : "search";
+        ensureCollectionLoaded(collectionName, deadlineNanos);
+        try {
+            return searchOnce(wrapper, options, collectionName, deadlineNanos);
+        } catch (RuntimeException exception) {
+            if (!isCollectionNotLoaded(exception)) {
+                throw propagateSearchFailure(
+                    operation, collectionName, exception);
+            }
+            clientManager.markCollectionUnloaded(collectionName);
+            try {
+                ensureCollectionLoaded(collectionName, deadlineNanos);
+                return searchOnce(
+                    wrapper, options, collectionName, deadlineNanos);
+            } catch (RuntimeException retryException) {
+                retryException.addSuppressed(exception);
+                throw propagateSearchFailure(
+                    operation, collectionName, retryException);
+            }
+        }
+    }
+
+    private List<Document> searchOnce(
+        SearchWrapper wrapper,
+        StoreOptions options,
+        String collectionName,
+        long deadlineNanos
+    ) {
+        if (wrapper.getVector() == null || wrapper.getVector().length == 0) {
+            return queryByCondition(
+                wrapper, options, collectionName, deadlineNanos);
+        }
+        return searchByVector(wrapper, options, collectionName, deadlineNanos);
     }
 
     private List<Document> searchByVector(
-        MilvusClientV2 client,
         SearchWrapper wrapper,
         StoreOptions options,
-        String collectionName
+        String collectionName,
+        long deadlineNanos
     ) {
         SearchReq searchReq = buildSearchReq(wrapper, options, collectionName);
-        try {
-            SearchResp resp = client.search(searchReq);
-            return parseSearchResults(resp, wrapper.getMinScore());
-        } catch (Exception e) {
-            if (isCollectionNotLoaded(e)) {
-                clientManager.markCollectionUnloaded(collectionName);
-                try {
-                    ensureCollectionLoaded(client, collectionName);
-                    SearchResp retryResp = client.search(searchReq);
-                    return parseSearchResults(retryResp, wrapper.getMinScore());
-                } catch (Exception retryException) {
-                    retryException.addSuppressed(e);
-                    throw propagateSearchFailure("search", collectionName, retryException);
-                }
-            }
-            throw propagateSearchFailure("search", collectionName, e);
-        }
+        SearchResp resp = withClientBeforeDeadline(
+            deadlineNanos, client -> client.search(searchReq));
+        return parseSearchResults(resp, wrapper.getMinScore());
     }
 
     private List<Document> queryByCondition(
-        MilvusClientV2 client,
         SearchWrapper wrapper,
         StoreOptions options,
-        String collectionName
+        String collectionName,
+        long deadlineNanos
     ) {
         QueryReq queryReq = buildQueryReq(wrapper, options, collectionName);
-        try {
-            QueryResp resp = client.query(queryReq);
-            return parseQueryResults(resp);
-        } catch (Exception e) {
-            if (isCollectionNotLoaded(e)) {
-                clientManager.markCollectionUnloaded(collectionName);
-                try {
-                    ensureCollectionLoaded(client, collectionName);
-                    QueryResp retryResp = client.query(queryReq);
-                    return parseQueryResults(retryResp);
-                } catch (Exception retryException) {
-                    retryException.addSuppressed(e);
-                    throw propagateSearchFailure("query", collectionName, retryException);
-                }
-            }
-            throw propagateSearchFailure("query", collectionName, e);
-        }
+        QueryResp resp = withClientBeforeDeadline(
+            deadlineNanos, client -> client.query(queryReq));
+        return parseQueryResults(resp);
     }
 
     private RuntimeException propagateSearchFailure(
         String operation,
         String collectionName,
-        Exception exception
+        RuntimeException exception
     ) {
         LOG.error("Milvus {} failed. collection={}, message={}",
             operation, collectionName, exception.getMessage(), exception);
-        if (exception instanceof RuntimeException runtimeException) {
-            return runtimeException;
-        }
-        return new IllegalStateException("Milvus " + operation + " failed", exception);
+        return exception;
     }
 
     private SearchReq buildSearchReq(SearchWrapper wrapper, StoreOptions options, String collectionName) {
@@ -450,46 +523,150 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
         }
     }
 
-    private void ensureCollectionLoaded(MilvusClientV2 client, String collectionName) {
-        if (clientManager.isCollectionLoaded(collectionName)) {
-            return;
-        }
-        synchronized (clientManager.loadedCollectionsLock()) {
-            if (clientManager.isCollectionLoaded(collectionName)) {
-                return;
+    private void ensureCollectionLoaded(
+        String collectionName,
+        long deadlineNanos
+    ) {
+        while (!clientManager.isCollectionLoaded(collectionName)) {
+            MilvusClientManager.CollectionLoadTicket ticket =
+                clientManager.beginCollectionLoad(collectionName);
+            if (!ticket.isLeader()) {
+                if (awaitCollectionLoad(
+                    ticket, collectionName, deadlineNanos)) {
+                    return;
+                }
+                continue;
             }
-            boolean loaded = false;
             try {
-                loaded = Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()));
-            } catch (Exception e) {
-                LOG.warn("Milvus getLoadState failed. collection={}, message={}", collectionName, e.getMessage());
+                if (clientManager.isCollectionLoaded(collectionName)) {
+                    clientManager.completeCollectionLoad(ticket);
+                    return;
+                }
+                withClientBeforeDeadline(deadlineNanos, client -> {
+                    boolean loaded = Boolean.TRUE.equals(client.getLoadState(
+                        GetLoadStateReq.builder()
+                            .collectionName(collectionName)
+                            .build()
+                    ));
+                    if (!loaded) {
+                        client.loadCollection(LoadCollectionReq.builder()
+                            .collectionName(collectionName)
+                            .async(false)
+                            .build());
+                        waitForCollectionLoaded(
+                            client, collectionName, deadlineNanos);
+                    }
+                    return null;
+                });
+                clientManager.completeCollectionLoad(ticket);
+                return;
+            } catch (RuntimeException | Error failure) {
+                clientManager.failCollectionLoad(
+                    ticket, failure, isLeaderLocalAbort(failure));
+                throw failure;
+            } finally {
+                clientManager.endCollectionLoad(ticket);
             }
-
-            if (!loaded) {
-                client.loadCollection(LoadCollectionReq.builder().collectionName(collectionName).build());
-                waitForCollectionLoaded(client, collectionName);
-            }
-            clientManager.markCollectionLoaded(collectionName);
         }
     }
 
-    private void waitForCollectionLoaded(MilvusClientV2 client, String collectionName) {
-        long deadline = System.currentTimeMillis() + LOAD_TIMEOUT_MS;
-        while (System.currentTimeMillis() < deadline) {
+    private boolean isLeaderLocalAbort(Throwable failure) {
+        if (clientManager.isClosed()) {
+            return false;
+        }
+        if (failure instanceof StoreTimeoutException
+            || Thread.currentThread().isInterrupted()
+            || Context.current().isCancelled()) {
+            return true;
+        }
+        Throwable current = failure;
+        while (current != null) {
+            if (current instanceof InterruptedException) {
+                return true;
+            }
+            current = current.getCause();
+        }
+        return false;
+    }
+
+    private boolean awaitCollectionLoad(
+        MilvusClientManager.CollectionLoadTicket ticket,
+        String collectionName,
+        long deadlineNanos
+    ) {
+        Context currentContext = Context.current();
+        CompletableFuture<Void> cancelled = new CompletableFuture<Void>();
+        Context.CancellationListener cancellationListener = context ->
+            cancelled.completeExceptionally(new CancellationException(
+                "Milvus search was cancelled"));
+        currentContext.addListener(cancellationListener, Runnable::run);
+        try {
+            CompletableFuture.anyOf(ticket.completion(), cancelled).get(
+                remainingNanos(deadlineNanos, collectionName),
+                TimeUnit.NANOSECONDS
+            );
+            return true;
+        } catch (InterruptedException exception) {
+            Thread.currentThread().interrupt();
+            throw new IllegalStateException(
+                "Interrupted while loading Milvus collection: "
+                    + collectionName,
+                exception
+            );
+        } catch (TimeoutException exception) {
+            throw timeoutException(collectionName, exception);
+        } catch (ExecutionException exception) {
+            Throwable cause = exception.getCause();
+            if (cause instanceof MilvusClientManager
+                .RetryableCollectionLoadException) {
+                remainingNanos(deadlineNanos, collectionName);
+                return false;
+            }
+            if (cause instanceof RuntimeException runtimeException) {
+                throw runtimeException;
+            }
+            if (cause instanceof Error error) {
+                throw error;
+            }
+            throw new IllegalStateException(
+                "Unable to load Milvus collection: " + collectionName,
+                cause
+            );
+        } catch (CancellationException exception) {
+            throw new IllegalStateException(
+                "Milvus collection load was cancelled: " + collectionName,
+                exception
+            );
+        } finally {
+            currentContext.removeListener(cancellationListener);
+        }
+    }
+
+    private void waitForCollectionLoaded(
+        MilvusClientV2 client,
+        String collectionName,
+        long deadlineNanos
+    ) {
+        while (true) {
+            long remainingNanos = remainingNanos(
+                deadlineNanos, collectionName);
             if (Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()))) {
                 return;
             }
             try {
-                Thread.sleep(LOAD_POLL_INTERVAL_MS);
+                long sleepMillis = Math.min(
+                    LOAD_POLL_INTERVAL_MS,
+                    Math.max(1L, TimeUnit.NANOSECONDS.toMillis(remainingNanos))
+                );
+                Thread.sleep(sleepMillis);
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
                 throw new IllegalStateException("Interrupted while loading Milvus collection: " + collectionName, e);
             }
         }
-        throw new IllegalStateException("Timeout waiting for Milvus collection loaded: " + collectionName);
     }
 
-    private boolean isCollectionNotLoaded(Exception e) {
+    private boolean isCollectionNotLoaded(Throwable e) {
         Throwable current = e;
         while (current != null) {
             String message = current.getMessage();
@@ -501,6 +678,90 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
         return false;
     }
 
+    private <T> T withClientBeforeDeadline(
+        long deadlineNanos,
+        Function<MilvusClientV2, T> operation
+    ) {
+        long remainingNanos = remainingNanos(deadlineNanos, null);
+        long poolWaitNanos = TimeUnit.MILLISECONDS.toNanos(
+            config.getPoolMaxWaitMillis());
+        try {
+            return clientManager.withClient(
+                Duration.ofNanos(Math.min(remainingNanos, poolWaitNanos)),
+                operation
+            );
+        } catch (RuntimeException exception) {
+            if (deadlineExpired(deadlineNanos, exception)) {
+                throw timeoutException(null, exception);
+            }
+            throw exception;
+        }
+    }
+
+    private long resolveSearchTimeoutMillis(StoreOptions options) {
+        long timeoutMillis = config.getSearchTimeoutMillis();
+        Long requestedTimeoutMillis = options.getTimeoutMillis();
+        if (requestedTimeoutMillis != null) {
+            timeoutMillis = Math.min(timeoutMillis, requestedTimeoutMillis);
+        }
+        if (timeoutMillis <= DEADLINE_SAFETY_MARGIN_MS) {
+            throw new StoreTimeoutException(
+                "Insufficient time remaining for Milvus search"
+            );
+        }
+        return timeoutMillis;
+    }
+
+    private static long deadlineAfterMillis(long timeoutMillis) {
+        long now = System.nanoTime();
+        long timeoutNanos = TimeUnit.MILLISECONDS.toNanos(timeoutMillis);
+        if (now > Long.MAX_VALUE - timeoutNanos) {
+            return Long.MAX_VALUE;
+        }
+        return now + timeoutNanos;
+    }
+
+    private static long remainingNanos(
+        long deadlineNanos,
+        String collectionName
+    ) {
+        if (deadlineNanos == Long.MAX_VALUE) {
+            return Long.MAX_VALUE;
+        }
+        long remaining = deadlineNanos - System.nanoTime();
+        if (remaining <= 0L) {
+            throw timeoutException(collectionName, null);
+        }
+        return remaining;
+    }
+
+    private static boolean deadlineExpired(
+        long deadlineNanos,
+        Throwable failure
+    ) {
+        if (deadlineNanos != Long.MAX_VALUE
+            && System.nanoTime() >= deadlineNanos) {
+            return true;
+        }
+        Throwable cancellationCause = Context.current().cancellationCause();
+        return cancellationCause instanceof TimeoutException
+            || Status.fromThrowable(failure).getCode()
+                == Status.Code.DEADLINE_EXCEEDED;
+    }
+
+    private static StoreTimeoutException timeoutException(
+        String collectionName,
+        Throwable cause
+    ) {
+        String suffix = StringUtil.hasText(collectionName)
+            ? ": " + collectionName
+            : "";
+        return new StoreTimeoutException(
+            "Timeout waiting for Milvus search" + suffix,
+            cause
+        );
+    }
+
     private void createCollection(MilvusClientV2 client, String collectionName, int dimension) {
         List<CreateCollectionReq.FieldSchema> fieldSchemaList = new ArrayList<CreateCollectionReq.FieldSchema>();
         fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder()
@@ -547,7 +808,30 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
             .indexParams(indexParams)
             .build();
         client.createCollection(createCollectionReq);
-        ensureCollectionLoaded(client, collectionName);
+        ensureCollectionLoadedForWrite(client, collectionName);
+    }
+
+    private void ensureCollectionLoadedForWrite(
+        MilvusClientV2 client,
+        String collectionName
+    ) {
+        if (clientManager.isCollectionLoaded(collectionName)) {
+            return;
+        }
+        boolean loaded = Boolean.TRUE.equals(client.getLoadState(
+            GetLoadStateReq.builder().collectionName(collectionName).build()));
+        if (!loaded) {
+            client.loadCollection(LoadCollectionReq.builder()
+                .collectionName(collectionName)
+                .async(false)
+                .build());
+            waitForCollectionLoaded(
+                client,
+                collectionName,
+                deadlineAfterMillis(LOAD_TIMEOUT_MS)
+            );
+        }
+        clientManager.markCollectionLoaded(collectionName);
     }
 
     public boolean checkAvailable() {
@@ -568,25 +852,29 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
      * Prefer store operations so pooled lifecycle management remains automatic.
      */
     @Deprecated
-    public MilvusClientV2 getClient() {
-        MilvusClientV2 current = compatibilityClient;
-        if (current != null) {
-            return current;
+    public synchronized MilvusClientV2 getClient() {
+        if (closed) {
+            throw new IllegalStateException("Milvus vector store is closed");
         }
-        synchronized (this) {
-            if (compatibilityClient == null) {
-                compatibilityClient = new MilvusClientV2(
-                    MilvusClientManager.buildConnectConfig(config)
-                );
-            }
-            return compatibilityClient;
+        if (compatibilityClient == null) {
+            compatibilityClient = new MilvusClientV2(
+                MilvusClientManager.buildConnectConfig(config)
+            );
         }
+        return compatibilityClient;
     }
 
     @Override
     public void close() {
-        MilvusClientV2 legacyClient = compatibilityClient;
-        compatibilityClient = null;
+        MilvusClientV2 legacyClient;
+        synchronized (this) {
+            if (closed) {
+                return;
+            }
+            closed = true;
+            legacyClient = compatibilityClient;
+            compatibilityClient = null;
+        }
         if (legacyClient != null) {
             try {
                 legacyClient.close(1L);

easy-agents-store/easy-agents-store-milvus/src/main/java/com/easyagents/store/milvus/MilvusVectorStoreConfig.java

@@ -37,6 +37,7 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig {
     private long poolMaxWaitMillis = 3_000L;
     private long poolEvictionIntervalMillis = 60_000L;
     private long poolMinEvictableIdleMillis = 300_000L;
+    private long searchTimeoutMillis = 10_000L;
 
     public String getUri() {
         return uri;
@@ -150,6 +151,14 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig {
         this.poolMinEvictableIdleMillis = poolMinEvictableIdleMillis;
     }
 
+    public long getSearchTimeoutMillis() {
+        return searchTimeoutMillis;
+    }
+
+    public void setSearchTimeoutMillis(long searchTimeoutMillis) {
+        this.searchTimeoutMillis = searchTimeoutMillis;
+    }
+
     @Override
     public boolean checkAvailable() {
         return StringUtil.hasText(this.uri);

easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreCompatibilityTest.java

@@ -72,4 +72,21 @@ public class MilvusVectorStoreCompatibilityTest {
         Assert.assertEquals(Double.valueOf(0.0D), MilvusVectorStore.normalizeScore(-1.0F));
         Assert.assertNull(MilvusVectorStore.normalizeScore(null));
     }
+
+    @Test
+    public void shouldRejectCompatibilityClientAfterStoreCloses() {
+        MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
+        config.setUri("http://127.0.0.1:19530");
+        MilvusVectorStore store = new MilvusVectorStore(config);
+
+        store.close();
+
+        try {
+            store.getClient();
+            Assert.fail("A closed store must not recreate a compatibility client");
+        } catch (IllegalStateException expected) {
+            Assert.assertEquals(
+                "Milvus vector store is closed", expected.getMessage());
+        }
+    }
 }

easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreConfigTest.java

@@ -44,5 +44,14 @@ public class MilvusVectorStoreConfigTest {
         Assert.assertEquals(1, config.getPoolMinIdlePerKey());
         Assert.assertEquals(3_000L, config.getPoolMaxWaitMillis());
         Assert.assertEquals(300_000L, config.getPoolMinEvictableIdleMillis());
+        Assert.assertEquals(10_000L, config.getSearchTimeoutMillis());
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testSearchTimeoutMustLeaveCleanupMargin() {
+        MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
+        config.setUri("http://127.0.0.1:19530");
+        config.setSearchTimeoutMillis(200L);
+        new MilvusVectorStore(config);
     }
 }

easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreGrpcTest.java

@@ -0,0 +1,935 @@
+package com.easyagents.store.milvus;
+
+import com.easyagents.core.document.Document;
+import com.easyagents.core.store.SearchWrapper;
+import com.easyagents.core.store.StoreOptions;
+import com.easyagents.core.store.StoreTimeoutException;
+import io.grpc.Context;
+import io.grpc.Server;
+import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
+import io.grpc.stub.ServerCallStreamObserver;
+import io.grpc.stub.StreamObserver;
+import io.milvus.grpc.CheckHealthRequest;
+import io.milvus.grpc.CheckHealthResponse;
+import io.milvus.grpc.ConnectRequest;
+import io.milvus.grpc.ConnectResponse;
+import io.milvus.grpc.CollectionSchema;
+import io.milvus.grpc.DataType;
+import io.milvus.grpc.DescribeCollectionRequest;
+import io.milvus.grpc.DescribeCollectionResponse;
+import io.milvus.grpc.ErrorCode;
+import io.milvus.grpc.FieldSchema;
+import io.milvus.grpc.GetLoadStateRequest;
+import io.milvus.grpc.GetLoadStateResponse;
+import io.milvus.grpc.ListDatabasesRequest;
+import io.milvus.grpc.ListDatabasesResponse;
+import io.milvus.grpc.LoadCollectionRequest;
+import io.milvus.grpc.LoadState;
+import io.milvus.grpc.MilvusServiceGrpc;
+import io.milvus.grpc.QueryRequest;
+import io.milvus.grpc.QueryResults;
+import io.milvus.v2.client.MilvusClientV2;
+import io.milvus.v2.service.vector.request.QueryReq;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BooleanSupplier;
+
+public class MilvusVectorStoreGrpcTest {
+
+    private static final long AWAIT_SECONDS = 5L;
+
+    @Test(timeout = 10_000L)
+    public void shouldUseOneSdkAttemptAndKeepClientReusable() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 2_000L)) {
+            fixture.manager.markCollectionLoaded("docs");
+            Assert.assertEquals(
+                0L,
+                MilvusClientManager.buildConnectConfig(fixture.config).getRpcDeadlineMs()
+            );
+
+            server.failQueriesWithUnavailable();
+            assertSearchFails(fixture.store, "docs");
+
+            Assert.assertEquals(1, server.queryCalls.get());
+            Assert.assertEquals(0, fixture.manager.getActiveClientCount());
+            Assert.assertEquals(1, fixture.manager.getIdleClientCount());
+
+            server.succeedQueries();
+            Assert.assertTrue(search(fixture.store, "docs").isEmpty());
+            Assert.assertEquals(2, server.queryCalls.get());
+            Assert.assertEquals(1, server.connectCalls.get());
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldReleaseAndReuseClientAfterContextDeadline() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 1_000L)) {
+            fixture.manager.markCollectionLoaded("docs");
+            fixture.manager.withClient(client -> client);
+            QueryBlock block = server.blockQueries();
+            ExecutorService executor = Executors.newSingleThreadExecutor();
+            try {
+                Future<List<Document>> search = executor.submit(() ->
+                    search(fixture.store, "docs"));
+
+                block.awaitEntered();
+                block.awaitCancelled();
+                Throwable failure = futureFailure(search);
+                Assert.assertTrue(failure.toString(),
+                    failure instanceof StoreTimeoutException);
+
+                Assert.assertTrue(server.querySawDeadline.get());
+                Assert.assertEquals(0, fixture.manager.getActiveClientCount());
+                Assert.assertEquals(1, fixture.manager.getIdleClientCount());
+
+                server.succeedQueries();
+                Assert.assertTrue(search(fixture.store, "docs").isEmpty());
+                Assert.assertEquals(1, server.connectCalls.get());
+            } finally {
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldReleaseAndReuseClientAfterThreadInterrupt() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 5_000L)) {
+            fixture.manager.markCollectionLoaded("docs");
+            fixture.manager.withClient(client -> client);
+            QueryBlock block = server.blockQueries();
+            CountDownLatch finished = new CountDownLatch(1);
+            AtomicReference<Throwable> failure = new AtomicReference<>();
+            AtomicBoolean interrupted = new AtomicBoolean();
+            Thread searchThread = new Thread(() -> {
+                try {
+                    search(fixture.store, "docs");
+                } catch (Throwable exception) {
+                    failure.set(exception);
+                } finally {
+                    interrupted.set(Thread.currentThread().isInterrupted());
+                    finished.countDown();
+                }
+            }, "milvus-interrupt-test");
+            searchThread.start();
+
+            block.awaitEntered();
+            searchThread.interrupt();
+            Assert.assertTrue(finished.await(AWAIT_SECONDS, TimeUnit.SECONDS));
+            block.awaitCancelled();
+
+            Assert.assertNotNull(failure.get());
+            Assert.assertTrue(interrupted.get());
+            Assert.assertEquals(0, fixture.manager.getActiveClientCount());
+            Assert.assertEquals(1, fixture.manager.getIdleClientCount());
+
+            server.succeedQueries();
+            Assert.assertTrue(search(fixture.store, "docs").isEmpty());
+            Assert.assertEquals(1, server.connectCalls.get());
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldKeepPoolAfterOrdinaryBusinessFailure() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 2_000L)) {
+            MilvusClientV2 first = fixture.manager.withClient(client -> client);
+
+            try {
+                fixture.manager.withClient(client -> {
+                    throw new IllegalArgumentException("synthetic business failure");
+                });
+                Assert.fail("The business failure must be propagated");
+            } catch (IllegalArgumentException expected) {
+                Assert.assertEquals("synthetic business failure", expected.getMessage());
+            }
+
+            MilvusClientV2 second = fixture.manager.withClient(client -> client);
+            Assert.assertSame(first, second);
+            Assert.assertEquals(1, server.connectCalls.get());
+            Assert.assertEquals(0, fixture.manager.getActiveClientCount());
+            Assert.assertEquals(1, fixture.manager.getIdleClientCount());
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldAttachContextToEveryClientOperation() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 2_000L)) {
+            Context callerContext = Context.current();
+            Context operationContext = fixture.manager.withClient(client ->
+                Context.current());
+
+            Assert.assertNotSame(callerContext, operationContext);
+            Assert.assertTrue(operationContext.isCancelled());
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldCapPoolWaitByRemainingSearchDeadline() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 5_000L)) {
+            fixture.manager.markCollectionLoaded("docs");
+            ExecutorService executor = Executors.newSingleThreadExecutor();
+            CountDownLatch borrowed = new CountDownLatch(1);
+            CountDownLatch release = new CountDownLatch(1);
+            try {
+                Future<?> holder = executor.submit(() ->
+                    fixture.manager.withClient(client -> {
+                        borrowed.countDown();
+                        try {
+                            Assert.assertTrue(release.await(
+                                AWAIT_SECONDS, TimeUnit.SECONDS));
+                        } catch (InterruptedException exception) {
+                            Thread.currentThread().interrupt();
+                            throw new IllegalStateException(exception);
+                        }
+                        return null;
+                    }));
+                Assert.assertTrue(borrowed.await(
+                    AWAIT_SECONDS, TimeUnit.SECONDS));
+
+                StoreOptions options = StoreOptions.ofCollectionName("docs");
+                options.setTimeoutMillis(500L);
+                long startedAt = System.nanoTime();
+                try {
+                    search(fixture.store, options);
+                    Assert.fail("Pool wait must respect the remaining deadline");
+                } catch (RuntimeException expected) {
+                    Assert.assertTrue(expected.toString(),
+                        expected instanceof StoreTimeoutException);
+                    long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(
+                        System.nanoTime() - startedAt);
+                    Assert.assertTrue("elapsedMillis=" + elapsedMillis,
+                        elapsedMillis < 800L);
+                }
+
+                release.countDown();
+                holder.get(AWAIT_SECONDS, TimeUnit.SECONDS);
+                Assert.assertEquals(0, fixture.manager.getActiveClientCount());
+            } finally {
+                release.countDown();
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldInvalidateOnlyClosedClient() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 2, 2_000L)) {
+            ExecutorService executor = Executors.newSingleThreadExecutor();
+            CountDownLatch firstBorrowed = new CountDownLatch(1);
+            CountDownLatch releaseFirst = new CountDownLatch(1);
+            try {
+                Future<?> holder = executor.submit(() ->
+                    fixture.manager.withClient(client -> {
+                        firstBorrowed.countDown();
+                        try {
+                            Assert.assertTrue(releaseFirst.await(
+                                AWAIT_SECONDS, TimeUnit.SECONDS));
+                        } catch (InterruptedException exception) {
+                            Thread.currentThread().interrupt();
+                            throw new IllegalStateException(exception);
+                        }
+                        return null;
+                    }));
+                Assert.assertTrue(firstBorrowed.await(
+                    AWAIT_SECONDS, TimeUnit.SECONDS));
+                fixture.manager.withClient(client -> client);
+                releaseFirst.countDown();
+                holder.get(AWAIT_SECONDS, TimeUnit.SECONDS);
+                Assert.assertEquals(2, fixture.manager.getIdleClientCount());
+
+                AtomicReference<MilvusClientV2> closed = new AtomicReference<>();
+                try {
+                    fixture.manager.withClient(client -> {
+                        closed.set(client);
+                        client.close();
+                        throw new IllegalStateException("synthetic closed client");
+                    });
+                    Assert.fail("The closed-client failure must be propagated");
+                } catch (IllegalStateException expected) {
+                    Assert.assertEquals(
+                        "synthetic closed client", expected.getMessage());
+                }
+
+                Assert.assertEquals(1, fixture.manager.getIdleClientCount());
+                MilvusClientV2 remaining = fixture.manager.withClient(
+                    client -> client);
+                Assert.assertNotSame(closed.get(), remaining);
+                Assert.assertEquals(2, server.connectCalls.get());
+            } finally {
+                releaseFirst.countDown();
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldLoadSameCollectionOnlyOnce() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 2, 3_000L)) {
+            LoadGate gate = server.blockLoad("shared");
+            ExecutorService executor = Executors.newFixedThreadPool(2);
+            try {
+                Future<List<Document>> first = executor.submit(() ->
+                    search(fixture.store, "shared"));
+                gate.awaitEntered();
+
+                Future<List<Document>> second = executor.submit(() ->
+                    search(fixture.store, "shared"));
+                MilvusClientManager.CollectionLoadTicket follower =
+                    fixture.manager.beginCollectionLoad("shared");
+                Assert.assertFalse(follower.isLeader());
+                awaitCondition(() -> follower.completion().getNumberOfDependents() > 0);
+
+                Assert.assertEquals(1, server.loadCalls("shared"));
+                Assert.assertEquals(1, fixture.manager.getActiveClientCount());
+                gate.release();
+
+                Assert.assertTrue(first.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
+                Assert.assertTrue(second.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
+                Assert.assertEquals(1, server.loadCalls("shared"));
+                Assert.assertTrue(fixture.manager.isCollectionLoaded("shared"));
+            } finally {
+                gate.release();
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldRemoveLeaderTicketAfterLateCacheHit() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer()) {
+            MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
+            config.setUri(server.uri());
+            config.setDefaultCollectionName("docs");
+            config.setPoolMinIdlePerKey(0);
+            config.setSearchTimeoutMillis(2_000L);
+            RacingMilvusClientManager manager =
+                new RacingMilvusClientManager(config, "late-hit");
+            MilvusVectorStore store = new MilvusVectorStore(config, manager);
+            try {
+                Assert.assertTrue(search(store, "late-hit").isEmpty());
+                manager.markCollectionUnloaded("late-hit");
+
+                Assert.assertTrue(search(store, "late-hit").isEmpty());
+
+                Assert.assertEquals(1, server.loadCalls("late-hit"));
+                Assert.assertTrue(manager.isCollectionLoaded("late-hit"));
+            } finally {
+                store.close();
+                manager.close();
+            }
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldLoadDifferentCollectionsInParallel() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 2, 3_000L)) {
+            LoadGate firstGate = server.blockLoad("first");
+            LoadGate secondGate = server.blockLoad("second");
+            ExecutorService executor = Executors.newFixedThreadPool(2);
+            try {
+                Future<List<Document>> first = executor.submit(() ->
+                    search(fixture.store, "first"));
+                Future<List<Document>> second = executor.submit(() ->
+                    search(fixture.store, "second"));
+
+                firstGate.awaitEntered();
+                secondGate.awaitEntered();
+                Assert.assertEquals(2, server.activeLoads.get());
+                Assert.assertEquals(2, server.maxConcurrentLoads.get());
+                Assert.assertEquals(2, fixture.manager.getActiveClientCount());
+
+                firstGate.release();
+                secondGate.release();
+                Assert.assertTrue(first.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
+                Assert.assertTrue(second.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
+            } finally {
+                firstGate.release();
+                secondGate.release();
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldReelectFollowerAfterLoadLeaderTimesOut() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 3_000L)) {
+            LoadGate gate = server.blockLoad("reelect");
+            ExecutorService executor = Executors.newFixedThreadPool(2);
+            try {
+                StoreOptions shortBudget =
+                    StoreOptions.ofCollectionName("reelect");
+                shortBudget.setTimeoutMillis(500L);
+                Future<List<Document>> first = executor.submit(() ->
+                    search(fixture.store, shortBudget));
+                gate.awaitEntered();
+
+                Future<List<Document>> second = executor.submit(() ->
+                    search(fixture.store, "reelect"));
+                awaitCondition(() -> server.loadCalls("reelect") == 2);
+                gate.release();
+
+                Throwable firstFailure = futureFailure(first);
+                Assert.assertTrue(firstFailure.toString(),
+                    firstFailure instanceof StoreTimeoutException);
+                Assert.assertTrue(second.get(
+                    AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
+                Assert.assertEquals(2, server.loadCalls("reelect"));
+                Assert.assertTrue(
+                    fixture.manager.isCollectionLoaded("reelect"));
+            } finally {
+                gate.release();
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldCancelWaitingFollowerWithoutBorrowingClient() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 5_000L)) {
+            MilvusClientManager.CollectionLoadTicket leader =
+                fixture.manager.beginCollectionLoad("waiting");
+            CountDownLatch finished = new CountDownLatch(1);
+            AtomicReference<Throwable> failure = new AtomicReference<>();
+            Thread follower = new Thread(() -> {
+                try {
+                    search(fixture.store, "waiting");
+                } catch (Throwable exception) {
+                    failure.set(exception);
+                } finally {
+                    finished.countDown();
+                }
+            }, "milvus-load-follower-test");
+            try {
+                follower.start();
+                awaitCondition(() -> leader.completion().getNumberOfDependents() > 0);
+
+                Assert.assertEquals(0, fixture.manager.getActiveClientCount());
+                Assert.assertEquals(0, server.connectCalls.get());
+                follower.interrupt();
+                Assert.assertTrue(finished.await(AWAIT_SECONDS, TimeUnit.SECONDS));
+
+                Assert.assertNotNull(failure.get());
+                Assert.assertFalse(leader.completion().isDone());
+                Assert.assertEquals(0, fixture.manager.getActiveClientCount());
+                Assert.assertEquals(0, server.connectCalls.get());
+            } finally {
+                follower.interrupt();
+                fixture.manager.failCollectionLoad(
+                    leader, new IllegalStateException("test cleanup"), false);
+                fixture.manager.endCollectionLoad(leader);
+            }
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldRecoverAfterCollectionLoadFailure() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 2_000L)) {
+            server.failNextLoad("recoverable");
+            LoadGate gate = server.blockLoad("recoverable");
+            ExecutorService executor = Executors.newFixedThreadPool(2);
+            try {
+                Future<List<Document>> first = executor.submit(() ->
+                    search(fixture.store, "recoverable"));
+                gate.awaitEntered();
+                Future<List<Document>> second = executor.submit(() ->
+                    search(fixture.store, "recoverable"));
+                MilvusClientManager.CollectionLoadTicket follower =
+                    fixture.manager.beginCollectionLoad("recoverable");
+                Assert.assertFalse(follower.isLeader());
+                awaitCondition(() ->
+                    follower.completion().getNumberOfDependents() > 0);
+                gate.release();
+
+                assertFutureFails(first);
+                assertFutureFails(second);
+            } finally {
+                gate.release();
+                executor.shutdownNow();
+            }
+
+            Assert.assertEquals(1, server.loadCalls("recoverable"));
+            Assert.assertFalse(fixture.manager.isCollectionLoaded("recoverable"));
+            Assert.assertEquals(0, fixture.manager.getActiveClientCount());
+            Assert.assertEquals(1, fixture.manager.getIdleClientCount());
+
+            Assert.assertTrue(search(fixture.store, "recoverable").isEmpty());
+            Assert.assertEquals(2, server.loadCalls("recoverable"));
+            Assert.assertTrue(fixture.manager.isCollectionLoaded("recoverable"));
+            Assert.assertEquals(1, server.connectCalls.get());
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldCancelActiveRpcWhenManagerCloses() throws Exception {
+        FakeMilvusServer server = new FakeMilvusServer();
+        Fixture fixture = new Fixture(server, 1, 5_000L);
+        ExecutorService executor = Executors.newFixedThreadPool(2);
+        try {
+            fixture.manager.markCollectionLoaded("docs");
+            QueryBlock block = server.blockQueries();
+            Future<?> operation = executor.submit(() ->
+                query(fixture.manager));
+            block.awaitEntered();
+
+            Future<?> close = executor.submit(fixture.manager::close);
+            block.awaitCancelled();
+            close.get(AWAIT_SECONDS, TimeUnit.SECONDS);
+            assertFutureFails(operation);
+
+            try {
+                fixture.manager.withClient(client -> null);
+                Assert.fail("A closed manager must reject client borrows");
+            } catch (IllegalStateException expected) {
+                Assert.assertEquals("Milvus client pool is closed", expected.getMessage());
+            }
+        } finally {
+            executor.shutdownNow();
+            fixture.close();
+            server.close();
+        }
+    }
+
+    @Test(timeout = 10_000L)
+    public void shouldCancelActiveRpcWhenManagerReconfigures() throws Exception {
+        try (FakeMilvusServer server = new FakeMilvusServer();
+             Fixture fixture = new Fixture(server, 1, 5_000L)) {
+            ExecutorService executor = Executors.newFixedThreadPool(2);
+            try {
+                QueryBlock block = server.blockQueries();
+                Future<?> operation = executor.submit(() ->
+                    query(fixture.manager));
+                block.awaitEntered();
+
+                fixture.config.setPoolMaxTotal(2);
+                Future<Boolean> reconfigure = executor.submit(() ->
+                    fixture.manager.reconfigureIfNeeded(fixture.config));
+                block.awaitCancelled();
+
+                Assert.assertTrue(reconfigure.get(
+                    AWAIT_SECONDS, TimeUnit.SECONDS));
+                assertFutureFails(operation);
+
+                server.succeedQueries();
+                query(fixture.manager);
+            } finally {
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    private static Object query(MilvusClientManager manager) {
+        return manager.withClient(client -> client.query(
+            QueryReq.builder()
+                .collectionName("docs")
+                .filter("id == \"synthetic-id\"")
+                .outputFields(List.of("id"))
+                .build()
+        ));
+    }
+
+    private static List<Document> search(MilvusVectorStore store, String collection) {
+        return search(store, StoreOptions.ofCollectionName(collection));
+    }
+
+    private static List<Document> search(
+        MilvusVectorStore store,
+        StoreOptions options
+    ) {
+        SearchWrapper wrapper = new SearchWrapper();
+        wrapper.setWithVector(false);
+        wrapper.eq("id", "synthetic-id");
+        return store.search(wrapper, options);
+    }
+
+    private static void assertSearchFails(MilvusVectorStore store, String collection) {
+        try {
+            search(store, collection);
+            Assert.fail("The synthetic Milvus failure must be propagated");
+        } catch (RuntimeException expected) {
+            Assert.assertNotNull(expected);
+        }
+    }
+
+    private static void assertFutureFails(Future<?> future)
+        throws InterruptedException, TimeoutException {
+        futureFailure(future);
+    }
+
+    private static Throwable futureFailure(Future<?> future)
+        throws InterruptedException, TimeoutException {
+        try {
+            future.get(AWAIT_SECONDS, TimeUnit.SECONDS);
+            Assert.fail("The synthetic Milvus failure must be propagated");
+        } catch (ExecutionException expected) {
+            Assert.assertNotNull(expected.getCause());
+            return expected.getCause();
+        }
+        throw new AssertionError("Expected future to fail");
+    }
+
+    private static void awaitCondition(BooleanSupplier condition)
+        throws InterruptedException, TimeoutException {
+        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(AWAIT_SECONDS);
+        while (!condition.getAsBoolean()) {
+            if (System.nanoTime() >= deadline) {
+                throw new TimeoutException("Timed out waiting for test condition");
+            }
+            if (Thread.interrupted()) {
+                throw new InterruptedException();
+            }
+            Thread.onSpinWait();
+        }
+    }
+
+    private static final class Fixture implements AutoCloseable {
+
+        private final MilvusVectorStoreConfig config;
+        private final MilvusClientManager manager;
+        private final MilvusVectorStore store;
+
+        private Fixture(FakeMilvusServer server, int poolSize, long searchTimeoutMillis) {
+            config = new MilvusVectorStoreConfig();
+            config.setUri(server.uri());
+            config.setDefaultCollectionName("docs");
+            config.setPoolMaxTotal(poolSize);
+            config.setPoolMaxTotalPerKey(poolSize);
+            config.setPoolMaxIdlePerKey(poolSize);
+            config.setPoolMinIdlePerKey(0);
+            config.setPoolMaxWaitMillis(1_000L);
+            config.setSearchTimeoutMillis(searchTimeoutMillis);
+            manager = new MilvusClientManager(config);
+            store = new MilvusVectorStore(config, manager);
+        }
+
+        @Override
+        public void close() {
+            store.close();
+            manager.close();
+        }
+    }
+
+    private static final class RacingMilvusClientManager
+        extends MilvusClientManager {
+
+        private final String collectionName;
+        private final AtomicInteger observations = new AtomicInteger();
+
+        private RacingMilvusClientManager(
+            MilvusVectorStoreConfig config,
+            String collectionName
+        ) {
+            super(config);
+            this.collectionName = collectionName;
+        }
+
+        @Override
+        boolean isCollectionLoaded(String requestedCollectionName) {
+            if (collectionName.equals(requestedCollectionName)) {
+                int observation = observations.getAndIncrement();
+                if (observation == 0) {
+                    return false;
+                }
+                if (observation == 1) {
+                    return true;
+                }
+            }
+            return super.isCollectionLoaded(requestedCollectionName);
+        }
+    }
+
+    private static final class QueryBlock {
+
+        private final CountDownLatch entered = new CountDownLatch(1);
+        private final CountDownLatch cancelled = new CountDownLatch(1);
+
+        private void awaitEntered() throws InterruptedException {
+            Assert.assertTrue(entered.await(AWAIT_SECONDS, TimeUnit.SECONDS));
+        }
+
+        private void awaitCancelled() throws InterruptedException {
+            Assert.assertTrue(cancelled.await(AWAIT_SECONDS, TimeUnit.SECONDS));
+        }
+    }
+
+    private static final class LoadGate {
+
+        private final CountDownLatch entered = new CountDownLatch(1);
+        private final CountDownLatch release = new CountDownLatch(1);
+
+        private void awaitEntered() throws InterruptedException {
+            Assert.assertTrue(entered.await(AWAIT_SECONDS, TimeUnit.SECONDS));
+        }
+
+        private void release() {
+            release.countDown();
+        }
+    }
+
+    private static final class FakeMilvusServer implements AutoCloseable {
+
+        private static final io.milvus.grpc.Status SUCCESS =
+            io.milvus.grpc.Status.newBuilder()
+                .setErrorCode(ErrorCode.Success)
+                .setCode(0)
+                .build();
+
+        private final AtomicInteger connectCalls = new AtomicInteger();
+        private final AtomicInteger queryCalls = new AtomicInteger();
+        private final AtomicInteger activeLoads = new AtomicInteger();
+        private final AtomicInteger maxConcurrentLoads = new AtomicInteger();
+        private final AtomicBoolean querySawDeadline = new AtomicBoolean();
+        private final ConcurrentHashMap<String, AtomicInteger> loadCalls =
+            new ConcurrentHashMap<>();
+        private final ConcurrentHashMap<String, LoadGate> loadGates =
+            new ConcurrentHashMap<>();
+        private final Set<String> loadedCollections = ConcurrentHashMap.newKeySet();
+        private final Set<String> failNextLoads = ConcurrentHashMap.newKeySet();
+        private final ExecutorService rpcExecutor = Executors.newCachedThreadPool();
+        private final Server server;
+        private volatile QueryAction queryAction = QueryAction.SUCCESS;
+        private volatile QueryBlock queryBlock;
+
+        private FakeMilvusServer() throws IOException {
+            server = NettyServerBuilder.forPort(0)
+                .executor(rpcExecutor)
+                .addService(new Service())
+                .build()
+                .start();
+        }
+
+        private String uri() {
+            return "http://127.0.0.1:" + server.getPort();
+        }
+
+        private void succeedQueries() {
+            queryAction = QueryAction.SUCCESS;
+            queryBlock = null;
+        }
+
+        private void failQueriesWithUnavailable() {
+            queryAction = QueryAction.UNAVAILABLE;
+            queryBlock = null;
+        }
+
+        private QueryBlock blockQueries() {
+            QueryBlock block = new QueryBlock();
+            queryBlock = block;
+            queryAction = QueryAction.BLOCK;
+            return block;
+        }
+
+        private LoadGate blockLoad(String collectionName) {
+            LoadGate gate = new LoadGate();
+            loadGates.put(collectionName, gate);
+            return gate;
+        }
+
+        private void failNextLoad(String collectionName) {
+            failNextLoads.add(collectionName);
+        }
+
+        private int loadCalls(String collectionName) {
+            AtomicInteger calls = loadCalls.get(collectionName);
+            return calls == null ? 0 : calls.get();
+        }
+
+        @Override
+        public void close() {
+            for (LoadGate gate : new ArrayList<>(loadGates.values())) {
+                gate.release();
+            }
+            server.shutdownNow();
+            try {
+                server.awaitTermination(AWAIT_SECONDS, TimeUnit.SECONDS);
+            } catch (InterruptedException exception) {
+                Thread.currentThread().interrupt();
+            } finally {
+                rpcExecutor.shutdownNow();
+            }
+        }
+
+        private enum QueryAction {
+            SUCCESS,
+            UNAVAILABLE,
+            BLOCK
+        }
+
+        private final class Service extends MilvusServiceGrpc.MilvusServiceImplBase {
+
+            @Override
+            public void connect(
+                ConnectRequest request,
+                StreamObserver<ConnectResponse> observer
+            ) {
+                connectCalls.incrementAndGet();
+                observer.onNext(ConnectResponse.newBuilder()
+                    .setStatus(SUCCESS)
+                    .setIdentifier(1L)
+                    .build());
+                observer.onCompleted();
+            }
+
+            @Override
+            public void listDatabases(
+                ListDatabasesRequest request,
+                StreamObserver<ListDatabasesResponse> observer
+            ) {
+                observer.onNext(ListDatabasesResponse.newBuilder()
+                    .setStatus(SUCCESS)
+                    .addDbNames("default")
+                    .build());
+                observer.onCompleted();
+            }
+
+            @Override
+            public void checkHealth(
+                CheckHealthRequest request,
+                StreamObserver<CheckHealthResponse> observer
+            ) {
+                observer.onNext(CheckHealthResponse.newBuilder()
+                    .setStatus(SUCCESS)
+                    .setIsHealthy(true)
+                    .build());
+                observer.onCompleted();
+            }
+
+            @Override
+            public void getLoadState(
+                GetLoadStateRequest request,
+                StreamObserver<GetLoadStateResponse> observer
+            ) {
+                LoadState state = loadedCollections.contains(request.getCollectionName())
+                    ? LoadState.LoadStateLoaded
+                    : LoadState.LoadStateNotLoad;
+                observer.onNext(GetLoadStateResponse.newBuilder()
+                    .setStatus(SUCCESS)
+                    .setState(state)
+                    .build());
+                observer.onCompleted();
+            }
+
+            @Override
+            public void describeCollection(
+                DescribeCollectionRequest request,
+                StreamObserver<DescribeCollectionResponse> observer
+            ) {
+                CollectionSchema schema = CollectionSchema.newBuilder()
+                    .setName(request.getCollectionName())
+                    .addFields(FieldSchema.newBuilder()
+                        .setName("id")
+                        .setIsPrimaryKey(true)
+                        .setDataType(DataType.VarChar)
+                        .build())
+                    .build();
+                observer.onNext(DescribeCollectionResponse.newBuilder()
+                    .setStatus(SUCCESS)
+                    .setCollectionName(request.getCollectionName())
+                    .setSchema(schema)
+                    .build());
+                observer.onCompleted();
+            }
+
+            @Override
+            public void loadCollection(
+                LoadCollectionRequest request,
+                StreamObserver<io.milvus.grpc.Status> observer
+            ) {
+                String collectionName = request.getCollectionName();
+                loadCalls.computeIfAbsent(
+                    collectionName, ignored -> new AtomicInteger()).incrementAndGet();
+                LoadGate gate = loadGates.get(collectionName);
+                if (gate != null) {
+                    int active = activeLoads.incrementAndGet();
+                    maxConcurrentLoads.accumulateAndGet(active, Math::max);
+                    gate.entered.countDown();
+                    try {
+                        if (!gate.release.await(AWAIT_SECONDS, TimeUnit.SECONDS)) {
+                            observer.onError(io.grpc.Status.DEADLINE_EXCEEDED
+                                .withDescription("test load gate timed out")
+                                .asRuntimeException());
+                            return;
+                        }
+                    } catch (InterruptedException exception) {
+                        Thread.currentThread().interrupt();
+                        observer.onError(io.grpc.Status.CANCELLED
+                            .withCause(exception)
+                            .asRuntimeException());
+                        return;
+                    } finally {
+                        activeLoads.decrementAndGet();
+                    }
+                }
+                if (failNextLoads.remove(collectionName)) {
+                    observer.onNext(io.milvus.grpc.Status.newBuilder()
+                        .setErrorCode(ErrorCode.UnexpectedError)
+                        .setCode(1)
+                        .setReason("synthetic load failure")
+                        .build());
+                    observer.onCompleted();
+                    return;
+                }
+                loadedCollections.add(collectionName);
+                observer.onNext(SUCCESS);
+                observer.onCompleted();
+            }
+
+            @Override
+            public void query(
+                QueryRequest request,
+                StreamObserver<QueryResults> observer
+            ) {
+                queryCalls.incrementAndGet();
+                querySawDeadline.compareAndSet(
+                    false, Context.current().getDeadline() != null);
+                QueryAction action = queryAction;
+                if (action == QueryAction.UNAVAILABLE) {
+                    observer.onError(io.grpc.Status.UNAVAILABLE
+                        .withDescription("synthetic query failure")
+                        .asRuntimeException());
+                    return;
+                }
+                if (action == QueryAction.BLOCK) {
+                    QueryBlock block = queryBlock;
+                    @SuppressWarnings("unchecked")
+                    ServerCallStreamObserver<QueryResults> serverObserver =
+                        (ServerCallStreamObserver<QueryResults>) observer;
+                    serverObserver.setOnCancelHandler(block.cancelled::countDown);
+                    block.entered.countDown();
+                    return;
+                }
+                observer.onNext(QueryResults.newBuilder()
+                    .setStatus(SUCCESS)
+                    .setCollectionName(request.getCollectionName())
+                    .build());
+                observer.onCompleted();
+            }
+        }
+    }
+}

easy-agents-store/easy-agents-store-milvus/src/test/java/com/easyagents/store/milvus/MilvusVectorStoreIntegrationTest.java

@@ -63,7 +63,7 @@ public class MilvusVectorStoreIntegrationTest {
 
             assertQueryFailureIsNotReportedAsEmpty(store, options);
             assertPoolExhaustionIsBounded(manager);
-            assertFailedClientIsNotReused(manager);
+            assertBusinessFailurePreservesClient(manager);
             Assert.assertTrue(store.checkAvailable());
         } finally {
             try {
@@ -104,10 +104,11 @@ public class MilvusVectorStoreIntegrationTest {
         }
     }
 
-    private static void assertFailedClientIsNotReused(MilvusClientManager manager)
+    private static void assertBusinessFailurePreservesClient(MilvusClientManager manager)
         throws InterruptedException, ExecutionException {
         ExecutorService executor = Executors.newFixedThreadPool(2);
         CountDownLatch borrowed = new CountDownLatch(1);
+        CountDownLatch waiterStarted = new CountDownLatch(1);
         CountDownLatch fail = new CountDownLatch(1);
         AtomicReference<Object> failedClient = new AtomicReference<>();
         try {
@@ -132,11 +133,14 @@ public class MilvusVectorStoreIntegrationTest {
                 }
             });
             Assert.assertTrue(borrowed.await(2, TimeUnit.SECONDS));
-            Future<Object> waiting = executor.submit(() -> manager.withClient(client -> client));
-            Thread.sleep(100L);
+            Future<Object> waiting = executor.submit(() -> {
+                waiterStarted.countDown();
+                return manager.withClient(client -> client);
+            });
+            Assert.assertTrue(waiterStarted.await(2, TimeUnit.SECONDS));
             fail.countDown();
             failing.get();
-            Assert.assertNotSame(failedClient.get(), waiting.get());
+            Assert.assertSame(failedClient.get(), waiting.get());
         } finally {
             fail.countDown();
             executor.shutdownNow();