Merge pull request #1389 from trheyi/main

Refactor KB Search Handler and Update Test Cases

.github/workflows/pr-test.yml

@@ -336,6 +336,11 @@ jobs:
       - name: Run KB Tests (kb)
         run: make unit-test-kb
 
+      - name: Codecov Report
+        uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+
       - name: "Comment on PR - KB Tests Done"
         uses: actions/github-script@v7
         with:
@@ -528,6 +533,11 @@ jobs:
       - name: Run AI Tests (agent, aigc)
         run: make unit-test-ai
 
+      - name: Codecov Report
+        uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+
       - name: "Comment on PR - AI Tests Done"
         uses: actions/github-script@v7
         with:

.github/workflows/unit-test.yml

@@ -295,6 +295,11 @@ jobs:
       - name: Run KB Tests (kb)
         run: make unit-test-kb
 
+      - name: Codecov Report
+        uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+
   # =============================================================================
   # AI Tests (agent, aigc) - Run once with SQLite
   # =============================================================================
@@ -425,6 +430,11 @@ jobs:
       - name: Run AI Tests (agent, aigc)
         run: make unit-test-ai
 
+      - name: Codecov Report
+        uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+
   # =============================================================================
   # Benchmark & Memory Leak Tests - Run with MySQL8.0 and SQLite3
   # =============================================================================

Makefile

@@ -76,7 +76,7 @@ unit-test-core:
 # AI Unit Test (agent, aigc)
 .PHONY: unit-test-ai
 unit-test-ai:
-	echo "mode: count" > coverage-ai.out
+	echo "mode: count" > coverage.out
 	for d in $(TESTFOLDER_AI); do \
 		$(GO) test -tags $(TESTTAGS) -v -timeout=20m -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal' $$d > tmp.out; \
 		cat tmp.out; \
@@ -100,7 +100,7 @@ unit-test-ai:
 			exit 1; \
 		fi; \
 		if [ -f profile.out ]; then \
-			cat profile.out | grep -v "mode:" >> coverage-ai.out; \
+			cat profile.out | grep -v "mode:" >> coverage.out; \
 			rm profile.out; \
 		fi; \
 	done
@@ -108,7 +108,7 @@ unit-test-ai:
 # KB Unit Test (kb)
 .PHONY: unit-test-kb
 unit-test-kb:
-	echo "mode: count" > coverage-kb.out
+	echo "mode: count" > coverage.out
 	for d in $(TESTFOLDER_KB); do \
 		$(GO) test -tags $(TESTTAGS) -v -timeout=20m -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") -skip='TestMemoryLeak|TestIsolateDisposal|TestSearchCleanup' $$d > tmp.out; \
 		cat tmp.out; \
@@ -132,7 +132,7 @@ unit-test-kb:
 			exit 1; \
 		fi; \
 		if [ -f profile.out ]; then \
-			cat profile.out | grep -v "mode:" >> coverage-kb.out; \
+			cat profile.out | grep -v "mode:" >> coverage.out; \
 			rm profile.out; \
 		fi; \
 	done

agent/assistant/search.go

@@ -801,15 +801,25 @@ func (ast *Assistant) buildSearchRequests(ctx *context.Context, query string, co
 				threshold = config.KB.Threshold
 			}
 		}
-		requests = append(requests, &searchTypes.Request{
+
-			Type:        searchTypes.SearchTypeKB,
+		// Filter collections by authorization (Collection-level permission check)
-			Query:       query, // KB uses original query for semantic search
+		allowedCollections := FilterKBCollectionsByAuth(ctx, ast.KB.Collections)
-			Source:      searchTypes.SourceAuto,
+		if len(allowedCollections) == 0 {
-			Limit:       limit,
+			ctx.Logger.Info("No accessible KB collections after auth filter")
-			Collections: ast.KB.Collections,
+		} else {
-			Threshold:   threshold,
+			// Build KB request
-			Graph:       config != nil && config.KB != nil && config.KB.Graph,
+			kbReq := &searchTypes.Request{
-		})
+				Type:        searchTypes.SearchTypeKB,
+				Query:       query, // KB uses original query for semantic search
+				Source:      searchTypes.SourceAuto,
+				Limit:       limit,
+				Collections: allowedCollections,
+				Threshold:   threshold,
+				Graph:       config != nil && config.KB != nil && config.KB.Graph,
+			}
+
+			requests = append(requests, kbReq)
+		}
 	}
 
 	// DB search - check if DB is configured and allowed by intent
@@ -818,13 +828,22 @@ func (ast *Assistant) buildSearchRequests(ctx *context.Context, query string, co
 		if config != nil && config.DB != nil && config.DB.MaxResults > 0 {
 			limit = config.DB.MaxResults
 		}
-		requests = append(requests, &searchTypes.Request{
+
+		// Build DB request with auth where clauses
+		dbReq := &searchTypes.Request{
 			Type:   searchTypes.SearchTypeDB,
 			Query:  query, // DB uses original query for QueryDSL generation
 			Source: searchTypes.SourceAuto,
 			Limit:  limit,
 			Models: ast.DB.Models,
-		})
+		}
+
+		// Apply authorization where clauses
+		if authWheres := BuildDBAuthWheres(ctx); authWheres != nil {
+			dbReq.Wheres = authWheres
+		}
+
+		requests = append(requests, dbReq)
 	}
 
 	return requests, extractedKeywords

agent/assistant/search_auth_db.go

@@ -0,0 +1,103 @@
+package assistant
+
+import (
+	"github.com/yaoapp/gou/query/gou"
+	"github.com/yaoapp/yao/agent/context"
+)
+
+// BuildDBAuthWheres builds where clauses for DB search based on authorization
+// This applies permission-based filtering to database queries
+// Returns gou.Where clauses to filter records by authorization scope
+func BuildDBAuthWheres(ctx *context.Context) []gou.Where {
+	if ctx == nil || ctx.Authorized == nil {
+		return nil
+	}
+
+	authInfo := ctx.Authorized
+
+	// No constraints, no filter needed
+	if !authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
+		return nil
+	}
+
+	var wheres []gou.Where
+
+	// Team only - User can access:
+	// 1. Public records (public = true)
+	// 2. Records in their team where:
+	//    - They created the record (__yao_created_by matches)
+	//    - OR the record is shared with team (share = "team")
+	if authInfo.Constraints.TeamOnly && authInfo.TeamID != "" {
+		wheres = append(wheres, gou.Where{
+			Wheres: []gou.Where{
+				// Public records
+				{Condition: gou.Condition{
+					Field: &gou.Expression{Field: "public"},
+					Value: true,
+					OP:    "=",
+					OR:    true,
+				}},
+				// Team records
+				{
+					Wheres: []gou.Where{
+						{Condition: gou.Condition{
+							Field: &gou.Expression{Field: "__yao_team_id"},
+							Value: authInfo.TeamID,
+							OP:    "=",
+						}},
+						{Wheres: []gou.Where{
+							{Condition: gou.Condition{
+								Field: &gou.Expression{Field: "__yao_created_by"},
+								Value: authInfo.UserID,
+								OP:    "=",
+							}},
+							{Condition: gou.Condition{
+								Field: &gou.Expression{Field: "share"},
+								Value: "team",
+								OP:    "=",
+								OR:    true,
+							}},
+						}},
+					},
+				},
+			},
+		})
+		return wheres
+	}
+
+	// Owner only - User can access:
+	// 1. Public records (public = true)
+	// 2. Records they created where:
+	//    - __yao_team_id is null (not team records)
+	//    - __yao_created_by matches their user ID
+	if authInfo.Constraints.OwnerOnly && authInfo.UserID != "" {
+		wheres = append(wheres, gou.Where{
+			Wheres: []gou.Where{
+				// Public records
+				{Condition: gou.Condition{
+					Field: &gou.Expression{Field: "public"},
+					Value: true,
+					OP:    "=",
+					OR:    true,
+				}},
+				// Owner records
+				{
+					Wheres: []gou.Where{
+						{Condition: gou.Condition{
+							Field: &gou.Expression{Field: "__yao_team_id"},
+							OP:    "null",
+						}},
+						{Condition: gou.Condition{
+							Field: &gou.Expression{Field: "__yao_created_by"},
+							Value: authInfo.UserID,
+							OP:    "=",
+						}},
+					},
+				},
+			},
+		})
+		return wheres
+	}
+
+	return wheres
+}

agent/assistant/search_auth_integration_test.go

@@ -0,0 +1,631 @@
+package assistant_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	graphragtypes "github.com/yaoapp/gou/graphrag/types"
+	"github.com/yaoapp/yao/agent/assistant"
+	agentContext "github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/search"
+	searchTypes "github.com/yaoapp/yao/agent/search/types"
+	"github.com/yaoapp/yao/agent/testutils"
+	"github.com/yaoapp/yao/kb"
+	"github.com/yaoapp/yao/kb/api"
+	oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
+)
+
+// ========== Test Constants ==========
+
+const (
+	// Collection IDs for auth testing
+	AuthTestCollectionTeam1  = "auth_test_team1"
+	AuthTestCollectionTeam2  = "auth_test_team2"
+	AuthTestCollectionPublic = "auth_test_public"
+
+	// Test users and teams
+	TestUserA = "user_a"
+	TestUserB = "user_b"
+	TestTeam1 = "team_1"
+	TestTeam2 = "team_2"
+)
+
+// ========== Setup Test ==========
+
+// TestAuthSearchSetup creates test collections with different permissions.
+// Run once before running auth tests:
+//
+//	go test -v -run "TestAuthSearchSetup" ./agent/assistant/...
+func TestAuthSearchSetup(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	if kb.API == nil {
+		t.Fatal("KB API not initialized")
+	}
+
+	ctx := context.Background()
+
+	// Check if collections already exist
+	team1Exists := collectionReady(ctx, AuthTestCollectionTeam1, 2)
+	team2Exists := collectionReady(ctx, AuthTestCollectionTeam2, 2)
+	publicExists := collectionReady(ctx, AuthTestCollectionPublic, 2)
+
+	if team1Exists && team2Exists && publicExists {
+		t.Log("✓ All auth test collections already exist")
+		t.Log("  Run TestAuthSearchCleanup to recreate")
+		return
+	}
+
+	// Cleanup existing
+	t.Log("Cleaning up existing collections...")
+	cleanupAuthCollections(ctx, t)
+	time.Sleep(1 * time.Second)
+
+	// Create Team1 collection (owned by UserA, Team1)
+	t.Log("Creating Team1 collection...")
+	createAuthCollection(ctx, t, AuthTestCollectionTeam1, TestUserA, TestTeam1, false, "team")
+	addAuthDocument(ctx, t, AuthTestCollectionTeam1, "Team1 Doc1", "Team1 private document about quantum physics and relativity theory.")
+	addAuthDocument(ctx, t, AuthTestCollectionTeam1, "Team1 Doc2", "Team1 shared document about machine learning and neural networks.")
+
+	// Create Team2 collection (owned by UserB, Team2)
+	t.Log("Creating Team2 collection...")
+	createAuthCollection(ctx, t, AuthTestCollectionTeam2, TestUserB, TestTeam2, false, "team")
+	addAuthDocument(ctx, t, AuthTestCollectionTeam2, "Team2 Doc1", "Team2 private document about deep learning algorithms.")
+	addAuthDocument(ctx, t, AuthTestCollectionTeam2, "Team2 Doc2", "Team2 shared document about computer vision techniques.")
+
+	// Create Public collection
+	t.Log("Creating Public collection...")
+	createAuthCollection(ctx, t, AuthTestCollectionPublic, TestUserA, TestTeam1, true, "")
+	addAuthDocument(ctx, t, AuthTestCollectionPublic, "Public Doc1", "Public document about artificial intelligence and robotics.")
+	addAuthDocument(ctx, t, AuthTestCollectionPublic, "Public Doc2", "Public document about natural language processing.")
+
+	// Wait for indexing
+	t.Log("Waiting for indexing...")
+	time.Sleep(2 * time.Second)
+
+	t.Log("✓ Auth test setup complete!")
+}
+
+// TestAuthSearchCleanup removes auth test collections.
+func TestAuthSearchCleanup(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	if kb.API == nil {
+		t.Fatal("KB API not initialized")
+	}
+
+	ctx := context.Background()
+	cleanupAuthCollections(ctx, t)
+	t.Log("✓ Auth test cleanup complete!")
+}
+
+// ========== KB Collection-Level Auth Filter Tests ==========
+
+// Note: KB permission filtering works at the Collection level.
+// The Collection metadata contains __yao_team_id, __yao_created_by, public, share fields.
+// FilterKBCollectionsByAuth filters collections based on user authorization.
+
+func TestKBCollectionAuthFilter(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	// Ensure test data exists (auto-creates if not)
+	ensureAuthTestData(t)
+
+	t.Run("TeamMemberCanAccessTeamCollection", func(t *testing.T) {
+		// UserA from Team1 should access Team1 collection
+		ctx := createAuthContext(TestUserA, TestTeam1, true, false)
+		collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2}
+
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
+		assert.Contains(t, allowed, AuthTestCollectionTeam1, "Team1 member should access Team1 collection")
+		t.Logf("  Allowed collections: %v", allowed)
+	})
+
+	t.Run("TeamMemberCannotAccessOtherTeamCollection", func(t *testing.T) {
+		// UserA from Team1 should NOT access Team2 collection
+		ctx := createAuthContext(TestUserA, TestTeam1, true, false)
+		collections := []string{AuthTestCollectionTeam2}
+
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
+		assert.NotContains(t, allowed, AuthTestCollectionTeam2, "Team1 member should NOT access Team2 collection")
+		t.Logf("  Allowed collections: %v (expected empty)", allowed)
+	})
+
+	t.Run("OwnerCanAccessOwnCollection", func(t *testing.T) {
+		// UserA with OwnerOnly should access collections they created
+		ctx := createAuthContext(TestUserA, "", false, true)
+		collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2}
+
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
+		assert.Contains(t, allowed, AuthTestCollectionTeam1, "Owner should access own collection")
+		assert.NotContains(t, allowed, AuthTestCollectionTeam2, "Owner should NOT access other's collection")
+		t.Logf("  Allowed collections: %v", allowed)
+	})
+
+	t.Run("PublicCollectionAccessibleToAll", func(t *testing.T) {
+		// Note: The 'public' field in Metadata is not automatically saved to the database
+		// by the current KB API. This test documents the expected behavior.
+		// When public=true is properly set in DB, this should pass.
+
+		// First, check the collection metadata
+		bgCtx := context.Background()
+		collection, err := kb.API.GetCollection(bgCtx, AuthTestCollectionPublic)
+		assert.NoError(t, err)
+
+		// Check if public is set correctly
+		publicVal := collection["public"]
+		t.Logf("  Public collection public field: %v (type: %T)", publicVal, publicVal)
+
+		// If public is not set (0 or false), the test documents current behavior
+		// The collection should be accessible via owner check since UserA created it
+		ctx := createAuthContext(TestUserA, TestTeam1, false, true) // Owner check
+		collections := []string{AuthTestCollectionPublic}
+
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
+		assert.Contains(t, allowed, AuthTestCollectionPublic, "Owner should access their collection")
+		t.Logf("  Allowed collections (owner check): %v", allowed)
+	})
+
+	t.Run("NoConstraintsMeansFullAccess", func(t *testing.T) {
+		// User with no constraints should access all collections
+		ctx := createAuthContext(TestUserA, TestTeam1, false, false)
+		collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
+
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
+		assert.Len(t, allowed, 3, "No constraints should allow all collections")
+		t.Logf("  Allowed collections: %v", allowed)
+	})
+
+	t.Run("NilContextMeansFullAccess", func(t *testing.T) {
+		collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2}
+
+		allowed := assistant.FilterKBCollectionsByAuth(nil, collections)
+		assert.Len(t, allowed, 2, "Nil context should allow all collections")
+		t.Logf("  Allowed collections: %v", allowed)
+	})
+}
+
+// ========== DB Auth Wheres Tests ==========
+
+func TestDBAuthWheresFilter(t *testing.T) {
+	// Note: This test doesn't need KB, just tests the BuildDBAuthWheres function
+	t.Run("TeamOnlyGeneratesCorrectWheres", func(t *testing.T) {
+		ctx := createAuthContext(TestUserA, TestTeam1, true, false)
+		wheres := assistant.BuildDBAuthWheres(ctx)
+
+		assert.NotNil(t, wheres)
+		assert.Len(t, wheres, 1)
+
+		// Verify structure: should have 2 top-level conditions (public OR team filter)
+		where := wheres[0]
+		assert.Len(t, where.Wheres, 2, "Should have 2 conditions: public OR team")
+
+		// First condition: public = true (OR)
+		publicCond := where.Wheres[0]
+		assert.NotNil(t, publicCond.Condition.Field)
+		assert.Equal(t, "public", publicCond.Condition.Field.Field)
+		assert.Equal(t, true, publicCond.Condition.Value)
+		assert.True(t, publicCond.Condition.OR)
+
+		// Second condition: team filter with nested conditions
+		teamCond := where.Wheres[1]
+		assert.Len(t, teamCond.Wheres, 2, "Team filter should have team_id and (created_by OR share)")
+
+		// Team ID check
+		teamIDCond := teamCond.Wheres[0]
+		assert.Equal(t, "__yao_team_id", teamIDCond.Condition.Field.Field)
+		assert.Equal(t, TestTeam1, teamIDCond.Condition.Value)
+
+		// Created by OR share = team
+		ownerOrShareCond := teamCond.Wheres[1]
+		assert.Len(t, ownerOrShareCond.Wheres, 2)
+		assert.Equal(t, "__yao_created_by", ownerOrShareCond.Wheres[0].Condition.Field.Field)
+		assert.Equal(t, TestUserA, ownerOrShareCond.Wheres[0].Condition.Value)
+		assert.Equal(t, "share", ownerOrShareCond.Wheres[1].Condition.Field.Field)
+		assert.Equal(t, "team", ownerOrShareCond.Wheres[1].Condition.Value)
+		assert.True(t, ownerOrShareCond.Wheres[1].Condition.OR)
+
+		t.Logf("  TeamOnly: Verified team_id=%s, created_by=%s", TestTeam1, TestUserA)
+	})
+
+	t.Run("OwnerOnlyGeneratesCorrectWheres", func(t *testing.T) {
+		ctx := createAuthContext(TestUserA, "", false, true)
+		wheres := assistant.BuildDBAuthWheres(ctx)
+
+		assert.NotNil(t, wheres)
+		assert.Len(t, wheres, 1)
+
+		// Verify structure: should have 2 top-level conditions (public OR owner filter)
+		where := wheres[0]
+		assert.Len(t, where.Wheres, 2, "Should have 2 conditions: public OR owner")
+
+		// First condition: public = true (OR)
+		publicCond := where.Wheres[0]
+		assert.NotNil(t, publicCond.Condition.Field)
+		assert.Equal(t, "public", publicCond.Condition.Field.Field)
+		assert.Equal(t, true, publicCond.Condition.Value)
+		assert.True(t, publicCond.Condition.OR)
+
+		// Second condition: owner filter with nested conditions
+		ownerCond := where.Wheres[1]
+		assert.Len(t, ownerCond.Wheres, 2, "Owner filter should have team_id IS NULL and created_by")
+
+		// Team ID is null check
+		teamNullCond := ownerCond.Wheres[0]
+		assert.Equal(t, "__yao_team_id", teamNullCond.Condition.Field.Field)
+		assert.Equal(t, "null", teamNullCond.Condition.OP)
+
+		// Created by check
+		createdByCond := ownerCond.Wheres[1]
+		assert.Equal(t, "__yao_created_by", createdByCond.Condition.Field.Field)
+		assert.Equal(t, TestUserA, createdByCond.Condition.Value)
+
+		t.Logf("  OwnerOnly: Verified created_by=%s, team_id IS NULL", TestUserA)
+	})
+
+	t.Run("NoConstraintsReturnsNil", func(t *testing.T) {
+		ctx := createAuthContext(TestUserA, TestTeam1, false, false)
+		wheres := assistant.BuildDBAuthWheres(ctx)
+
+		assert.Nil(t, wheres, "No constraints should return nil")
+		t.Log("  No constraints: nil wheres (no filter)")
+	})
+
+	t.Run("EmptyTeamIDReturnsNil", func(t *testing.T) {
+		ctx := createAuthContext(TestUserA, "", true, false)
+		wheres := assistant.BuildDBAuthWheres(ctx)
+
+		assert.Nil(t, wheres, "Empty TeamID with TeamOnly should return nil")
+		t.Log("  Empty TeamID with TeamOnly: nil wheres")
+	})
+
+	t.Run("EmptyUserIDReturnsNil", func(t *testing.T) {
+		ctx := createAuthContext("", TestTeam1, false, true)
+		wheres := assistant.BuildDBAuthWheres(ctx)
+
+		assert.Nil(t, wheres, "Empty UserID with OwnerOnly should return nil")
+		t.Log("  Empty UserID with OwnerOnly: nil wheres")
+	})
+
+	t.Run("NilContextReturnsNil", func(t *testing.T) {
+		wheres := assistant.BuildDBAuthWheres(nil)
+
+		assert.Nil(t, wheres, "Nil context should return nil")
+		t.Log("  Nil context: nil wheres")
+	})
+
+	t.Run("NilAuthorizedReturnsNil", func(t *testing.T) {
+		ctx := &agentContext.Context{Authorized: nil}
+		wheres := assistant.BuildDBAuthWheres(ctx)
+
+		assert.Nil(t, wheres, "Nil Authorized should return nil")
+		t.Log("  Nil Authorized: nil wheres")
+	})
+}
+
+// ========== KB Search Integration Tests ==========
+
+func TestKBSearchIntegration(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	// Ensure test data exists (auto-creates if not)
+	ensureAuthTestData(t)
+
+	t.Run("TeamMemberSearchOnlyFindsTeamData", func(t *testing.T) {
+		// UserA from Team1 searches - should ONLY find Team1 data
+		ctx := createAuthContext(TestUserA, TestTeam1, true, false)
+
+		// Filter collections first
+		allCollections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2}
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, allCollections)
+
+		// Should only allow Team1
+		assert.Contains(t, allowed, AuthTestCollectionTeam1)
+		assert.NotContains(t, allowed, AuthTestCollectionTeam2)
+		assert.Len(t, allowed, 1, "Should only have 1 allowed collection")
+
+		// Search on allowed collections
+		result := executeKBSearchOnCollections(t, allowed, "quantum physics deep learning")
+		assert.Greater(t, len(result.Items), 0, "Should find Team1 documents")
+
+		// Verify ALL results are from Team1 collection only
+		for _, item := range result.Items {
+			assert.Equal(t, AuthTestCollectionTeam1, item.Collection,
+				"All results should be from Team1 collection, got: %s", item.Collection)
+		}
+		t.Logf("  ✓ Team1 member found %d items, all from Team1 collection", len(result.Items))
+	})
+
+	t.Run("TeamMemberCannotAccessOtherTeamData", func(t *testing.T) {
+		// UserA from Team1 tries to access Team2 - should be blocked
+		ctx := createAuthContext(TestUserA, TestTeam1, true, false)
+
+		// Try to filter Team2 collection
+		collections := []string{AuthTestCollectionTeam2}
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, collections)
+
+		// Should be empty - no access
+		assert.Empty(t, allowed, "Team1 member should NOT have access to Team2 collection")
+		t.Log("  ✓ Team1 member correctly blocked from Team2 collection")
+	})
+
+	t.Run("OwnerSearchOnlyFindsOwnData", func(t *testing.T) {
+		// UserA with OwnerOnly - should only find collections they created
+		ctx := createAuthContext(TestUserA, "", false, true)
+
+		// Filter all collections
+		allCollections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, allCollections)
+
+		// UserA created Team1 and Public, not Team2
+		assert.Contains(t, allowed, AuthTestCollectionTeam1, "Owner should access Team1 (created by UserA)")
+		assert.Contains(t, allowed, AuthTestCollectionPublic, "Owner should access Public (created by UserA)")
+		assert.NotContains(t, allowed, AuthTestCollectionTeam2, "Owner should NOT access Team2 (created by UserB)")
+
+		// Search and verify results
+		result := executeKBSearchOnCollections(t, allowed, "quantum artificial intelligence")
+		assert.Greater(t, len(result.Items), 0, "Should find owner's documents")
+
+		// Verify NO results from Team2
+		for _, item := range result.Items {
+			assert.NotEqual(t, AuthTestCollectionTeam2, item.Collection,
+				"Should NOT have results from Team2, got: %s", item.Collection)
+		}
+		t.Logf("  ✓ Owner found %d items, none from Team2", len(result.Items))
+	})
+
+	t.Run("NoConstraintsSearchFindsAllData", func(t *testing.T) {
+		// User with no constraints - should find all data
+		ctx := createAuthContext(TestUserA, TestTeam1, false, false)
+
+		// Filter all collections
+		allCollections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, allCollections)
+
+		// Should have access to all
+		assert.Len(t, allowed, 3, "No constraints should allow all collections")
+
+		// Search and verify results from multiple collections
+		result := executeKBSearchOnCollections(t, allowed, "quantum deep learning artificial")
+
+		// Should find results from multiple collections
+		collectionsFound := make(map[string]bool)
+		for _, item := range result.Items {
+			collectionsFound[item.Collection] = true
+		}
+		assert.Greater(t, len(collectionsFound), 1, "Should find results from multiple collections")
+		t.Logf("  ✓ No constraints: found %d items from %d collections", len(result.Items), len(collectionsFound))
+	})
+
+	t.Run("SearchResultsMatchCollectionFilter", func(t *testing.T) {
+		// Verify that search results ONLY come from allowed collections
+		ctx := createAuthContext(TestUserB, TestTeam2, true, false)
+
+		// UserB from Team2 - should only access Team2
+		allCollections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
+		allowed := assistant.FilterKBCollectionsByAuth(ctx, allCollections)
+
+		assert.Contains(t, allowed, AuthTestCollectionTeam2, "Team2 member should access Team2")
+		assert.NotContains(t, allowed, AuthTestCollectionTeam1, "Team2 member should NOT access Team1")
+
+		// Search
+		result := executeKBSearchOnCollections(t, allowed, "deep learning computer vision")
+
+		// Verify results
+		if len(result.Items) > 0 {
+			for _, item := range result.Items {
+				// Results should only be from allowed collections
+				assert.Contains(t, allowed, item.Collection,
+					"Result from %s should be in allowed list %v", item.Collection, allowed)
+			}
+			t.Logf("  ✓ Team2 member found %d items, all from allowed collections", len(result.Items))
+		} else {
+			t.Log("  ✓ Team2 member found 0 items (collection may be empty)")
+		}
+	})
+}
+
+// ========== Helper Functions ==========
+
+// ensureAuthTestData checks if auth test data exists, creates if not.
+// This is a utility function that can be called from any test.
+// Note: testutils.Prepare must be called before this function.
+func ensureAuthTestData(t *testing.T) {
+	if kb.API == nil {
+		t.Fatal("KB API not initialized - call testutils.Prepare first")
+	}
+
+	ctx := context.Background()
+
+	// Check if all collections already exist with required documents
+	team1Ready := collectionReady(ctx, AuthTestCollectionTeam1, 2)
+	team2Ready := collectionReady(ctx, AuthTestCollectionTeam2, 2)
+	publicReady := collectionReady(ctx, AuthTestCollectionPublic, 2)
+
+	if team1Ready && team2Ready && publicReady {
+		// All data exists, skip creation
+		return
+	}
+
+	// Data not ready, create it
+	t.Log("Auth test data not found, creating...")
+	createAuthTestData(t, ctx)
+}
+
+// createAuthTestData creates the test collections and documents
+func createAuthTestData(t *testing.T, ctx context.Context) {
+	// Cleanup existing
+	t.Log("Cleaning up existing collections...")
+	cleanupAuthCollections(ctx, t)
+	time.Sleep(1 * time.Second)
+
+	// Create Team1 collection (owned by UserA, Team1)
+	t.Log("Creating Team1 collection...")
+	createAuthCollection(ctx, t, AuthTestCollectionTeam1, TestUserA, TestTeam1, false, "team")
+	addAuthDocument(ctx, t, AuthTestCollectionTeam1, "Team1 Doc1", "Team1 private document about quantum physics and relativity theory.")
+	addAuthDocument(ctx, t, AuthTestCollectionTeam1, "Team1 Doc2", "Team1 shared document about machine learning and neural networks.")
+
+	// Create Team2 collection (owned by UserB, Team2)
+	t.Log("Creating Team2 collection...")
+	createAuthCollection(ctx, t, AuthTestCollectionTeam2, TestUserB, TestTeam2, false, "team")
+	addAuthDocument(ctx, t, AuthTestCollectionTeam2, "Team2 Doc1", "Team2 private document about deep learning algorithms.")
+	addAuthDocument(ctx, t, AuthTestCollectionTeam2, "Team2 Doc2", "Team2 shared document about computer vision techniques.")
+
+	// Create Public collection
+	t.Log("Creating Public collection...")
+	createAuthCollection(ctx, t, AuthTestCollectionPublic, TestUserA, TestTeam1, true, "")
+	addAuthDocument(ctx, t, AuthTestCollectionPublic, "Public Doc1", "Public document about artificial intelligence and robotics.")
+	addAuthDocument(ctx, t, AuthTestCollectionPublic, "Public Doc2", "Public document about natural language processing.")
+
+	// Wait for indexing
+	t.Log("Waiting for indexing...")
+	time.Sleep(2 * time.Second)
+
+	t.Log("✓ Auth test data created!")
+}
+
+func createAuthContext(userID, teamID string, teamOnly, ownerOnly bool) *agentContext.Context {
+	return &agentContext.Context{
+		Authorized: &oauthtypes.AuthorizedInfo{
+			UserID: userID,
+			TeamID: teamID,
+			Constraints: oauthtypes.DataConstraints{
+				TeamOnly:  teamOnly,
+				OwnerOnly: ownerOnly,
+			},
+		},
+	}
+}
+
+func collectionReady(ctx context.Context, collectionID string, minDocs int) bool {
+	collection, err := kb.API.GetCollection(ctx, collectionID)
+	if err != nil || collection == nil {
+		return false
+	}
+
+	docs, err := kb.API.ListDocuments(ctx, &api.ListDocumentsFilter{
+		Page:         1,
+		PageSize:     20,
+		CollectionID: collectionID,
+	})
+	if err != nil || docs == nil {
+		return false
+	}
+
+	return len(docs.Data) >= minDocs
+}
+
+func cleanupAuthCollections(ctx context.Context, t *testing.T) {
+	collections := []string{AuthTestCollectionTeam1, AuthTestCollectionTeam2, AuthTestCollectionPublic}
+	for _, id := range collections {
+		if result, err := kb.API.RemoveCollection(ctx, id); err == nil && result.Removed {
+			t.Logf("  Removed: %s", id)
+		}
+	}
+}
+
+func createAuthCollection(ctx context.Context, t *testing.T, id, userID, teamID string, public bool, share string) {
+	params := &api.CreateCollectionParams{
+		ID: id,
+		Metadata: map[string]interface{}{
+			"name":   id,
+			"public": public,
+			"share":  share,
+		},
+		EmbeddingProviderID: "__yao.openai",
+		EmbeddingOptionID:   "text-embedding-3-small",
+		Locale:              "en",
+		Config: &graphragtypes.CreateCollectionOptions{
+			Distance:  "cosine",
+			IndexType: "hnsw",
+		},
+		AuthScope: map[string]interface{}{
+			"__yao_created_by": userID,
+			"__yao_team_id":    teamID,
+		},
+	}
+
+	_, err := kb.API.CreateCollection(ctx, params)
+	if err != nil {
+		t.Fatalf("Failed to create collection %s: %v", id, err)
+	}
+	t.Logf("  ✓ Created: %s", id)
+}
+
+func addAuthDocument(ctx context.Context, t *testing.T, collectionID, title, content string) {
+	params := &api.AddTextParams{
+		CollectionID: collectionID,
+		Text:         content,
+		DocID:        fmt.Sprintf("%s__%s", collectionID, sanitizeForID(title)),
+		Metadata: map[string]interface{}{
+			"title": title,
+		},
+		Chunking: &api.ProviderConfigParams{
+			ProviderID: "__yao.structured",
+			OptionID:   "standard",
+		},
+		Embedding: &api.ProviderConfigParams{
+			ProviderID: "__yao.openai",
+			OptionID:   "text-embedding-3-small",
+		},
+	}
+
+	_, err := kb.API.AddText(ctx, params)
+	if err != nil {
+		t.Logf("  Warning: Failed to add document '%s': %v", title, err)
+		return
+	}
+	t.Logf("    ✓ Added: %s", title)
+}
+
+func sanitizeForID(s string) string {
+	result := ""
+	for _, c := range s {
+		if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') {
+			result += string(c)
+		} else if c == ' ' {
+			result += "_"
+		}
+	}
+	return result
+}
+
+func executeKBSearch(t *testing.T, collectionID, query string, metadata map[string]interface{}) *searchTypes.Result {
+	return executeKBSearchOnCollections(t, []string{collectionID}, query)
+}
+
+func executeKBSearchOnCollections(t *testing.T, collections []string, query string) *searchTypes.Result {
+	if len(collections) == 0 {
+		return &searchTypes.Result{Items: []*searchTypes.ResultItem{}}
+	}
+
+	cfg := &searchTypes.Config{
+		KB: &searchTypes.KBConfig{
+			Collections: collections,
+			Threshold:   0.3,
+		},
+	}
+	searcher := search.New(cfg, nil)
+
+	req := &searchTypes.Request{
+		Type:        searchTypes.SearchTypeKB,
+		Query:       query,
+		Collections: collections,
+		Threshold:   0.3,
+		Limit:       20,
+		Source:      searchTypes.SourceAuto,
+	}
+
+	result, err := searcher.Search(nil, req)
+	if err != nil {
+		t.Fatalf("Search failed: %v", err)
+	}
+	return result
+}

agent/assistant/search_auth_kb.go

@@ -0,0 +1,123 @@
+package assistant
+
+import (
+	"context"
+
+	agentContext "github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/kb"
+	oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
+)
+
+// FilterKBCollectionsByAuth filters collections based on user authorization.
+// Returns only collections that the user has permission to access.
+// Permission is determined by Collection's metadata (public, share, __yao_team_id, __yao_created_by).
+func FilterKBCollectionsByAuth(ctx *agentContext.Context, collections []string) []string {
+	if ctx == nil || ctx.Authorized == nil {
+		return collections // No auth context, return all
+	}
+
+	authInfo := ctx.Authorized
+
+	// No constraints, return all collections
+	if !authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
+		return collections
+	}
+
+	// Check KB API
+	if kb.API == nil {
+		return collections // KB not initialized, return all
+	}
+
+	var allowed []string
+	bgCtx := context.Background()
+
+	for _, collectionID := range collections {
+		// Get collection metadata
+		collection, err := kb.API.GetCollection(bgCtx, collectionID)
+		if err != nil {
+			continue // Skip if can't get collection
+		}
+
+		if hasCollectionAccess(authInfo, collection) {
+			allowed = append(allowed, collectionID)
+		}
+	}
+
+	return allowed
+}
+
+// hasCollectionAccess checks if user has access to a collection based on its metadata.
+func hasCollectionAccess(authInfo *oauthtypes.AuthorizedInfo, collection map[string]interface{}) bool {
+	if authInfo == nil {
+		return true
+	}
+
+	// No constraints, allow access
+	if !authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
+		return true
+	}
+
+	// Check public access (handle different types: bool, int, float64)
+	if isPublicValue(collection["public"]) {
+		return true
+	}
+
+	// Get metadata for permission fields
+	metadata, _ := collection["metadata"].(map[string]interface{})
+	if metadata == nil {
+		metadata = collection
+	}
+
+	// Team only check
+	if authInfo.Constraints.TeamOnly && authInfo.TeamID != "" {
+		teamID, _ := metadata["__yao_team_id"].(string)
+		if teamID == "" {
+			teamID, _ = collection["__yao_team_id"].(string)
+		}
+
+		if teamID == authInfo.TeamID {
+			createdBy, _ := metadata["__yao_created_by"].(string)
+			if createdBy == "" {
+				createdBy, _ = collection["__yao_created_by"].(string)
+			}
+			share, _ := metadata["share"].(string)
+			if share == "" {
+				share, _ = collection["share"].(string)
+			}
+
+			if createdBy == authInfo.UserID || share == "team" {
+				return true
+			}
+		}
+	}
+
+	// Owner only check
+	if authInfo.Constraints.OwnerOnly && authInfo.UserID != "" {
+		createdBy, _ := metadata["__yao_created_by"].(string)
+		if createdBy == "" {
+			createdBy, _ = collection["__yao_created_by"].(string)
+		}
+		if createdBy == authInfo.UserID {
+			return true
+		}
+	}
+
+	return false
+}
+
+// isPublicValue checks if a value represents "public" access
+func isPublicValue(v interface{}) bool {
+	switch val := v.(type) {
+	case bool:
+		return val
+	case int:
+		return val == 1
+	case int64:
+		return val == 1
+	case float64:
+		return val == 1
+	case string:
+		return val == "true" || val == "1"
+	}
+	return false
+}

agent/search/handlers/kb/handler.go

@@ -1,12 +1,16 @@
 package kb
 
 import (
+	"context"
+	"fmt"
 	"time"
 
 	"github.com/yaoapp/yao/agent/search/types"
+	"github.com/yaoapp/yao/kb"
+	kbapi "github.com/yaoapp/yao/kb/api"
 )
 
-// Handler implements KB search
+// Handler implements KB search using the KB API
 type Handler struct {
 	config *types.KBConfig // KB search configuration
 }
@@ -22,7 +26,6 @@ func (h *Handler) Type() types.SearchType {
 }
 
 // Search executes vector search and optional graph association
-// TODO: Implement actual vector search and graph association logic
 func (h *Handler) Search(req *types.Request) (*types.Result, error) {
 	start := time.Now()
 
@@ -39,6 +42,19 @@ func (h *Handler) Search(req *types.Request) (*types.Result, error) {
 		}, nil
 	}
 
+	// Check if KB API is available
+	if kb.API == nil {
+		return &types.Result{
+			Type:     types.SearchTypeKB,
+			Query:    req.Query,
+			Source:   req.Source,
+			Items:    []*types.ResultItem{},
+			Total:    0,
+			Duration: time.Since(start).Milliseconds(),
+			Error:    "knowledge base not initialized",
+		}, nil
+	}
+
 	// Get collections from request or config
 	collections := req.Collections
 	if len(collections) == 0 && h.config != nil {
@@ -72,24 +88,96 @@ func (h *Handler) Search(req *types.Request) (*types.Result, error) {
 		limit = 10 // default
 	}
 
-	// TODO: Implement actual vector search
+	// Determine search mode
-	// 1. Generate embedding for query using collection's embedding config
+	mode := kbapi.SearchModeVector
-	// 2. Search each collection with vector similarity
+	if req.Graph {
-	// 3. If req.Graph is true, perform graph association
+		mode = kbapi.SearchModeExpand
-	// 4. Merge and return results
+	}
-
+	if h.config != nil && h.config.Graph {
-	// For now, return empty result (skeleton)
+		mode = kbapi.SearchModeExpand
-	result := &types.Result{
-		Type:     types.SearchTypeKB,
-		Query:    req.Query,
-		Source:   req.Source,
-		Items:    []*types.ResultItem{},
-		Total:    0,
-		Duration: time.Since(start).Milliseconds(),
 	}
 
-	// Store threshold in result metadata for debugging
+	// Build KB API queries - one per collection
-	_ = threshold
+	var queries []kbapi.Query
+	for _, collectionID := range collections {
+		queries = append(queries, kbapi.Query{
+			CollectionID: collectionID,
+			Input:        req.Query,
+			Mode:         mode,
+			Threshold:    threshold,
+			PageSize:     limit,
+			Metadata:     req.Metadata,
+		})
+	}
+
+	// Execute search using KB API
+	ctx := context.Background()
+	searchResult, err := kb.API.Search(ctx, queries)
+	if err != nil {
+		return &types.Result{
+			Type:     types.SearchTypeKB,
+			Query:    req.Query,
+			Source:   req.Source,
+			Items:    []*types.ResultItem{},
+			Total:    0,
+			Duration: time.Since(start).Milliseconds(),
+			Error:    fmt.Sprintf("search failed: %v", err),
+		}, nil
+	}
+
+	// Convert segments to result items
+	// Note: MinScore filtering is already done by KB API, no need to filter again
+	items := make([]*types.ResultItem, 0, len(searchResult.Segments))
+	for _, seg := range searchResult.Segments {
+		item := &types.ResultItem{
+			Type:       types.SearchTypeKB,
+			Source:     req.Source,
+			Score:      seg.Score,
+			Content:    seg.Text,
+			DocumentID: seg.DocumentID,
+			Collection: seg.CollectionID,
+			Metadata:   seg.Metadata,
+		}
+
+		// Extract title from metadata if available
+		if seg.Metadata != nil {
+			if title, ok := seg.Metadata["title"].(string); ok {
+				item.Title = title
+			}
+		}
+
+		items = append(items, item)
+	}
+
+	// Convert graph data if available
+	var graphNodes []*types.GraphNode
+	if searchResult.Graph != nil {
+		for _, node := range searchResult.Graph.Nodes {
+			// Extract name from properties if available
+			name := ""
+			if node.Properties != nil {
+				if n, ok := node.Properties["name"].(string); ok {
+					name = n
+				}
+			}
+			graphNodes = append(graphNodes, &types.GraphNode{
+				ID:       node.ID,
+				Type:     node.EntityType,
+				Name:     name,
+				Metadata: node.Properties,
+			})
+		}
+	}
+
+	result := &types.Result{
+		Type:       types.SearchTypeKB,
+		Query:      req.Query,
+		Source:     req.Source,
+		Items:      items,
+		Total:      len(items),
+		Duration:   time.Since(start).Milliseconds(),
+		GraphNodes: graphNodes,
+	}
 
 	return result, nil
 }

agent/search/handlers/kb/handler_test.go

@@ -31,13 +31,12 @@ func TestHandler_Type(t *testing.T) {
 	assert.Equal(t, types.SearchTypeKB, h.Type())
 }
 
-func TestHandler_Search(t *testing.T) {
+func TestHandler_Search_Validation(t *testing.T) {
 	tests := []struct {
 		name        string
 		config      *types.KBConfig
 		req         *types.Request
 		expectError string
-		expectItems int
 	}{
 		{
 			name:   "empty query",
@@ -47,85 +46,15 @@ func TestHandler_Search(t *testing.T) {
 				Query: "",
 			},
 			expectError: "query is required",
-			expectItems: 0,
 		},
 		{
-			name:   "no collections in request or config",
+			name:   "no collections - KB not initialized",
 			config: nil,
 			req: &types.Request{
 				Type:  types.SearchTypeKB,
 				Query: "test query",
 			},
-			expectError: "",
+			expectError: "knowledge base not initialized",
-			expectItems: 0,
-		},
-		{
-			name: "collections from config",
-			config: &types.KBConfig{
-				Collections: []string{"docs"},
-				Threshold:   0.7,
-			},
-			req: &types.Request{
-				Type:  types.SearchTypeKB,
-				Query: "test query",
-			},
-			expectError: "",
-			expectItems: 0, // skeleton returns empty
-		},
-		{
-			name:   "collections from request",
-			config: nil,
-			req: &types.Request{
-				Type:        types.SearchTypeKB,
-				Query:       "test query",
-				Collections: []string{"docs", "faq"},
-			},
-			expectError: "",
-			expectItems: 0, // skeleton returns empty
-		},
-		{
-			name: "with threshold from request",
-			config: &types.KBConfig{
-				Collections: []string{"docs"},
-				Threshold:   0.7,
-			},
-			req: &types.Request{
-				Type:        types.SearchTypeKB,
-				Query:       "test query",
-				Threshold:   0.9,
-				Collections: []string{"docs"},
-			},
-			expectError: "",
-			expectItems: 0, // skeleton returns empty
-		},
-		{
-			name: "with graph enabled",
-			config: &types.KBConfig{
-				Collections: []string{"docs"},
-				Graph:       true,
-			},
-			req: &types.Request{
-				Type:        types.SearchTypeKB,
-				Query:       "test query",
-				Collections: []string{"docs"},
-				Graph:       true,
-			},
-			expectError: "",
-			expectItems: 0, // skeleton returns empty
-		},
-		{
-			name: "with limit",
-			config: &types.KBConfig{
-				Collections: []string{"docs"},
-			},
-			req: &types.Request{
-				Type:        types.SearchTypeKB,
-				Query:       "test query",
-				Collections: []string{"docs"},
-				Limit:       5,
-			},
-			expectError: "",
-			expectItems: 0, // skeleton returns empty
 		},
 	}
 
@@ -138,7 +67,6 @@ func TestHandler_Search(t *testing.T) {
 			assert.NotNil(t, result)
 			assert.Equal(t, types.SearchTypeKB, result.Type)
 			assert.Equal(t, tt.req.Query, result.Query)
-			assert.Equal(t, tt.expectItems, len(result.Items))
 
 			if tt.expectError != "" {
 				assert.Equal(t, tt.expectError, result.Error)
@@ -168,3 +96,166 @@ func TestHandler_Search_SourcePreserved(t *testing.T) {
 		assert.Equal(t, source, result.Source)
 	}
 }
+
+func TestHandler_Search_CollectionsFromConfig(t *testing.T) {
+	cfg := &types.KBConfig{
+		Collections: []string{"docs", "faq"},
+		Threshold:   0.7,
+	}
+	h := NewHandler(cfg)
+
+	// Request without collections should use config collections
+	req := &types.Request{
+		Type:  types.SearchTypeKB,
+		Query: "test query",
+	}
+	result, err := h.Search(req)
+
+	assert.NoError(t, err)
+	assert.NotNil(t, result)
+	// Without KB initialized, we get "knowledge base not initialized" error
+	assert.Equal(t, "knowledge base not initialized", result.Error)
+}
+
+func TestHandler_Search_CollectionsFromRequest(t *testing.T) {
+	h := NewHandler(nil)
+
+	// Request with collections
+	req := &types.Request{
+		Type:        types.SearchTypeKB,
+		Query:       "test query",
+		Collections: []string{"docs", "faq"},
+	}
+	result, err := h.Search(req)
+
+	assert.NoError(t, err)
+	assert.NotNil(t, result)
+	// Without KB initialized, we get "knowledge base not initialized" error
+	assert.Equal(t, "knowledge base not initialized", result.Error)
+}
+
+func TestHandler_Search_ThresholdHandling(t *testing.T) {
+	tests := []struct {
+		name            string
+		configThreshold float64
+		reqThreshold    float64
+	}{
+		{
+			name:            "threshold from request",
+			configThreshold: 0.7,
+			reqThreshold:    0.9,
+		},
+		{
+			name:            "threshold from config",
+			configThreshold: 0.8,
+			reqThreshold:    0,
+		},
+		{
+			name:            "default threshold",
+			configThreshold: 0,
+			reqThreshold:    0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var cfg *types.KBConfig
+			if tt.configThreshold > 0 {
+				cfg = &types.KBConfig{
+					Collections: []string{"docs"},
+					Threshold:   tt.configThreshold,
+				}
+			}
+			h := NewHandler(cfg)
+
+			req := &types.Request{
+				Type:        types.SearchTypeKB,
+				Query:       "test query",
+				Threshold:   tt.reqThreshold,
+				Collections: []string{"docs"},
+			}
+			result, err := h.Search(req)
+
+			assert.NoError(t, err)
+			assert.NotNil(t, result)
+		})
+	}
+}
+
+func TestHandler_Search_GraphMode(t *testing.T) {
+	tests := []struct {
+		name        string
+		configGraph bool
+		reqGraph    bool
+	}{
+		{
+			name:        "graph from request",
+			configGraph: false,
+			reqGraph:    true,
+		},
+		{
+			name:        "graph from config",
+			configGraph: true,
+			reqGraph:    false,
+		},
+		{
+			name:        "no graph",
+			configGraph: false,
+			reqGraph:    false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg := &types.KBConfig{
+				Collections: []string{"docs"},
+				Graph:       tt.configGraph,
+			}
+			h := NewHandler(cfg)
+
+			req := &types.Request{
+				Type:        types.SearchTypeKB,
+				Query:       "test query",
+				Collections: []string{"docs"},
+				Graph:       tt.reqGraph,
+			}
+			result, err := h.Search(req)
+
+			assert.NoError(t, err)
+			assert.NotNil(t, result)
+		})
+	}
+}
+
+func TestHandler_Search_LimitHandling(t *testing.T) {
+	h := NewHandler(&types.KBConfig{Collections: []string{"docs"}})
+
+	tests := []struct {
+		name  string
+		limit int
+	}{
+		{
+			name:  "custom limit",
+			limit: 5,
+		},
+		{
+			name:  "default limit",
+			limit: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := &types.Request{
+				Type:        types.SearchTypeKB,
+				Query:       "test query",
+				Collections: []string{"docs"},
+				Limit:       tt.limit,
+			}
+			result, err := h.Search(req)
+
+			assert.NoError(t, err)
+			assert.NotNil(t, result)
+		})
+	}
+}

agent/search/jsapi_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/yaoapp/yao/agent/context"
 	"github.com/yaoapp/yao/agent/search"
 	"github.com/yaoapp/yao/agent/search/types"
+	"github.com/yaoapp/yao/agent/testutils"
 )
 
 func TestNewJSAPI(t *testing.T) {
@@ -51,6 +52,9 @@ func TestJSAPI_Web_WithOptions(t *testing.T) {
 }
 
 func TestJSAPI_KB(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
 	api := search.NewJSAPI(nil, &types.Config{
 		KB: &types.KBConfig{Collections: []string{"docs"}},
 	}, nil)
@@ -66,6 +70,9 @@ func TestJSAPI_KB(t *testing.T) {
 }
 
 func TestJSAPI_KB_WithOptions(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
 	api := search.NewJSAPI(nil, &types.Config{
 		KB: &types.KBConfig{Collections: []string{"docs"}},
 	}, nil)
@@ -87,6 +94,9 @@ func TestJSAPI_KB_WithOptions(t *testing.T) {
 }
 
 func TestJSAPI_DB(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
 	api := search.NewJSAPI(nil, &types.Config{
 		DB: &types.DBConfig{Models: []string{"product"}},
 	}, &search.Uses{QueryDSL: "builtin"})
@@ -102,6 +112,9 @@ func TestJSAPI_DB(t *testing.T) {
 }
 
 func TestJSAPI_DB_WithOptions(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
 	api := search.NewJSAPI(nil, &types.Config{
 		DB: &types.DBConfig{Models: []string{"product"}},
 	}, &search.Uses{QueryDSL: "builtin"})
@@ -122,6 +135,9 @@ func TestJSAPI_DB_WithOptions(t *testing.T) {
 }
 
 func TestJSAPI_All(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
 	api := search.NewJSAPI(nil, &types.Config{
 		KB: &types.KBConfig{Collections: []string{"docs"}},
 		DB: &types.DBConfig{Models: []string{"product"}},
@@ -155,6 +171,9 @@ func TestJSAPI_All(t *testing.T) {
 }
 
 func TestJSAPI_Any(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
 	api := search.NewJSAPI(nil, &types.Config{
 		KB: &types.KBConfig{Collections: []string{"docs"}},
 		DB: &types.DBConfig{Models: []string{"product"}},
@@ -186,6 +205,9 @@ func TestJSAPI_Any(t *testing.T) {
 }
 
 func TestJSAPI_Race(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
 	api := search.NewJSAPI(nil, &types.Config{
 		KB: &types.KBConfig{Collections: []string{"docs"}},
 		DB: &types.DBConfig{Models: []string{"product"}},

agent/search/search.go

@@ -151,14 +151,24 @@ func (s *Searcher) parallelAny(ctx *context.Context, reqs []*types.Request) ([]*
 		wg.Add(1)
 		go func(idx int, r *types.Request) {
 			defer wg.Done()
-			result, _ := s.Search(ctx, r)
+
+			// Check if done before starting
 			select {
+			case <-done:
+				return
+			default:
+			}
+
+			result, _ := s.Search(ctx, r)
+
+			// Try to send result
+			select {
+			case <-done:
+				// Already found a successful result
 			case resultChan <- struct {
 				idx    int
 				result *types.Result
 			}{idx, result}:
-			case <-done:
-				// Already found a successful result, discard this one
 			}
 		}(i, req)
 	}
@@ -170,25 +180,23 @@ func (s *Searcher) parallelAny(ctx *context.Context, reqs []*types.Request) ([]*
 	}()
 
 	// Collect results until we find one with items (success)
-	var mu sync.Mutex
+	var foundSuccess bool
 	for res := range resultChan {
-		mu.Lock()
 		results[res.idx] = res.result
 		// Check if this result has items (success = has results and no error)
-		if res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" {
+		if !foundSuccess && res.result != nil && len(res.result.Items) > 0 && res.result.Error == "" {
-			mu.Unlock()
+			foundSuccess = true
-			close(done) // Signal other goroutines to stop sending
+			close(done) // Signal other goroutines to stop
-			return results, nil
 		}
-		mu.Unlock()
 	}
 
-	// No successful result found, return all results
+	// All goroutines have completed (resultChan is closed)
 	return results, nil
 }
 
 // parallelRace returns as soon as any search completes (like Promise.race)
 // Returns immediately when first result arrives, regardless of success/failure
+// Note: Still waits for all goroutines to complete before returning to avoid resource leaks
 func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]*types.Result, error) {
 	results := make([]*types.Result, len(reqs))
 	resultChan := make(chan struct {
@@ -203,14 +211,24 @@ func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]
 		wg.Add(1)
 		go func(idx int, r *types.Request) {
 			defer wg.Done()
-			result, _ := s.Search(ctx, r)
+
+			// Check if done before starting
 			select {
+			case <-done:
+				return
+			default:
+			}
+
+			result, _ := s.Search(ctx, r)
+
+			// Try to send result
+			select {
+			case <-done:
+				// Already got first result
 			case resultChan <- struct {
 				idx    int
 				result *types.Result
 			}{idx, result}:
-			case <-done:
-				// Already got first result, discard this one
 			}
 		}(i, req)
 	}
@@ -221,14 +239,17 @@ func (s *Searcher) parallelRace(ctx *context.Context, reqs []*types.Request) ([]
 		close(resultChan)
 	}()
 
-	// Return immediately when first result arrives
+	// Get first result and signal others to stop
-	if res, ok := <-resultChan; ok {
+	var gotFirst bool
+	for res := range resultChan {
 		results[res.idx] = res.result
-		close(done) // Signal other goroutines to stop sending
+		if !gotFirst {
-		return results, nil
+			gotFirst = true
+			close(done) // Signal other goroutines to stop
+		}
 	}
 
-	// No results (shouldn't happen with valid requests)
+	// All goroutines have completed (resultChan is closed)
 	return results, nil
 }
 

agent/search/types/types.go

@@ -48,9 +48,10 @@ type Request struct {
 	TimeRange string   `json:"time_range,omitempty"` // "day", "week", "month", "year"
 
 	// Knowledge base specific
-	Collections []string `json:"collections,omitempty"` // KB collection IDs
+	Collections []string               `json:"collections,omitempty"` // KB collection IDs
-	Threshold   float64  `json:"threshold,omitempty"`   // Similarity threshold (0-1)
+	Threshold   float64                `json:"threshold,omitempty"`   // Similarity threshold (0-1)
-	Graph       bool     `json:"graph,omitempty"`       // Enable graph association
+	Graph       bool                   `json:"graph,omitempty"`       // Enable graph association
+	Metadata    map[string]interface{} `json:"metadata,omitempty"`    // Metadata filter for KB search
 
 	// Database search specific
 	Models   []string     `json:"models,omitempty"`   // Model IDs (e.g., "user", "agents.mybot.product")

kb/api/README.md

@@ -184,10 +184,10 @@ result, err := kb.API.RemoveDocuments(ctx, params)
 
 The Search API supports batch queries with three search modes:
 
-| Mode | Description |
+| Mode     | Description                                            |
-|------|-------------|
+| -------- | ------------------------------------------------------ |
-| `vector` | Pure vector similarity search |
+| `vector` | Pure vector similarity search                          |
-| `graph` | Graph traversal to find related segments via entities |
+| `graph`  | Graph traversal to find related segments via entities  |
 | `expand` | Graph-based entity expansion + vector search (default) |
 
 ### Basic Vector Search
@@ -274,7 +274,7 @@ queries := []api.Query{
         CollectionID: "my_collection",
         Input:        "physics",
         DocumentID:   "specific_doc_id", // filter to specific document
-        MinScore:     0.5,               // minimum similarity score
+        Threshold:    0.5,               // similarity threshold
         Metadata: map[string]interface{}{
             "category": "science",
         },
@@ -288,18 +288,18 @@ result, err := kb.API.Search(ctx, queries)
 
 ## Query Parameters
 
-| Field | Type | Description |
+| Field          | Type          | Description                                            |
-|-------|------|-------------|
+| -------------- | ------------- | ------------------------------------------------------ |
-| `CollectionID` | string | Collection to search (required) |
+| `CollectionID` | string        | Collection to search (required)                        |
-| `Input` | string | Direct query text |
+| `Input`        | string        | Direct query text                                      |
-| `Messages` | []ChatMessage | Conversation history (last user message used as query) |
+| `Messages`     | []ChatMessage | Conversation history (last user message used as query) |
-| `Mode` | SearchMode | `vector`, `graph`, or `expand` (default: `expand`) |
+| `Mode`         | SearchMode    | `vector`, `graph`, or `expand` (default: `expand`)     |
-| `DocumentID` | string | Filter to specific document |
+| `DocumentID`   | string        | Filter to specific document                            |
-| `MinScore` | float64 | Minimum similarity threshold |
+| `Threshold`    | float64       | Similarity threshold (0-1)                             |
-| `Metadata` | map | Filter by metadata fields |
+| `Metadata`     | map           | Filter by metadata fields                              |
-| `MaxDepth` | int | Graph traversal depth (default: 2) |
+| `MaxDepth`     | int           | Graph traversal depth (default: 2)                     |
-| `Page` | int | Page number (1-based) |
+| `Page`         | int           | Page number (1-based)                                  |
-| `PageSize` | int | Results per page |
+| `PageSize`     | int           | Results per page                                       |
 
 ## Search Result
 
@@ -328,7 +328,7 @@ type ProviderConfigParams struct {
 ```
 
 Common providers:
+
 - **Chunking**: `__yao.structured` - text splitting
 - **Embedding**: `__yao.openai` - vector embeddings
 - **Extraction**: `__yao.openai` - entity/relationship extraction for graph
-

kb/api/search.go

@@ -15,7 +15,7 @@ import (
 const (
 	DefaultSearchK        = 10
 	DefaultMaxDepth       = 2
-	DefaultMinScore       = 0.0
+	DefaultThreshold      = 0.0
 	MaxSearchK            = 100
 	DefaultSearchPageSize = 20
 )
@@ -257,7 +257,7 @@ func (kb *KBInstance) searchVector(ctx context.Context, collectionID string, que
 		DocumentID:   query.DocumentID,
 		Query:        queryText,
 		K:            k,
-		MinScore:     query.MinScore,
+		MinScore:     query.Threshold,
 		Embedding:    embedding,
 	}
 
@@ -356,7 +356,7 @@ func (kb *KBInstance) searchExpand(ctx context.Context, collectionID string, que
 		DocumentID:   query.DocumentID,
 		Query:        queryText,
 		K:            k,
-		MinScore:     query.MinScore,
+		MinScore:     query.Threshold,
 		Embedding:    embedding,
 	}
 

kb/api/search_test.go

@@ -314,30 +314,30 @@ func TestSearchQuery(t *testing.T) {
 		}
 	})
 
-	t.Run("Search_WithMinScore", func(t *testing.T) {
+	t.Run("Search_WithThreshold", func(t *testing.T) {
-		// Test: Filter by minimum score
+		// Test: Filter by similarity threshold
 		queries := []api.Query{
 			{
 				CollectionID: SearchTestScienceCollection,
 				Input:        "Einstein relativity",
 				Mode:         api.SearchModeVector,
-				MinScore:     0.5,
+				Threshold:    0.5,
 				PageSize:     10,
 			},
 		}
 
 		result, err := kb.API.Search(ctx, queries)
 		if err != nil {
-			t.Logf("MinScore search error: %v", err)
+			t.Logf("Threshold search error: %v", err)
 			return
 		}
 
 		assert.NotNil(t, result)
-		t.Logf("MinScore search returned %d segments", len(result.Segments))
+		t.Logf("Threshold search returned %d segments", len(result.Segments))
 
-		// Verify all results meet minimum score
+		// Verify all results meet threshold
 		for _, seg := range result.Segments {
-			assert.GreaterOrEqual(t, seg.Score, 0.5, "All segments should meet minimum score")
+			assert.GreaterOrEqual(t, seg.Score, 0.5, "All segments should meet threshold")
 		}
 	})
 

kb/api/types.go

@@ -233,8 +233,8 @@ type Query struct {
 	// DocumentID filters results to a specific document (optional)
 	DocumentID string `json:"document_id,omitempty" yaml:"document_id,omitempty"`
 
-	// MinScore filters results below this similarity threshold (optional)
+	// Threshold filters results below this similarity threshold (optional)
-	MinScore float64 `json:"min_score,omitempty" yaml:"min_score,omitempty"`
+	Threshold float64 `json:"threshold,omitempty" yaml:"threshold,omitempty"`
 
 	// Metadata filters segments by metadata fields (optional)
 	Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"`