feat(mcp): support document-scoped knowledge search

internal/mcpserver/scope.go

@@ -172,6 +172,44 @@ func searchTargetsFor(kbs []*types.KnowledgeBase) types.SearchTargets {
 	return targets
 }
 
+// documentSearchTargets narrows an already authorized KB selection to explicit
+// documents. Never leave whole-KB targets for bases with no selected documents.
+func (s *Server) documentSearchTargets(
+	ctx context.Context, kbs []*types.KnowledgeBase, knowledgeIDs []string,
+) (types.SearchTargets, error) {
+	if len(knowledgeIDs) == 0 {
+		return nil, fmt.Errorf("knowledge_ids must contain at least one document id; " +
+			"omit it to search whole knowledge bases")
+	}
+	targets := make(types.SearchTargets, 0)
+	byKB := make(map[string]*types.SearchTarget)
+	seen := make(map[string]bool, len(knowledgeIDs))
+	for _, id := range knowledgeIDs {
+		id = strings.TrimSpace(id)
+		if id == "" {
+			return nil, fmt.Errorf("knowledge_ids must not contain blank document ids")
+		}
+		if seen[id] {
+			continue
+		}
+		k, kb, err := s.knowledgeInKnowledgeBases(ctx, kbs, id)
+		if err != nil {
+			return nil, err
+		}
+		target := byKB[kb.ID]
+		if target == nil {
+			target = &types.SearchTarget{
+				Type: types.SearchTargetTypeKnowledge, KnowledgeBaseID: kb.ID, TenantID: kb.TenantID,
+			}
+			byKB[kb.ID] = target
+			targets = append(targets, target)
+		}
+		target.KnowledgeIDs = append(target.KnowledgeIDs, k.ID)
+		seen[id] = true
+	}
+	return targets, nil
+}
+
 func knowledgeBaseIDs(kbs []*types.KnowledgeBase) []string {
 	ids := make([]string, 0, len(kbs))
 	for _, kb := range kbs {
@@ -213,6 +251,18 @@ func wikiScopesFor(kbs []*types.KnowledgeBase) []tools.WikiScope {
 // then match an allowed knowledge base and that knowledge base's tenant.
 func (s *Server) knowledgeInScope(
 	ctx context.Context, ep *types.MCPEndpoint, knowledgeID string,
+) (*types.Knowledge, *types.KnowledgeBase, error) {
+	allowed, err := s.allowedKnowledgeBases(ctx, ep)
+	if err != nil {
+		return nil, nil, err
+	}
+	return s.knowledgeInKnowledgeBases(ctx, allowed, knowledgeID)
+}
+
+// knowledgeInKnowledgeBases applies the same document/owner checks to either
+// the full endpoint scope or a narrower, already authorized KB selection.
+func (s *Server) knowledgeInKnowledgeBases(
+	ctx context.Context, kbs []*types.KnowledgeBase, knowledgeID string,
 ) (*types.Knowledge, *types.KnowledgeBase, error) {
 	knowledgeID = strings.TrimSpace(knowledgeID)
 	if knowledgeID == "" {
@@ -222,13 +272,16 @@ func (s *Server) knowledgeInScope(
 	if err != nil || k == nil {
 		return nil, nil, fmt.Errorf("document %q was not found", knowledgeID)
 	}
-	allowed, err := s.allowedKnowledgeBases(ctx, ep)
-	if err != nil {
-		return nil, nil, err
+	var kb *types.KnowledgeBase
+	for _, candidate := range kbs {
+		if candidate.ID == k.KnowledgeBaseID {
+			kb = candidate
+			break
+		}
 	}
-	kb := matchKnowledgeBase(allowed, k.KnowledgeBaseID)
 	if kb == nil || kb.TenantID != k.TenantID {
-		return nil, nil, fmt.Errorf("document %q is outside this endpoint's scope", knowledgeID)
+		return nil, nil, fmt.Errorf("document %q is outside this endpoint's scope or the selected knowledge bases",
+			knowledgeID)
 	}
 	return k, kb, nil
 }

internal/mcpserver/tools_retrieve.go

@@ -53,6 +53,10 @@ func searchKnowledgeTool() mcp.Tool {
 		mcp.WithArray("knowledge_base_ids", mcp.WithStringItems(),
 			mcp.Description("Optional knowledge base ids or names to restrict the search; defaults to every "+
 				"knowledge base in scope")),
+		mcp.WithArray("knowledge_ids", mcp.MinItems(1), mcp.WithStringItems(mcp.MinLength(1)),
+			mcp.Description("Optional non-empty list of document ids from list_documents or search results. "+
+				"Restricts retrieval before ranking; every document must be in scope and belong to the selected "+
+				"knowledge bases. Omit to search whole knowledge bases; invalid or conflicting ids return an error.")),
 		readOnlyEndpointAnnotations(),
 	)
 }
@@ -165,31 +169,36 @@ func (s *Server) handleSearchKnowledge(ctx context.Context, req mcp.CallToolRequ
 		return mcp.NewToolResultError(
 			"none of the selected knowledge bases supports semantic or keyword retrieval"), nil
 	}
-	return s.runSearchKnowledge(ctx, kbs, strings.TrimSpace(query),
-		req.GetString("mode", tools.SearchModeHybrid), req.GetInt("limit", 0))
+	targets := searchTargetsFor(kbs)
+	if _, supplied := req.GetArguments()["knowledge_ids"]; supplied {
+		ids, err := req.RequireStringSlice("knowledge_ids")
+		if err != nil {
+			return mcp.NewToolResultError(err.Error()), nil
+		}
+		targets, err = s.documentSearchTargets(ctx, kbs, ids)
+		if err != nil {
+			return mcp.NewToolResultError(err.Error()), nil
+		}
+	}
+	return s.runSearchKnowledgeWithFilter(ctx, targets, strings.TrimSpace(query),
+		req.GetString("mode", tools.SearchModeHybrid), req.GetInt("limit", 0), nil)
 }
 
-// runSearchKnowledge executes the shared search_knowledge agent tool for
+// runSearchKnowledgeWithFilter executes the shared search_knowledge agent tool for
 // both the search_knowledge and grep_chunks endpoint tools.
-func (s *Server) runSearchKnowledge(
-	ctx context.Context, kbs []*types.KnowledgeBase, query, mode string, limit int,
-) (*mcp.CallToolResult, error) {
-	return s.runSearchKnowledgeWithFilter(ctx, kbs, query, mode, limit, nil)
-}
-
 func (s *Server) runSearchKnowledgeWithFilter(
-	ctx context.Context, kbs []*types.KnowledgeBase, query, mode string, limit int, pattern *regexp.Regexp,
+	ctx context.Context, targets types.SearchTargets, query, mode string, limit int, pattern *regexp.Regexp,
 ) (*mcp.CallToolResult, error) {
 	tool := tools.NewSearchKnowledgeTool(
-		s.kbService, s.knowledgeService, s.chunkService, searchTargetsFor(kbs), nil, s.cfg,
+		s.kbService, s.knowledgeService, s.chunkService, targets, nil, s.cfg,
 	)
 	if pattern != nil {
 		tool.WithPatternFilter(pattern)
 	}
 	callArgs := map[string]any{
 		"query":              query,
 		"mode":               mode,
-		"knowledge_base_ids": knowledgeBaseIDs(kbs),
+		"knowledge_base_ids": targets.GetAllKnowledgeBaseIDs(),
 	}
 	if limit > 0 {
 		callArgs["limit"] = limit
@@ -223,7 +232,7 @@ func (s *Server) handleGrepChunks(ctx context.Context, req mcp.CallToolRequest)
 			"word, identifier or phrase (for example \"timeout|deadline\" instead of \"^\\d+$\")"), nil
 	}
 	return s.runSearchKnowledgeWithFilter(
-		ctx, kbs, terms, tools.SearchModeKeyword, req.GetInt("limit", 0), pattern,
+		ctx, searchTargetsFor(kbs), terms, tools.SearchModeKeyword, req.GetInt("limit", 0), pattern,
 	)
 }