Restore RAG indexing event forwarding to TUI after toolset refactor

The PR #2210 moved RAG from agent-level config to standard toolset type
(tools.Startable) but removed the event forwarding that showed indexing
progress in the TUI sidebar.

This restores event forwarding by:
- Adding an EventCallback to RAGTool that forwards rag.Manager events
  during Start() initialization
- Having StartBackgroundRAGInit discover RAG tools from agent toolsets
  and wire up the event callback before initialization happens
- Converting RAG manager events (indexing started/progress/completed,
  usage, errors) back to runtime events for the TUI

Assisted-By: docker-agent

Author: dgageot

pkg/runtime/rag.go

@@ -1,15 +1,63 @@
 package runtime
 
-import "context"
+import (
+	"context"
+	"fmt"
+	"log/slog"
 
-// StartBackgroundRAGInit is a no-op. RAG initialization is now handled
-// per-toolset via the tools.Startable interface.
-func (r *LocalRuntime) StartBackgroundRAGInit(_ context.Context, _ func(Event)) {
-	// RAG toolsets are initialized lazily when first used.
+	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))
+		}
+	}
 }
 
-// InitializeRAG is a no-op. RAG initialization is now handled
-// per-toolset via the tools.Startable interface.
-func (r *LocalRuntime) InitializeRAG(_ context.Context, _ chan Event) {
-	// RAG toolsets are initialized lazily when first used.
+// 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 {
+	return func(ragEvent ragtypes.Event) {
+		agentName := r.CurrentAgentName()
+		slog.Debug("Forwarding RAG event", "type", ragEvent.Type, "rag", ragName, "agent", agentName)
+
+		switch ragEvent.Type {
+		case ragtypes.EventTypeIndexingStarted:
+			sendEvent(RAGIndexingStarted(ragName, ragEvent.StrategyName))
+		case ragtypes.EventTypeIndexingProgress:
+			if ragEvent.Progress != nil {
+				sendEvent(RAGIndexingProgress(ragName, ragEvent.StrategyName, ragEvent.Progress.Current, ragEvent.Progress.Total, agentName))
+			}
+		case ragtypes.EventTypeIndexingComplete:
+			sendEvent(RAGIndexingCompleted(ragName, ragEvent.StrategyName))
+		case ragtypes.EventTypeUsage:
+			sendEvent(NewTokenUsageEvent("", agentName, &Usage{
+				InputTokens:   ragEvent.TotalTokens,
+				ContextLength: ragEvent.TotalTokens,
+				Cost:          ragEvent.Cost,
+			}))
+		case ragtypes.EventTypeError:
+			if ragEvent.Error != nil {
+				sendEvent(Error(fmt.Sprintf("RAG %s error: %v", ragName, ragEvent.Error)))
+			}
+		default:
+			slog.Debug("Unhandled RAG event type", "type", ragEvent.Type, "rag", ragName)
+		}
+
+		_ = ctx // available for future use
+	}
 }

pkg/tools/builtin/rag.go

@@ -10,13 +10,18 @@ import (
 	"slices"
 
 	"github.com/docker/docker-agent/pkg/rag"
+	ragtypes "github.com/docker/docker-agent/pkg/rag/types"
 	"github.com/docker/docker-agent/pkg/tools"
 )
 
+// 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
 type RAGTool struct {
-	manager  *rag.Manager
-	toolName string
+	manager       *rag.Manager
+	toolName      string
+	eventCallback RAGEventCallback
 }
 
 // Verify interface compliance
@@ -35,6 +40,11 @@ func NewRAGTool(manager *rag.Manager, toolName string) *RAGTool {
 	}
 }
 
+// Name returns the tool name for this RAG source.
+func (t *RAGTool) Name() string {
+	return t.toolName
+}
+
 type QueryRAGArgs struct {
 	Query string `json:"query" jsonschema:"Search query"`
 }
@@ -46,12 +56,24 @@ type QueryResult struct {
 	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.
+func (t *RAGTool) SetEventCallback(cb RAGEventCallback) {
+	t.eventCallback = cb
+}
+
 // Start initializes the RAG manager (indexes documents).
 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
+	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)
 	}
@@ -64,6 +86,21 @@ func (t *RAGTool) Start(ctx context.Context) error {
 	return nil
 }
 
+// forwardEvents reads events from the RAG manager and forwards them via the callback.
+func (t *RAGTool) forwardEvents(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case event, ok := <-t.manager.Events():
+			if !ok {
+				return
+			}
+			t.eventCallback(event)
+		}
+	}
+}
+
 // Stop closes the RAG manager and releases resources.
 func (t *RAGTool) Stop(_ context.Context) error {
 	if t.manager == nil {