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

easy-agents-core/src/main/java/com/easyagents/core/store/StoreOptions.java

@@ -43,6 +43,11 @@ public class StoreOptions extends Metadata {
         public void setEmbeddingOptions(EmbeddingOptions embeddingOptions) {
             throw new IllegalStateException("Can not set embeddingOptions to the default instance.");
         }
+
+        @Override
+        public void setTimeoutMillis(Long timeoutMillis) {
+            throw new IllegalStateException("Can not set timeoutMillis to the default instance.");
+        }
     };
 
     /**
@@ -65,6 +70,11 @@ public class StoreOptions extends Metadata {
      */
     private EmbeddingOptions embeddingOptions = EmbeddingOptions.DEFAULT;
 
+    /**
+     * Optional upper bound for one store operation.
+     */
+    private Long timeoutMillis;
+
 
     public String getCollectionName() {
         return collectionName;
@@ -111,6 +121,17 @@ public class StoreOptions extends Metadata {
         this.embeddingOptions = embeddingOptions;
     }
 
+    public Long getTimeoutMillis() {
+        return timeoutMillis;
+    }
+
+    public void setTimeoutMillis(Long timeoutMillis) {
+        if (timeoutMillis != null && timeoutMillis <= 0L) {
+            throw new IllegalArgumentException("timeoutMillis must be greater than zero");
+        }
+        this.timeoutMillis = timeoutMillis;
+    }
+
 
     public static StoreOptions ofCollectionName(String collectionName) {
         StoreOptions storeOptions = new StoreOptions();

easy-agents-core/src/main/java/com/easyagents/core/store/StoreTimeoutException.java

@@ -0,0 +1,21 @@
+/*
+ *  Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ */
+package com.easyagents.core.store;
+
+/**
+ * Indicates that a store operation exhausted its caller-provided time budget.
+ */
+public class StoreTimeoutException extends RuntimeException {
+
+    public StoreTimeoutException(String message) {
+        super(message);
+    }
+
+    public StoreTimeoutException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeManager.java

@@ -18,6 +18,7 @@ package com.easyagents.flow.core.knowledge;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 
 public class KnowledgeManager {
 
@@ -51,4 +52,20 @@ public class KnowledgeManager {
         }
         return null;
     }
+
+    /**
+     * 将完整知识检索请求交给首个能够处理它的 Provider。
+     *
+     * @param request 检索请求
+     * @return 节点输出；没有 Provider 能处理时返回 null
+     */
+    public Map<String, Object> search(KnowledgeSearchRequest request) {
+        for (KnowledgeProvider provider : providers) {
+            Map<String, Object> result = provider.search(request);
+            if (result != null) {
+                return result;
+            }
+        }
+        return null;
+    }
 }

easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeProvider.java

@@ -15,6 +15,37 @@
  */
 package com.easyagents.flow.core.knowledge;
 
+import com.easyagents.flow.core.util.Maps;
+
+import java.util.List;
+import java.util.Map;
+
 public interface KnowledgeProvider {
     Knowledge getKnowledge(Object id);
+
+    /**
+     * 执行完整的知识库节点检索请求。
+     *
+     * <p>默认实现保留单知识库兼容。需要跨知识库汇总的业务 Provider
+     * 应覆盖本方法并返回完整节点输出。</p>
+     *
+     * @param request 检索请求
+     * @return 节点输出；当前 Provider 不支持该请求时返回 null
+     */
+    default Map<String, Object> search(KnowledgeSearchRequest request) {
+        if (request == null || request.getKnowledgeIds().size() != 1) {
+            return null;
+        }
+        Object knowledgeId = request.getKnowledgeIds().get(0);
+        Knowledge knowledge = getKnowledge(knowledgeId);
+        if (knowledge == null) {
+            return null;
+        }
+        List<Map<String, Object>> documents = knowledge.search(
+                request.getKeyword(),
+                request.getLimit(),
+                request.getKnowledgeNode(),
+                request.getChain());
+        return Maps.of("documents", documents);
+    }
 }

easy-agents-flow/src/main/java/com/easyagents/flow/core/knowledge/KnowledgeSearchRequest.java

@@ -0,0 +1,67 @@
+/**
+ * Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com).
+ * <p>
+ * Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0.
+ */
+package com.easyagents.flow.core.knowledge;
+
+import com.easyagents.flow.core.chain.Chain;
+import com.easyagents.flow.core.node.KnowledgeNode;
+
+import java.util.Collections;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * 工作流知识库节点的完整检索请求。
+ */
+public class KnowledgeSearchRequest {
+
+    private final List<Object> knowledgeIds;
+    private final String keyword;
+    private final int limit;
+    private final String retrievalMode;
+    private final KnowledgeNode knowledgeNode;
+    private final Chain chain;
+
+    public KnowledgeSearchRequest(
+            List<Object> knowledgeIds,
+            String keyword,
+            int limit,
+            String retrievalMode,
+            KnowledgeNode knowledgeNode,
+            Chain chain) {
+        this.knowledgeIds = knowledgeIds == null
+                ? Collections.emptyList()
+                : Collections.unmodifiableList(new ArrayList<>(knowledgeIds));
+        this.keyword = keyword;
+        this.limit = limit;
+        this.retrievalMode = retrievalMode;
+        this.knowledgeNode = knowledgeNode;
+        this.chain = chain;
+    }
+
+    public List<Object> getKnowledgeIds() {
+        return knowledgeIds;
+    }
+
+    public String getKeyword() {
+        return keyword;
+    }
+
+    public int getLimit() {
+        return limit;
+    }
+
+    public String getRetrievalMode() {
+        return retrievalMode;
+    }
+
+    public KnowledgeNode getKnowledgeNode() {
+        return knowledgeNode;
+    }
+
+    public Chain getChain() {
+        return chain;
+    }
+}

easy-agents-flow/src/main/java/com/easyagents/flow/core/node/KnowledgeNode.java

@@ -17,15 +17,15 @@ package com.easyagents.flow.core.node;
 
 import com.easyagents.flow.core.chain.Chain;
 import com.easyagents.flow.core.chain.ChainState;
-import com.easyagents.flow.core.knowledge.Knowledge;
 import com.easyagents.flow.core.knowledge.KnowledgeManager;
-import com.easyagents.flow.core.util.Maps;
+import com.easyagents.flow.core.knowledge.KnowledgeSearchRequest;
 import com.easyagents.flow.core.util.StringUtil;
 import com.easyagents.flow.core.util.TextTemplate;
 import org.slf4j.Logger;
 
-import java.util.Arrays;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 
@@ -36,6 +36,7 @@ public class KnowledgeNode extends BaseNode {
     private static final Logger logger = org.slf4j.LoggerFactory.getLogger(KnowledgeNode.class);
 
     private Object knowledgeId;
+    private List<Object> knowledgeIds = new ArrayList<>();
     private String keyword;
     private String limit;
     private String retrievalMode = "HYBRID";
@@ -48,6 +49,32 @@ public class KnowledgeNode extends BaseNode {
         this.knowledgeId = knowledgeId;
     }
 
+    /**
+     * 获取规范化的知识库集合，兼容历史单值字段。
+     *
+     * @return 去重后的知识库 ID
+     */
+    public List<Object> getKnowledgeIds() {
+        if (knowledgeIds != null && !knowledgeIds.isEmpty()) {
+            return Collections.unmodifiableList(knowledgeIds);
+        }
+        return knowledgeId == null
+                ? Collections.emptyList()
+                : Collections.singletonList(knowledgeId);
+    }
+
+    public void setKnowledgeIds(List<?> knowledgeIds) {
+        LinkedHashSet<Object> normalized = new LinkedHashSet<>();
+        if (knowledgeIds != null) {
+            for (Object id : knowledgeIds) {
+                if (id != null && StringUtil.hasText(String.valueOf(id))) {
+                    normalized.add(id);
+                }
+            }
+        }
+        this.knowledgeIds = new ArrayList<>(normalized);
+    }
+
     public String getKeyword() {
         return keyword;
     }
@@ -88,25 +115,44 @@ public class KnowledgeNode extends BaseNode {
         if (StringUtil.hasText(realLimitString)) {
             try {
                 realLimit = Integer.parseInt(realLimitString);
-            } catch (Exception e) {
-                logger.error(e.toString(), e);
+            } catch (NumberFormatException exception) {
+                throw new IllegalArgumentException(
+                        "知识库节点最终返回条数必须为正整数", exception);
             }
         }
-
-        Knowledge knowledge = KnowledgeManager.getInstance().getKnowledge(knowledgeId);
-
-        if (knowledge == null) {
-            return Collections.emptyMap();
+        if (realLimit <= 0) {
+            throw new IllegalArgumentException(
+                    "知识库节点最终返回条数必须为正整数");
         }
 
-        List<Map<String, Object>> result = knowledge.search(realKeyword, realLimit, this, chain);
-        return Maps.of("documents", result);
+        List<Object> resolvedKnowledgeIds = getKnowledgeIds();
+        if (resolvedKnowledgeIds.isEmpty()) {
+            throw new IllegalArgumentException("知识库节点至少需要选择一个知识库");
+        }
+        if (resolvedKnowledgeIds.size() > 1
+                && !"VECTOR".equalsIgnoreCase(retrievalMode)) {
+            throw new IllegalArgumentException("多知识库检索仅支持 VECTOR 模式");
+        }
+
+        Map<String, Object> result = KnowledgeManager.getInstance().search(
+                new KnowledgeSearchRequest(
+                        resolvedKnowledgeIds,
+                        realKeyword,
+                        realLimit,
+                        retrievalMode,
+                        this,
+                        chain));
+        if (result == null) {
+            throw new IllegalStateException("没有可用的知识库 Provider");
+        }
+        return result;
     }
 
     @Override
     public String toString() {
         return "KnowledgeNode{" +
                 "knowledgeId=" + knowledgeId +
+                ", knowledgeIds=" + knowledgeIds +
                 ", keyword='" + keyword + '\'' +
                 ", limit='" + limit + '\'' +
                 ", retrievalMode='" + retrievalMode + '\'' +

easy-agents-flow/src/main/java/com/easyagents/flow/core/parser/impl/KnowledgeNodeParser.java

@@ -16,15 +16,41 @@
 package com.easyagents.flow.core.parser.impl;
 
 import com.alibaba.fastjson.JSONObject;
+import com.alibaba.fastjson.JSONArray;
 import com.easyagents.flow.core.node.KnowledgeNode;
 import com.easyagents.flow.core.parser.BaseNodeParser;
 
+import java.util.ArrayList;
+
 public class KnowledgeNodeParser extends BaseNodeParser<KnowledgeNode> {
 
     @Override
     public KnowledgeNode doParse(JSONObject root, JSONObject data, JSONObject chainJSONObject) {
         KnowledgeNode knowledgeNode = new KnowledgeNode();
-        knowledgeNode.setKnowledgeId(data.get("knowledgeId"));
+        if (data.containsKey("knowledgeIds")) {
+            Object rawIds = data.get("knowledgeIds");
+            if (!(rawIds instanceof JSONArray)) {
+                throw new IllegalArgumentException("knowledgeIds 必须为数组");
+            }
+            JSONArray ids = (JSONArray) rawIds;
+            if (ids.isEmpty()) {
+                throw new IllegalArgumentException("knowledgeIds 不能为空");
+            }
+            java.util.LinkedHashSet<String> normalized =
+                    new java.util.LinkedHashSet<>();
+            for (Object id : ids) {
+                String value = id == null ? null : String.valueOf(id).trim();
+                if (!com.easyagents.flow.core.util.StringUtil.hasText(value)) {
+                    throw new IllegalArgumentException("knowledgeIds 不能包含空值");
+                }
+                if (!normalized.add(value)) {
+                    throw new IllegalArgumentException("knowledgeIds 不能包含重复值");
+                }
+            }
+            knowledgeNode.setKnowledgeIds(new ArrayList<>(normalized));
+        } else {
+            knowledgeNode.setKnowledgeId(data.get("knowledgeId"));
+        }
         knowledgeNode.setLimit(data.getString("limit"));
         knowledgeNode.setKeyword(data.getString("keyword"));
         knowledgeNode.setRetrievalMode(data.getString("retrievalMode"));

easy-agents-flow/src/test/java/com/easyagents/flow/core/test/KnowledgeNodeTest.java

@@ -0,0 +1,217 @@
+/**
+ * Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com).
+ * <p>
+ * Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0.
+ */
+package com.easyagents.flow.core.test;
+
+import com.alibaba.fastjson.JSONArray;
+import com.alibaba.fastjson.JSONObject;
+import com.easyagents.flow.core.chain.Chain;
+import com.easyagents.flow.core.chain.ChainDefinition;
+import com.easyagents.flow.core.chain.ChainState;
+import com.easyagents.flow.core.knowledge.Knowledge;
+import com.easyagents.flow.core.knowledge.KnowledgeManager;
+import com.easyagents.flow.core.knowledge.KnowledgeProvider;
+import com.easyagents.flow.core.knowledge.KnowledgeSearchRequest;
+import com.easyagents.flow.core.node.KnowledgeNode;
+import com.easyagents.flow.core.parser.impl.KnowledgeNodeParser;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * 知识库节点多来源契约测试。
+ */
+public class KnowledgeNodeTest {
+
+    @Test
+    public void shouldParseLegacyKnowledgeId() {
+        JSONObject data = baseData();
+        data.put("knowledgeId", "101");
+
+        KnowledgeNode node = parse(data);
+
+        Assert.assertEquals(List.of("101"), node.getKnowledgeIds());
+    }
+
+    @Test
+    public void shouldPreferKnowledgeIdsAndKeepOrder() {
+        JSONObject data = baseData();
+        data.put("knowledgeId", "legacy");
+        JSONArray ids = new JSONArray();
+        ids.addAll(List.of("201", "202"));
+        data.put("knowledgeIds", ids);
+
+        KnowledgeNode node = parse(data);
+
+        Assert.assertEquals(List.of("201", "202"), node.getKnowledgeIds());
+    }
+
+    @Test
+    public void shouldNormalizeKnowledgeIdsBeforeStoringThem() {
+        JSONObject data = baseData();
+        JSONArray ids = new JSONArray();
+        ids.addAll(List.of(" 201 ", "202"));
+        data.put("knowledgeIds", ids);
+
+        KnowledgeNode node = parse(data);
+
+        Assert.assertEquals(List.of("201", "202"), node.getKnowledgeIds());
+    }
+
+    @Test
+    public void shouldRejectInvalidKnowledgeIds() {
+        JSONObject data = baseData();
+        data.put("knowledgeIds", "[201,202]");
+        assertParseFailure(data, "必须为数组");
+
+        JSONArray duplicateIds = new JSONArray();
+        duplicateIds.addAll(List.of("201", "201"));
+        data.put("knowledgeIds", duplicateIds);
+        assertParseFailure(data, "重复值");
+    }
+
+    @Test
+    public void defaultProviderShouldKeepSingleKnowledgeCompatibility() {
+        KnowledgeProvider provider = id ->
+                (keyword, limit, node, chain) -> List.of(Map.of(
+                        "knowledgeId", id,
+                        "content", keyword));
+        KnowledgeNode node = new KnowledgeNode();
+        node.setKnowledgeId("301");
+
+        Map<String, Object> output = provider.search(
+                new KnowledgeSearchRequest(
+                        node.getKnowledgeIds(),
+                        "问题",
+                        3,
+                        "HYBRID",
+                        node,
+                        null));
+
+        Assert.assertNotNull(output);
+        Assert.assertEquals(1, ((List<?>) output.get("documents")).size());
+    }
+
+    @Test
+    public void defaultProviderShouldDeclineMultiKnowledgeRequest() {
+        KnowledgeProvider provider = id -> null;
+        KnowledgeNode node = new KnowledgeNode();
+        node.setKnowledgeIds(List.of("401", "402"));
+
+        Assert.assertNull(provider.search(new KnowledgeSearchRequest(
+                node.getKnowledgeIds(),
+                "问题",
+                3,
+                "VECTOR",
+                node,
+                null)));
+    }
+
+    @Test
+    public void shouldResolveVariableLimitAndDefaultBlankValueAtRuntime() {
+        Assert.assertEquals(7, executeAndCaptureLimit("{{start.limit}}", "7"));
+        Assert.assertEquals(10, executeAndCaptureLimit("{{start.limit}}", "   "));
+        Assert.assertEquals(10, executeAndCaptureLimit("{{start.limit ?? }}", null));
+    }
+
+    @Test
+    public void shouldRejectInvalidResolvedVariableLimitAtRuntime() {
+        assertRuntimeLimitFailure("abc");
+        assertRuntimeLimitFailure("0");
+        assertRuntimeLimitFailure("-2");
+    }
+
+    private static KnowledgeNode parse(JSONObject data) {
+        return new KnowledgeNodeParser().doParse(
+                new JSONObject(), data, new JSONObject());
+    }
+
+    private static JSONObject baseData() {
+        JSONObject data = new JSONObject();
+        data.put("keyword", "问题");
+        data.put("limit", "5");
+        data.put("retrievalMode", "VECTOR");
+        return data;
+    }
+
+    private static void assertParseFailure(
+            JSONObject data, String expectedMessage) {
+        try {
+            parse(data);
+            Assert.fail("invalid knowledgeIds must be rejected");
+        } catch (IllegalArgumentException expected) {
+            Assert.assertTrue(expected.getMessage().contains(expectedMessage));
+        }
+    }
+
+    private static int executeAndCaptureLimit(
+            String limitTemplate,
+            String runtimeValue) {
+        AtomicInteger capturedLimit = new AtomicInteger(-1);
+        KnowledgeProvider provider = new KnowledgeProvider() {
+            @Override
+            public Knowledge getKnowledge(Object id) {
+                return null;
+            }
+
+            @Override
+            public Map<String, Object> search(KnowledgeSearchRequest request) {
+                capturedLimit.set(request.getLimit());
+                return Map.of("documents", List.of());
+            }
+        };
+        KnowledgeManager.getInstance().registerProvider(provider);
+        try {
+            KnowledgeNode node = runtimeNode(limitTemplate);
+            ChainState state = new ChainState();
+            if (runtimeValue != null) {
+                state.getMemory().put("start.limit", runtimeValue);
+            }
+            node.execute(new FixedStateChain(state));
+            return capturedLimit.get();
+        } finally {
+            KnowledgeManager.getInstance().removeProvider(provider);
+        }
+    }
+
+    private static void assertRuntimeLimitFailure(String runtimeValue) {
+        KnowledgeNode node = runtimeNode("{{start.limit}}");
+        ChainState state = new ChainState();
+        state.getMemory().put("start.limit", runtimeValue);
+
+        IllegalArgumentException exception = Assert.assertThrows(
+                IllegalArgumentException.class,
+                () -> node.execute(new FixedStateChain(state)));
+
+        Assert.assertTrue(exception.getMessage().contains("必须为正整数"));
+    }
+
+    private static KnowledgeNode runtimeNode(String limitTemplate) {
+        KnowledgeNode node = new KnowledgeNode();
+        node.setKnowledgeIds(List.of("501", "502"));
+        node.setKeyword("问题");
+        node.setLimit(limitTemplate);
+        node.setRetrievalMode("VECTOR");
+        return node;
+    }
+
+    private static final class FixedStateChain extends Chain {
+
+        private final ChainState state;
+
+        private FixedStateChain(ChainState state) {
+            super(new ChainDefinition(), "knowledge-node-limit-test");
+            this.state = state;
+        }
+
+        @Override
+        public ChainState getExecutionState() {
+            return state;
+        }
+    }
+}

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();