Implement Auto Search Feature in Assistant Stream Method

- Added functionality to the Assistant's Stream method to execute auto search if enabled, enhancing the search capabilities based on user configuration.
- Introduced helper methods for determining auto search eligibility, executing the search, and injecting search context into messages.
- Updated DESIGN.md to reflect the new auto search logic, including detailed descriptions of the new methods and their integration points within the search process.

agent/assistant/agent.go

@@ -199,6 +199,16 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 			return nil, err
 		}
 
+		// ================================================
+		// Execute Auto Search (if enabled)
+		// ================================================
+		if ast.shouldAutoSearch(ctx, createResponse) {
+			refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse)
+			if refCtx != nil && len(refCtx.References) > 0 {
+				completionMessages = ast.injectSearchContext(completionMessages, refCtx)
+			}
+		}
+
 		// Begin step tracking for LLM call
 		ast.BeginStep(ctx, context.StepTypeLLM, map[string]interface{}{
 			"messages": completionMessages,

agent/assistant/search.go

@@ -0,0 +1,275 @@
+package assistant
+
+import (
+	"strings"
+
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/search"
+	searchTypes "github.com/yaoapp/yao/agent/search/types"
+)
+
+// shouldAutoSearch determines if auto search should be executed
+// Returns false if:
+// - uses.search is "disabled"
+// - assistant has no search configuration
+func (ast *Assistant) shouldAutoSearch(ctx *context.Context, createResponse *context.HookCreateResponse) bool {
+	// Get merged uses configuration
+	uses := ast.getMergedSearchUses(createResponse)
+
+	// Check if search is explicitly disabled
+	if uses != nil && uses.Search == "disabled" {
+		ctx.Logger.Info("Auto search disabled by uses.search=disabled")
+		return false
+	}
+
+	// Check if assistant has search configuration
+	if ast.Search == nil && (uses == nil || uses.Search == "") {
+		return false
+	}
+
+	// Check if search is enabled (builtin, agent, mcp, or empty means builtin)
+	return true
+}
+
+// getMergedSearchUses returns the merged uses configuration for search
+// Priority: createResponse > assistant
+func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResponse) *context.Uses {
+	// Start with assistant uses
+	var uses *context.Uses
+	if ast.Uses != nil {
+		uses = &context.Uses{
+			Search:   ast.Uses.Search,
+			Web:      ast.Uses.Web,
+			Keyword:  ast.Uses.Keyword,
+			QueryDSL: ast.Uses.QueryDSL,
+			Rerank:   ast.Uses.Rerank,
+		}
+	}
+
+	// Override with createResponse.Uses if provided (highest priority)
+	if createResponse != nil && createResponse.Uses != nil {
+		if uses == nil {
+			uses = &context.Uses{}
+		}
+		if createResponse.Uses.Search != "" {
+			uses.Search = createResponse.Uses.Search
+		}
+		if createResponse.Uses.Web != "" {
+			uses.Web = createResponse.Uses.Web
+		}
+		if createResponse.Uses.Keyword != "" {
+			uses.Keyword = createResponse.Uses.Keyword
+		}
+		if createResponse.Uses.QueryDSL != "" {
+			uses.QueryDSL = createResponse.Uses.QueryDSL
+		}
+		if createResponse.Uses.Rerank != "" {
+			uses.Rerank = createResponse.Uses.Rerank
+		}
+	}
+
+	return uses
+}
+
+// executeAutoSearch executes auto search based on configuration
+// Returns ReferenceContext with results and formatted context
+func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) *searchTypes.ReferenceContext {
+	ctx.Logger.Phase("Search")
+	defer ctx.Logger.PhaseComplete("Search")
+
+	// Get merged uses configuration
+	uses := ast.getMergedSearchUses(createResponse)
+
+	// Convert to search.Uses
+	searchUses := &search.Uses{}
+	if uses != nil {
+		searchUses.Search = uses.Search
+		searchUses.Web = uses.Web
+		searchUses.Keyword = uses.Keyword
+		searchUses.QueryDSL = uses.QueryDSL
+		searchUses.Rerank = uses.Rerank
+	}
+
+	// Get merged search config
+	searchConfig := ast.GetMergedSearchConfig()
+
+	// Create searcher
+	searcher := search.New(searchConfig, searchUses)
+
+	// Extract query from messages
+	query := extractQueryFromMessages(messages)
+	if query == "" {
+		ctx.Logger.Info("No query found in messages, skipping auto search")
+		return nil
+	}
+
+	// Build search requests based on configuration
+	requests := ast.buildSearchRequests(query, searchConfig)
+	if len(requests) == 0 {
+		ctx.Logger.Info("No search requests to execute")
+		return nil
+	}
+
+	// Execute searches in parallel
+	ctx.Logger.Info("Executing %d search requests for query: %s", len(requests), truncateString(query, 50))
+
+	results, err := searcher.All(ctx, requests)
+	if err != nil {
+		// Log error but don't fail - search errors shouldn't block the main flow
+		ctx.Logger.Error("Auto search failed: %v", err)
+		return nil
+	}
+
+	// Build reference context (includes references, XML, and prompt)
+	var citationConfig *searchTypes.CitationConfig
+	if searchConfig != nil {
+		citationConfig = searchConfig.Citation
+	}
+	refCtx := search.BuildReferenceContext(results, citationConfig)
+
+	if len(refCtx.References) == 0 {
+		ctx.Logger.Info("No search results found")
+		return nil
+	}
+
+	ctx.Logger.Info("Auto search completed: %d references", len(refCtx.References))
+	return refCtx
+}
+
+// buildSearchRequests builds search requests based on assistant configuration
+func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Config) []*searchTypes.Request {
+	var requests []*searchTypes.Request
+
+	// Web search - check if web search is configured
+	if config != nil && config.Web != nil {
+		requests = append(requests, &searchTypes.Request{
+			Type:   searchTypes.SearchTypeWeb,
+			Query:  query,
+			Source: searchTypes.SourceAuto,
+			Limit:  config.Web.MaxResults,
+		})
+	}
+
+	// KB search - check if KB is configured
+	if ast.KB != nil && len(ast.KB.Collections) > 0 {
+		limit := 10
+		threshold := 0.7
+		if config != nil && config.KB != nil {
+			if config.KB.Threshold > 0 {
+				threshold = config.KB.Threshold
+			}
+		}
+		requests = append(requests, &searchTypes.Request{
+			Type:        searchTypes.SearchTypeKB,
+			Query:       query,
+			Source:      searchTypes.SourceAuto,
+			Limit:       limit,
+			Collections: ast.KB.Collections,
+			Threshold:   threshold,
+			Graph:       config != nil && config.KB != nil && config.KB.Graph,
+		})
+	}
+
+	// DB search - check if DB is configured
+	if ast.DB != nil && len(ast.DB.Models) > 0 {
+		limit := 20
+		if config != nil && config.DB != nil && config.DB.MaxResults > 0 {
+			limit = config.DB.MaxResults
+		}
+		requests = append(requests, &searchTypes.Request{
+			Type:   searchTypes.SearchTypeDB,
+			Query:  query,
+			Source: searchTypes.SourceAuto,
+			Limit:  limit,
+			Models: ast.DB.Models,
+		})
+	}
+
+	return requests
+}
+
+// injectSearchContext injects search results into messages
+// Adds search context as a system message after existing system messages
+func (ast *Assistant) injectSearchContext(messages []context.Message, refCtx *searchTypes.ReferenceContext) []context.Message {
+	if refCtx == nil || len(refCtx.References) == 0 {
+		return messages
+	}
+
+	// Build the search context message
+	var contentParts []string
+
+	// Add citation prompt
+	if refCtx.Prompt != "" {
+		contentParts = append(contentParts, refCtx.Prompt)
+	}
+
+	// Add XML context
+	if refCtx.XML != "" {
+		contentParts = append(contentParts, refCtx.XML)
+	}
+
+	if len(contentParts) == 0 {
+		return messages
+	}
+
+	// Create system message with search context
+	searchMessage := context.Message{
+		Role:    "system",
+		Content: strings.Join(contentParts, "\n\n"),
+	}
+
+	// Find the position to insert the search message
+	// Insert after any existing system messages but before user messages
+	insertIndex := 0
+	for i, msg := range messages {
+		if msg.Role == "system" {
+			insertIndex = i + 1
+		} else {
+			break
+		}
+	}
+
+	// Insert the search message
+	result := make([]context.Message, 0, len(messages)+1)
+	result = append(result, messages[:insertIndex]...)
+	result = append(result, searchMessage)
+	result = append(result, messages[insertIndex:]...)
+
+	return result
+}
+
+// extractQueryFromMessages extracts the search query from messages
+// Uses the last user message as the query
+func extractQueryFromMessages(messages []context.Message) string {
+	// Find the last user message
+	for i := len(messages) - 1; i >= 0; i-- {
+		if messages[i].Role == "user" {
+			content := messages[i].Content
+			// Handle string content
+			if str, ok := content.(string); ok {
+				return str
+			}
+			// Handle content parts (array of objects)
+			if parts, ok := content.([]interface{}); ok {
+				for _, part := range parts {
+					if partMap, ok := part.(map[string]interface{}); ok {
+						if partMap["type"] == "text" {
+							if text, ok := partMap["text"].(string); ok {
+								return text
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+	return ""
+}
+
+// truncateString truncates a string to maxLen characters
+func truncateString(s string, maxLen int) string {
+	if len(s) <= maxLen {
+		return s
+	}
+	return s[:maxLen] + "..."
+}

agent/assistant/search_auto_disabled_test.go

@@ -0,0 +1,84 @@
+package assistant_test
+
+import (
+	stdContext "context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/yaoapp/yao/agent/assistant"
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/output/message"
+	"github.com/yaoapp/yao/agent/testutils"
+	"github.com/yaoapp/yao/openapi/oauth/types"
+)
+
+// newSearchAutoDisabledTestContext creates a test context
+func newSearchAutoDisabledTestContext(chatID, assistantID string) *context.Context {
+	authorized := &types.AuthorizedInfo{
+		Subject:  "test-user",
+		UserID:   "test-123",
+		TenantID: "test-tenant",
+	}
+
+	ctx := context.New(stdContext.Background(), authorized, chatID)
+	ctx.ID = chatID
+	ctx.AssistantID = assistantID
+	ctx.Locale = "en-us"
+	ctx.Client = context.Client{
+		Type: "web",
+		IP:   "127.0.0.1",
+	}
+	ctx.Referer = context.RefererAPI
+	ctx.Accept = context.AcceptWebCUI
+	ctx.IDGenerator = message.NewIDGenerator()
+	ctx.Metadata = make(map[string]interface{})
+	return ctx
+}
+
+func TestSearchAutoDisabled(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.LoadPath("/assistants/tests/search-auto-disabled")
+	require.NoError(t, err)
+	require.NotNil(t, ast)
+
+	t.Run("ShouldHaveSearchConfig", func(t *testing.T) {
+		// Search config is set but uses.search is disabled
+		assert.NotNil(t, ast.Search, "search config should be set")
+		assert.NotNil(t, ast.Search.Web, "web search config should be set")
+	})
+
+	t.Run("ShouldHaveDisabledUses", func(t *testing.T) {
+		assert.NotNil(t, ast.Uses, "uses config should be set")
+		assert.Equal(t, "disabled", ast.Uses.Search, "uses.search should be disabled")
+	})
+
+	t.Run("StreamShouldNotExecuteSearch", func(t *testing.T) {
+		// Get agent via assistant.Get (required for Stream)
+		agent, err := assistant.Get("tests.search-auto-disabled")
+		require.NoError(t, err)
+		require.NotNil(t, agent)
+
+		// Create context
+		ctx := newSearchAutoDisabledTestContext("test-search-auto-disabled", "tests.search-auto-disabled")
+
+		// Create messages
+		messages := []context.Message{
+			{
+				Role:    "user",
+				Content: "Hello, how are you?",
+			},
+		}
+
+		// Execute stream - should NOT trigger search because uses.search is "disabled"
+		response, err := agent.Stream(ctx, messages)
+		require.NoError(t, err)
+		require.NotNil(t, response)
+
+		resp := response.(*context.Response)
+		assert.NotNil(t, resp.Completion, "should have completion")
+		t.Logf("✓ Stream executed without search (disabled)")
+	})
+}

agent/assistant/search_auto_full_test.go

@@ -0,0 +1,125 @@
+package assistant_test
+
+import (
+	stdContext "context"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/yaoapp/yao/agent/assistant"
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/output/message"
+	"github.com/yaoapp/yao/agent/testutils"
+	"github.com/yaoapp/yao/openapi/oauth/types"
+)
+
+// newSearchAutoFullTestContext creates a test context
+func newSearchAutoFullTestContext(chatID, assistantID string) *context.Context {
+	authorized := &types.AuthorizedInfo{
+		Subject:  "test-user",
+		UserID:   "test-123",
+		TenantID: "test-tenant",
+	}
+
+	ctx := context.New(stdContext.Background(), authorized, chatID)
+	ctx.ID = chatID
+	ctx.AssistantID = assistantID
+	ctx.Locale = "en-us"
+	ctx.Client = context.Client{
+		Type: "web",
+		IP:   "127.0.0.1",
+	}
+	ctx.Referer = context.RefererAPI
+	ctx.Accept = context.AcceptWebCUI
+	ctx.IDGenerator = message.NewIDGenerator()
+	ctx.Metadata = make(map[string]interface{})
+	return ctx
+}
+
+func TestSearchAutoFull(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.LoadPath("/assistants/tests/search-auto-full")
+	require.NoError(t, err)
+	require.NotNil(t, ast)
+
+	t.Run("ShouldHaveWebSearchConfig", func(t *testing.T) {
+		assert.NotNil(t, ast.Search, "search config should be set")
+		assert.NotNil(t, ast.Search.Web, "web search config should be set")
+		assert.Equal(t, "tavily", ast.Search.Web.Provider)
+		assert.Equal(t, 3, ast.Search.Web.MaxResults)
+	})
+
+	t.Run("ShouldHaveKBSearchConfig", func(t *testing.T) {
+		assert.NotNil(t, ast.Search.KB, "kb search config should be set")
+		assert.Equal(t, 0.7, ast.Search.KB.Threshold)
+		assert.False(t, ast.Search.KB.Graph)
+	})
+
+	t.Run("ShouldHaveDBSearchConfig", func(t *testing.T) {
+		assert.NotNil(t, ast.Search.DB, "db search config should be set")
+		assert.Equal(t, 10, ast.Search.DB.MaxResults)
+	})
+
+	t.Run("ShouldHaveKBCollections", func(t *testing.T) {
+		assert.NotNil(t, ast.KB, "kb config should be set")
+		assert.Contains(t, ast.KB.Collections, "test-collection")
+	})
+
+	t.Run("ShouldHaveDBModels", func(t *testing.T) {
+		assert.NotNil(t, ast.DB, "db config should be set")
+		assert.Contains(t, ast.DB.Models, "user")
+		assert.Contains(t, ast.DB.Models, "article")
+	})
+
+	t.Run("ShouldHaveCitationConfig", func(t *testing.T) {
+		assert.NotNil(t, ast.Search.Citation, "citation config should be set")
+		assert.Equal(t, "xml", ast.Search.Citation.Format)
+		assert.True(t, ast.Search.Citation.AutoInjectPrompt)
+	})
+
+	t.Run("ShouldHaveUsesConfig", func(t *testing.T) {
+		assert.NotNil(t, ast.Uses, "uses config should be set")
+		assert.Equal(t, "builtin", ast.Uses.Search)
+		assert.Equal(t, "builtin", ast.Uses.Web)
+	})
+
+	t.Run("StreamShouldExecuteMultipleSearchTypes", func(t *testing.T) {
+		// Get agent via assistant.Get (required for Stream)
+		agent, err := assistant.Get("tests.search-auto-full")
+		require.NoError(t, err)
+		require.NotNil(t, agent)
+
+		// Create context
+		ctx := newSearchAutoFullTestContext("test-search-auto-full", "tests.search-auto-full")
+
+		// Create messages with a search query
+		messages := []context.Message{
+			{
+				Role:    "user",
+				Content: "Find information about machine learning",
+			},
+		}
+
+		// Execute stream - should trigger Web + KB + DB searches
+		response, err := agent.Stream(ctx, messages)
+
+		// Assert no error (if API key is configured)
+		if err != nil {
+			// If error contains "API key", it's expected in CI without keys
+			if strings.Contains(err.Error(), "API key") || strings.Contains(err.Error(), "api_key") {
+				t.Logf("Expected error without API key: %v", err)
+				return
+			}
+			// Other errors should fail
+			require.NoError(t, err)
+		}
+
+		require.NotNil(t, response)
+		resp := response.(*context.Response)
+		assert.NotNil(t, resp.Completion, "should have completion")
+		t.Logf("✓ Stream executed with full search config (Web + KB + DB)")
+	})
+}

agent/assistant/search_auto_hook_disable_test.go

@@ -0,0 +1,108 @@
+package assistant_test
+
+import (
+	stdContext "context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/yaoapp/yao/agent/assistant"
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/output/message"
+	"github.com/yaoapp/yao/agent/testutils"
+	"github.com/yaoapp/yao/openapi/oauth/types"
+)
+
+// newSearchAutoHookDisableTestContext creates a test context
+func newSearchAutoHookDisableTestContext(chatID, assistantID string) *context.Context {
+	authorized := &types.AuthorizedInfo{
+		Subject:  "test-user",
+		UserID:   "test-123",
+		TenantID: "test-tenant",
+	}
+
+	ctx := context.New(stdContext.Background(), authorized, chatID)
+	ctx.ID = chatID
+	ctx.AssistantID = assistantID
+	ctx.Locale = "en-us"
+	ctx.Client = context.Client{
+		Type: "web",
+		IP:   "127.0.0.1",
+	}
+	ctx.Referer = context.RefererAPI
+	ctx.Accept = context.AcceptWebCUI
+	ctx.IDGenerator = message.NewIDGenerator()
+	ctx.Metadata = make(map[string]interface{})
+	return ctx
+}
+
+func TestSearchAutoHookDisable(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.LoadPath("/assistants/tests/search-auto-hook-disable")
+	require.NoError(t, err)
+	require.NotNil(t, ast)
+
+	t.Run("ShouldHaveSearchConfigEnabled", func(t *testing.T) {
+		// Search config is enabled in package.yao
+		assert.NotNil(t, ast.Search, "search config should be set")
+		assert.NotNil(t, ast.Uses, "uses config should be set")
+		assert.Equal(t, "builtin", ast.Uses.Search, "uses.search should be builtin in config")
+	})
+
+	t.Run("ShouldHaveHookScript", func(t *testing.T) {
+		// Hook script should be loaded
+		assert.NotNil(t, ast.HookScript, "hook script should be loaded")
+	})
+
+	t.Run("HookShouldDisableSearch", func(t *testing.T) {
+		// Create context
+		ctx := newSearchAutoHookDisableTestContext("test-chat-id", "tests.search-auto-hook-disable")
+
+		// Create messages
+		messages := []context.Message{
+			{
+				Role:    "user",
+				Content: "Test message",
+			},
+		}
+
+		// Call Create hook directly
+		opts := &context.Options{}
+		response, _, err := ast.HookScript.Create(ctx, messages, opts)
+		require.NoError(t, err)
+		require.NotNil(t, response)
+
+		// Verify hook returns uses.search = "disabled"
+		assert.NotNil(t, response.Uses, "hook should return uses")
+		assert.Equal(t, "disabled", response.Uses.Search, "hook should disable search")
+	})
+
+	t.Run("StreamShouldRespectHookDisable", func(t *testing.T) {
+		// Get agent via assistant.Get (required for Stream)
+		agent, err := assistant.Get("tests.search-auto-hook-disable")
+		require.NoError(t, err)
+		require.NotNil(t, agent)
+
+		// Create context
+		ctx := newSearchAutoHookDisableTestContext("test-search-hook-disable", "tests.search-auto-hook-disable")
+
+		// Create messages
+		messages := []context.Message{
+			{
+				Role:    "user",
+				Content: "What is AI?",
+			},
+		}
+
+		// Execute stream - hook will disable search
+		response, err := agent.Stream(ctx, messages)
+		require.NoError(t, err)
+		require.NotNil(t, response)
+
+		resp := response.(*context.Response)
+		assert.NotNil(t, resp.Completion, "should have completion")
+		t.Logf("✓ Stream executed with hook disabling search")
+	})
+}

agent/assistant/search_auto_web_test.go

@@ -0,0 +1,103 @@
+package assistant_test
+
+import (
+	stdContext "context"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/yaoapp/yao/agent/assistant"
+	"github.com/yaoapp/yao/agent/context"
+	"github.com/yaoapp/yao/agent/output/message"
+	"github.com/yaoapp/yao/agent/testutils"
+	"github.com/yaoapp/yao/openapi/oauth/types"
+)
+
+// newSearchAutoTestContext creates a test context for search auto tests
+func newSearchAutoTestContext(chatID, assistantID string) *context.Context {
+	authorized := &types.AuthorizedInfo{
+		Subject:  "test-user",
+		UserID:   "test-123",
+		TenantID: "test-tenant",
+	}
+
+	ctx := context.New(stdContext.Background(), authorized, chatID)
+	ctx.ID = chatID
+	ctx.AssistantID = assistantID
+	ctx.Locale = "en-us"
+	ctx.Client = context.Client{
+		Type: "web",
+		IP:   "127.0.0.1",
+	}
+	ctx.Referer = context.RefererAPI
+	ctx.Accept = context.AcceptWebCUI
+	ctx.IDGenerator = message.NewIDGenerator()
+	ctx.Metadata = make(map[string]interface{})
+	return ctx
+}
+
+func TestSearchAutoWeb(t *testing.T) {
+	testutils.Prepare(t)
+	defer testutils.Clean(t)
+
+	ast, err := assistant.LoadPath("/assistants/tests/search-auto-web")
+	require.NoError(t, err)
+	require.NotNil(t, ast)
+
+	t.Run("ShouldHaveSearchConfig", func(t *testing.T) {
+		assert.NotNil(t, ast.Search, "search config should be set")
+		assert.NotNil(t, ast.Search.Web, "web search config should be set")
+		assert.Equal(t, "tavily", ast.Search.Web.Provider)
+		assert.Equal(t, 3, ast.Search.Web.MaxResults)
+	})
+
+	t.Run("ShouldHaveUsesConfig", func(t *testing.T) {
+		assert.NotNil(t, ast.Uses, "uses config should be set")
+		assert.Equal(t, "builtin", ast.Uses.Search)
+		assert.Equal(t, "builtin", ast.Uses.Web)
+	})
+
+	t.Run("ShouldHaveCitationConfig", func(t *testing.T) {
+		assert.NotNil(t, ast.Search.Citation, "citation config should be set")
+		assert.Equal(t, "xml", ast.Search.Citation.Format)
+		assert.True(t, ast.Search.Citation.AutoInjectPrompt)
+	})
+
+	t.Run("StreamShouldExecuteAutoSearch", func(t *testing.T) {
+		// Get agent via assistant.Get (required for Stream)
+		agent, err := assistant.Get("tests.search-auto-web")
+		require.NoError(t, err)
+		require.NotNil(t, agent)
+
+		// Create context
+		ctx := newSearchAutoTestContext("test-search-auto-web", "tests.search-auto-web")
+
+		// Create messages with a search query
+		messages := []context.Message{
+			{
+				Role:    "user",
+				Content: "What is the latest news about artificial intelligence?",
+			},
+		}
+
+		// Execute stream
+		response, err := agent.Stream(ctx, messages)
+
+		// Assert no error (if API key is configured)
+		if err != nil {
+			// If error contains "API key", it's expected in CI without keys
+			if strings.Contains(err.Error(), "API key") || strings.Contains(err.Error(), "api_key") {
+				t.Logf("Expected error without API key: %v", err)
+				return
+			}
+			// Other errors should fail
+			require.NoError(t, err)
+		}
+
+		require.NotNil(t, response)
+		resp := response.(*context.Response)
+		assert.NotNil(t, resp.Completion, "should have completion")
+		t.Logf("✓ Stream executed successfully with auto search")
+	})
+}

agent/search/DESIGN.md

@@ -1552,22 +1552,62 @@ Stream(ctx, messages, options)
   ├── 2. Create Hook (optional)
   │   └── Can call ctx.search.* and return search results
   │
-  ├── 3. Auto Search Decision
+  ├── 3. BuildRequest + BuildContent
+  │
+  ├── 4. Auto Search Decision (shouldAutoSearch)
   │   ├── IF Uses.Search == "disabled" → SKIP
   │   ├── IF Create Hook returned uses.search="disabled" → SKIP
-  │   └── ELSE → Execute Auto Search (based on Uses.Search mode)
+  │   └── ELSE → Execute Auto Search (executeAutoSearch)
-  │       ├── Read assistant's search config
+  │       ├── Read assistant's search config (GetMergedSearchConfig)
-  │       ├── Execute web/kb/db in parallel
+  │       ├── Build search requests (buildSearchRequests)
-  │       ├── Send search_start/search_result/search_complete to output
+  │       ├── Execute web/kb/db in parallel (searcher.All)
-  │       ├── Rerank results
+  │       ├── Build reference context (BuildReferenceContext)
-  │       ├── Generate citation IDs
+  │       └── Inject search context to messages (injectSearchContext)
-  │       └── Inject search context + citation prompt to messages
   │
-  ├── 4. LLM Call (with search context if any)
+  ├── 5. LLM Call (with search context if any)
   │
-  ├── 5. Next Hook (optional)
+  ├── 6. Next Hook (optional)
   │
-  └── 6. Output (response may contain #ref:xxx citations)
+  └── 7. Output (response may contain #ref:xxx citations)
+```
+
+**Implementation Files:**
+
+| File                  | Description                                     |
+| --------------------- | ----------------------------------------------- |
+| `assistant/search.go` | Core integration logic (shouldAutoSearch, etc.) |
+| `assistant/agent.go`  | Stream() integration point (after BuildContent) |
+| `search/reference.go` | BuildReferenceContext, FormatReferencesXML      |
+
+**Key Functions (`assistant/search.go`):**
+
+```go
+// shouldAutoSearch determines if auto search should be executed
+func (ast *Assistant) shouldAutoSearch(ctx *context.Context, createResponse *context.HookCreateResponse) bool
+
+// executeAutoSearch executes auto search based on configuration
+func (ast *Assistant) executeAutoSearch(ctx *context.Context, messages []context.Message, createResponse *context.HookCreateResponse) *searchTypes.ReferenceContext
+
+// injectSearchContext injects search results into messages
+func (ast *Assistant) injectSearchContext(messages []context.Message, refCtx *searchTypes.ReferenceContext) []context.Message
+
+// getMergedSearchUses returns the merged uses configuration for search
+func (ast *Assistant) getMergedSearchUses(createResponse *context.HookCreateResponse) *context.Uses
+
+// buildSearchRequests builds search requests based on assistant configuration
+func (ast *Assistant) buildSearchRequests(query string, config *searchTypes.Config) []*searchTypes.Request
+```
+
+**Integration in agent.go:**
+
+```go
+// In Stream(), after BuildContent:
+if ast.shouldAutoSearch(ctx, createResponse) {
+    refCtx := ast.executeAutoSearch(ctx, completionMessages, createResponse)
+    if refCtx != nil && len(refCtx.References) > 0 {
+        completionMessages = ast.injectSearchContext(completionMessages, refCtx)
+    }
+}
 ```
 
 ### Control via Uses.Search