Simplify RAG event forwarding and clean up RAGTool

- Remove RAGInitializer interface and StartBackgroundRAGInit indirection.
  RAG event callbacks are now wired in configureToolsetHandlers alongside
  other handler setup, using the same pattern as Elicitable/OAuthCapable.

- Remove deprecated NewManagers wrapper (no callers after toolset refactor).

- Clean up RAGTool: unexport internal types (QueryRAGArgs, QueryResult),
  inline sortResults, remove verbose debug logging from Tools(), simplify
  handleQueryRAG.

Assisted-By: docker-agent

Author: dgageot

pkg/app/app.go

@@ -30,12 +30,6 @@ import (
 	"github.com/docker/docker-agent/pkg/tui/messages"
 )
 
-// RAGInitializer is implemented by runtimes that support background RAG initialization.
-// Local runtimes use this to start indexing early; remote runtimes typically do not.
-type RAGInitializer interface {
-	StartBackgroundRAGInit(ctx context.Context, sendEvent func(runtime.Event))
-}
-
 type App struct {
 	runtime                runtime.Runtime
 	session                *session.Session
@@ -122,18 +116,6 @@ func New(ctx context.Context, rt runtime.Runtime, sess *session.Session, opts ..
 		}
 	}()
 
-	// If the runtime supports background RAG initialization, start it
-	// and forward events to the TUI. Remote runtimes typically handle RAG server-side
-	// and won't implement this optional interface.
-	if ragRuntime, ok := rt.(RAGInitializer); ok {
-		go ragRuntime.StartBackgroundRAGInit(ctx, func(event runtime.Event) {
-			select {
-			case app.events <- event:
-			case <-ctx.Done():
-			}
-		})
-	}
-
 	// Subscribe to tool list changes so the sidebar updates immediately
 	// when an MCP server adds or removes tools (outside of a RunStream).
 	if tcs, ok := rt.(runtime.ToolsChangeSubscriber); ok {

pkg/rag/builder.go

@@ -24,30 +24,6 @@ type ManagersBuildConfig struct {
 	Models        map[string]latest.ModelConfig // Model configurations from config
 }
 
-// NewManagers constructs all RAG managers defined in the config.
-//
-// Deprecated: Use NewManager for per-toolset creation instead.
-func NewManagers(ctx context.Context, cfg *latest.Config, buildCfg ManagersBuildConfig) ([]*Manager, error) {
-	if len(cfg.RAG) == 0 {
-		return nil, nil
-	}
-
-	var managers []*Manager
-
-	for ragName, ragToolset := range cfg.RAG {
-		if ragToolset.RAGConfig == nil {
-			continue
-		}
-		mgr, err := NewManager(ctx, ragName, ragToolset.RAGConfig, buildCfg)
-		if err != nil {
-			return nil, err
-		}
-		managers = append(managers, mgr)
-	}
-
-	return managers, nil
-}
-
 // NewManager constructs a single RAG manager from a RAGConfig.
 func NewManager(
 	ctx context.Context,

pkg/runtime/loop.go

@@ -96,11 +96,6 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 		// Emit team information
 		events <- TeamInfo(r.agentDetailsFromTeam(), a.Name())
 
-		// Initialize RAG and forward events
-		r.StartBackgroundRAGInit(ctx, func(event Event) {
-			events <- event
-		})
-
 		r.emitAgentWarnings(a, chanSend(events))
 		r.configureToolsetHandlers(a, events)
 
@@ -534,6 +529,11 @@ func (r *LocalRuntime) configureToolsetHandlers(a *agent.Agent, events chan Even
 			func() { events <- Authorization(tools.ElicitationActionAccept, a.Name()) },
 			r.managedOAuth,
 		)
+
+		// Wire RAG event forwarding so the TUI shows indexing progress.
+		if ragTool, ok := tools.As[*builtin.RAGTool](toolset); ok {
+			ragTool.SetEventCallback(ragEventForwarder(ragTool.Name(), r, chanSend(events)))
+		}
 	}
 }
 

pkg/runtime/rag.go

@@ -1,36 +1,15 @@
 package runtime
 
 import (
-	"context"
 	"fmt"
 	"log/slog"
 
 	ragtypes "github.com/docker/docker-agent/pkg/rag/types"
-	"github.com/docker/docker-agent/pkg/tools"
 	"github.com/docker/docker-agent/pkg/tools/builtin"
 )
 
-// StartBackgroundRAGInit discovers RAG toolsets from agents and wires up event
-// forwarding so the TUI can display indexing progress. Actual initialization
-// happens lazily when the tool is first used (via tools.Startable).
-func (r *LocalRuntime) StartBackgroundRAGInit(ctx context.Context, sendEvent func(Event)) {
-	for _, name := range r.team.AgentNames() {
-		a, err := r.team.Agent(name)
-		if err != nil {
-			continue
-		}
-		for _, ts := range a.ToolSets() {
-			ragTool, ok := tools.As[*builtin.RAGTool](ts)
-			if !ok {
-				continue
-			}
-			ragTool.SetEventCallback(ragEventForwarder(ctx, ragTool.Name(), r, sendEvent))
-		}
-	}
-}
-
 // ragEventForwarder returns a callback that converts RAG manager events to runtime events.
-func ragEventForwarder(ctx context.Context, ragName string, r *LocalRuntime, sendEvent func(Event)) builtin.RAGEventCallback {
+func ragEventForwarder(ragName string, r *LocalRuntime, sendEvent func(Event)) builtin.RAGEventCallback {
 	return func(ragEvent ragtypes.Event) {
 		agentName := r.CurrentAgentName()
 		slog.Debug("Forwarding RAG event", "type", ragEvent.Type, "rag", ragName, "agent", agentName)
@@ -57,7 +36,5 @@ func ragEventForwarder(ctx context.Context, ragName string, r *LocalRuntime, sen
 		default:
 			slog.Debug("Unhandled RAG event type", "type", ragEvent.Type, "rag", ragName)
 		}
-
-		_ = ctx // available for future use
 	}
 }

pkg/tools/builtin/rag.go

@@ -17,22 +17,21 @@ import (
 // RAGEventCallback is called to forward RAG manager events during initialization.
 type RAGEventCallback func(event ragtypes.Event)
 
-// RAGTool provides document querying capabilities for a single RAG source
+// RAGTool provides document querying capabilities for a single RAG source.
 type RAGTool struct {
 	manager       *rag.Manager
 	toolName      string
 	eventCallback RAGEventCallback
 }
 
-// Verify interface compliance
+// Verify interface compliance.
 var (
 	_ tools.ToolSet      = (*RAGTool)(nil)
 	_ tools.Instructable = (*RAGTool)(nil)
 	_ tools.Startable    = (*RAGTool)(nil)
 )
 
-// NewRAGTool creates a new RAG tool for a single RAG manager
-// toolName is the name to use for the tool (typically from config or manager name)
+// NewRAGTool creates a new RAG tool for a single RAG manager.
 func NewRAGTool(manager *rag.Manager, toolName string) *RAGTool {
 	return &RAGTool{
 		manager:  manager,
@@ -45,39 +44,28 @@ func (t *RAGTool) Name() string {
 	return t.toolName
 }
 
-type QueryRAGArgs struct {
-	Query string `json:"query" jsonschema:"Search query"`
-}
-
-type QueryResult struct {
-	SourcePath string  `json:"source_path" jsonschema:"Path to the source document"`
-	Content    string  `json:"content" jsonschema:"Relevant document chunk content"`
-	Similarity float64 `json:"similarity" jsonschema:"Similarity score (0-1)"`
-	ChunkIndex int     `json:"chunk_index" jsonschema:"Index of the chunk within the source document"`
-}
-
-// SetEventCallback sets a callback to receive RAG manager events during initialization.
-// This must be called before Start() to receive indexing progress events.
+// SetEventCallback sets a callback to receive RAG manager events during
+// initialization. Must be called before Start().
 func (t *RAGTool) SetEventCallback(cb RAGEventCallback) {
 	t.eventCallback = cb
 }
 
-// Start initializes the RAG manager (indexes documents).
+// Start initializes the RAG manager (indexes documents) and starts a
+// file watcher for incremental updates.
 func (t *RAGTool) Start(ctx context.Context) error {
 	if t.manager == nil {
 		return nil
 	}
-	slog.Debug("Starting RAG tool initialization", "tool", t.toolName)
 
-	// Forward RAG manager events if a callback is set
+	// Forward RAG manager events if a callback is set.
 	if t.eventCallback != nil {
 		go t.forwardEvents(ctx)
 	}
 
 	if err := t.manager.Initialize(ctx); err != nil {
 		return fmt.Errorf("failed to initialize RAG manager %q: %w", t.toolName, err)
 	}
-	// Start file watcher in background
+
 	go func() {
 		if err := t.manager.StartFileWatcher(ctx); err != nil {
 			slog.Error("Failed to start RAG file watcher", "tool", t.toolName, "error", err)
@@ -86,6 +74,14 @@ func (t *RAGTool) Start(ctx context.Context) error {
 	return nil
 }
 
+// Stop closes the RAG manager and releases resources.
+func (t *RAGTool) Stop(_ context.Context) error {
+	if t.manager == nil {
+		return nil
+	}
+	return t.manager.Close()
+}
+
 // forwardEvents reads events from the RAG manager and forwards them via the callback.
 func (t *RAGTool) forwardEvents(ctx context.Context) {
 	for {
@@ -101,27 +97,27 @@ func (t *RAGTool) forwardEvents(ctx context.Context) {
 	}
 }
 
-// Stop closes the RAG manager and releases resources.
-func (t *RAGTool) Stop(_ context.Context) error {
-	if t.manager == nil {
-		return nil
-	}
-	return t.manager.Close()
-}
-
 func (t *RAGTool) Instructions() string {
 	if t.manager != nil {
-		instruction := t.manager.ToolInstruction()
-		if instruction != "" {
+		if instruction := t.manager.ToolInstruction(); instruction != "" {
 			return instruction
 		}
 	}
-
-	// Default instruction if none provided
 	return fmt.Sprintf("Search documents in %s to find relevant code or documentation. "+
 		"Provide a clear search query describing what you need.", t.toolName)
 }
 
+type queryRAGArgs struct {
+	Query string `json:"query" jsonschema:"Search query"`
+}
+
+type queryResult struct {
+	SourcePath string  `json:"source_path" jsonschema:"Path to the source document"`
+	Content    string  `json:"content" jsonschema:"Relevant document chunk content"`
+	Similarity float64 `json:"similarity" jsonschema:"Similarity score (0-1)"`
+	ChunkIndex int     `json:"chunk_index" jsonschema:"Index of the chunk within the source document"`
+}
+
 func (t *RAGTool) Tools(context.Context) ([]tools.Tool, error) {
 	var description string
 	if t.manager != nil {
@@ -131,83 +127,53 @@ func (t *RAGTool) Tools(context.Context) ([]tools.Tool, error) {
 		"Provide a natural language query describing what you need. "+
 		"Returns the most relevant document chunks with file paths.", t.toolName))
 
-	paramsSchema := tools.MustSchemaFor[QueryRAGArgs]()
-	outputSchema := tools.MustSchemaFor[[]QueryResult]()
-
-	// Log schemas for debugging
-	if paramsJSON, err := json.Marshal(paramsSchema); err == nil {
-		slog.Debug("RAG tool parameters schema",
-			"tool_name", t.toolName,
-			"schema", string(paramsJSON))
-	}
-	if outputJSON, err := json.Marshal(outputSchema); err == nil {
-		slog.Debug("RAG tool output schema",
-			"tool_name", t.toolName,
-			"schema", string(outputJSON))
-	}
-
-	tool := tools.Tool{
+	return []tools.Tool{{
 		Name:         t.toolName,
 		Category:     "knowledge",
 		Description:  description,
-		Parameters:   paramsSchema,
-		OutputSchema: outputSchema,
+		Parameters:   tools.MustSchemaFor[queryRAGArgs](),
+		OutputSchema: tools.MustSchemaFor[[]queryResult](),
 		Handler:      tools.NewHandler(t.handleQueryRAG),
 		Annotations: tools.ToolAnnotations{
 			ReadOnlyHint: true,
 			Title:        "Query " + t.toolName,
 		},
-	}
-
-	slog.Debug("RAG tool registered",
-		"tool_name", tool.Name,
-		"category", tool.Category,
-		"description", description,
-		"title", tool.Annotations.Title,
-		"read_only", tool.Annotations.ReadOnlyHint)
-
-	return []tools.Tool{tool}, nil
+	}}, nil
 }
 
-func (t *RAGTool) handleQueryRAG(ctx context.Context, args QueryRAGArgs) (*tools.ToolCallResult, error) {
+func (t *RAGTool) handleQueryRAG(ctx context.Context, args queryRAGArgs) (*tools.ToolCallResult, error) {
 	if args.Query == "" {
 		return nil, errors.New("query cannot be empty")
 	}
 
 	results, err := t.manager.Query(ctx, args.Query)
 	if err != nil {
-		slog.Error("RAG query failed", "rag", t.manager.Name(), "error", err)
 		return nil, fmt.Errorf("RAG query failed: %w", err)
 	}
 
-	allResults := make([]QueryResult, 0, len(results))
-	for _, result := range results {
-		allResults = append(allResults, QueryResult{
-			SourcePath: result.Document.SourcePath,
-			Content:    result.Document.Content,
-			Similarity: result.Similarity,
-			ChunkIndex: result.Document.ChunkIndex,
+	out := make([]queryResult, 0, len(results))
+	for _, r := range results {
+		out = append(out, queryResult{
+			SourcePath: r.Document.SourcePath,
+			Content:    r.Document.Content,
+			Similarity: r.Similarity,
+			ChunkIndex: r.Document.ChunkIndex,
 		})
 	}
 
-	sortResults(allResults)
+	slices.SortFunc(out, func(a, b queryResult) int {
+		return cmp.Compare(b.Similarity, a.Similarity)
+	})
 
-	maxResults := 10
-	if len(allResults) > maxResults {
-		allResults = allResults[:maxResults]
+	const maxResults = 10
+	if len(out) > maxResults {
+		out = out[:maxResults]
 	}
 
-	resultJSON, err := json.Marshal(allResults)
+	resultJSON, err := json.Marshal(out)
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal results: %w", err)
 	}
 
 	return tools.ResultSuccess(string(resultJSON)), nil
 }
-
-// sortResults sorts query results by similarity in descending order
-func sortResults(results []QueryResult) {
-	slices.SortFunc(results, func(a, b QueryResult) int {
-		return cmp.Compare(b.Similarity, a.Similarity) // Descending order
-	})
-}