Add LSP multiplexer to support multiple LSP toolsets

pkg/teamloader/teamloader.go

@@ -423,8 +423,9 @@ func getFallbackModelsForAgent(ctx context.Context, cfg *latest.Config, a *lates
 // getToolsForAgent returns the tool definitions for an agent based on its configuration
 func getToolsForAgent(ctx context.Context, a *latest.AgentConfig, parentDir string, runConfig *config.RuntimeConfig, registry *ToolsetRegistry) ([]tools.ToolSet, []string) {
 	var (
-		toolSets []tools.ToolSet
-		warnings []string
+		toolSets    []tools.ToolSet
+		warnings    []string
+		lspBackends []builtin.LSPBackend
 	)
 
 	deferredToolset := builtin.NewDeferredToolset()
@@ -456,9 +457,29 @@ func getToolsForAgent(ctx context.Context, a *latest.AgentConfig, parentDir stri
 			}
 		}
 
+		// Collect LSP backends for multiplexing when there are multiple.
+		// Instead of adding them individually (which causes duplicate tool names),
+		// they are combined into a single LSPMultiplexer after the loop.
+		if toolset.Type == "lsp" {
+			if lspTool, ok := tool.(*builtin.LSPTool); ok {
+				lspBackends = append(lspBackends, builtin.LSPBackend{LSP: lspTool, Toolset: wrapped})
+				continue
+			}
+			slog.Warn("Toolset configured as type 'lsp' but registry returned unexpected type; treating as regular toolset",
+				"type", fmt.Sprintf("%T", tool), "command", toolset.Command)
+		}
+
 		toolSets = append(toolSets, wrapped)
 	}
 
+	// Merge LSP backends: if there are multiple, combine them into a single
+	// multiplexer so the LLM sees one set of lsp_* tools instead of duplicates.
+	if len(lspBackends) > 1 {
+		toolSets = append(toolSets, builtin.NewLSPMultiplexer(lspBackends))
+	} else if len(lspBackends) == 1 {
+		toolSets = append(toolSets, lspBackends[0].Toolset)
+	}
+
 	if deferredToolset.HasSources() {
 		toolSets = append(toolSets, deferredToolset)
 	}

pkg/teamloader/teamloader_test.go

@@ -386,6 +386,85 @@ agents:
 	assert.Equal(t, expected, rootAgent.AddPromptFiles())
 }
 
+func TestGetToolsForAgent_MultipleLSPToolsetsAreCombined(t *testing.T) {
+	t.Parallel()
+
+	a := &latest.AgentConfig{
+		Instruction: "test",
+		Toolsets: []latest.Toolset{
+			{
+				Type:      "lsp",
+				Command:   "gopls",
+				FileTypes: []string{".go"},
+			},
+			{
+				Type:      "lsp",
+				Command:   "gopls",
+				FileTypes: []string{".mod"},
+			},
+		},
+	}
+
+	runConfig := config.RuntimeConfig{
+		EnvProviderForTests: &noEnvProvider{},
+	}
+
+	got, warnings := getToolsForAgent(t.Context(), a, ".", &runConfig, NewDefaultToolsetRegistry())
+	require.Empty(t, warnings)
+
+	// Should have exactly one toolset (the multiplexer)
+	require.Len(t, got, 1)
+
+	// Verify that we get no duplicate tool names
+	allTools, err := got[0].Tools(t.Context())
+	require.NoError(t, err)
+
+	seen := make(map[string]bool)
+	for _, tool := range allTools {
+		assert.False(t, seen[tool.Name], "duplicate tool name: %s", tool.Name)
+		seen[tool.Name] = true
+	}
+
+	// Verify LSP tools are present
+	assert.True(t, seen["lsp_hover"])
+	assert.True(t, seen["lsp_definition"])
+}
+
+func TestGetToolsForAgent_SingleLSPToolsetNotWrapped(t *testing.T) {
+	t.Parallel()
+
+	a := &latest.AgentConfig{
+		Instruction: "test",
+		Toolsets: []latest.Toolset{
+			{
+				Type:      "lsp",
+				Command:   "gopls",
+				FileTypes: []string{".go"},
+			},
+		},
+	}
+
+	runConfig := config.RuntimeConfig{
+		EnvProviderForTests: &noEnvProvider{},
+	}
+
+	got, warnings := getToolsForAgent(t.Context(), a, ".", &runConfig, NewDefaultToolsetRegistry())
+	require.Empty(t, warnings)
+
+	// Should have exactly one toolset that provides LSP tools.
+	require.Len(t, got, 1)
+
+	allTools, err := got[0].Tools(t.Context())
+	require.NoError(t, err)
+
+	var names []string
+	for _, tool := range allTools {
+		names = append(names, tool.Name)
+	}
+	assert.Contains(t, names, "lsp_hover")
+	assert.Contains(t, names, "lsp_definition")
+}
+
 func TestExternalDepthContext(t *testing.T) {
 	t.Parallel()
 

pkg/tools/builtin/lsp_multiplexer.go

@@ -0,0 +1,172 @@
+package builtin
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/docker/cagent/pkg/tools"
+)
+
+// LSPMultiplexer combines multiple LSP backends into a single toolset.
+// It presents one set of lsp_* tools and routes each call to the appropriate
+// backend based on the file extension in the tool arguments.
+type LSPMultiplexer struct {
+	backends []LSPBackend
+}
+
+// LSPBackend pairs a raw LSPTool (used for file-type routing) with an
+// optionally-wrapped ToolSet (used for tool enumeration, so that per-toolset
+// config like tool filters, instructions, or toon wrappers are respected).
+type LSPBackend struct {
+	LSP     *LSPTool
+	Toolset tools.ToolSet
+}
+
+// lspRouteTarget pairs a backend with the tool handler it produced for a given tool name.
+type lspRouteTarget struct {
+	lsp     *LSPTool
+	handler tools.ToolHandler
+}
+
+// Verify interface compliance.
+var (
+	_ tools.ToolSet      = (*LSPMultiplexer)(nil)
+	_ tools.Startable    = (*LSPMultiplexer)(nil)
+	_ tools.Instructable = (*LSPMultiplexer)(nil)
+)
+
+// NewLSPMultiplexer creates a multiplexer that routes LSP tool calls
+// to the appropriate backend based on file type.
+func NewLSPMultiplexer(backends []LSPBackend) *LSPMultiplexer {
+	return &LSPMultiplexer{backends: append([]LSPBackend{}, backends...)}
+}
+
+func (m *LSPMultiplexer) Start(ctx context.Context) error {
+	var started int
+	for _, b := range m.backends {
+		if err := b.LSP.Start(ctx); err != nil {
+			// Clean up previously started backends to avoid resource leaks.
+			for _, s := range m.backends[:started] {
+				_ = s.LSP.Stop(ctx)
+			}
+			return fmt.Errorf("starting LSP backend %q: %w", b.LSP.handler.command, err)
+		}
+		started++
+	}
+	return nil
+}
+
+func (m *LSPMultiplexer) Stop(ctx context.Context) error {
+	var errs []error
+	for _, b := range m.backends {
+		if err := b.LSP.Stop(ctx); err != nil {
+			errs = append(errs, fmt.Errorf("stopping LSP backend %q: %w", b.LSP.handler.command, err))
+		}
+	}
+	return errors.Join(errs...)
+}
+
+func (m *LSPMultiplexer) Instructions() string {
+	// Combine instructions from all backends, deduplicating identical ones.
+	// Typically they share the same base LSP instructions, but individual
+	// toolsets may override them via the Instruction config field.
+	var parts []string
+	seen := make(map[string]bool)
+	for _, b := range m.backends {
+		instr := tools.GetInstructions(b.Toolset)
+		if instr != "" && !seen[instr] {
+			seen[instr] = true
+			parts = append(parts, instr)
+		}
+	}
+	return strings.Join(parts, "\n\n")
+}
+
+func (m *LSPMultiplexer) Tools(ctx context.Context) ([]tools.Tool, error) {
+	// Collect each backend's tools keyed by name. We build the union of all
+	// tool names (not just the first backend's) so that per-backend tool
+	// filters don't accidentally hide tools that other backends expose.
+	handlersByName := make(map[string][]lspRouteTarget)
+	seenTools := make(map[string]tools.Tool) // first definition wins (for schema/description)
+	var toolOrder []string                   // preserve insertion order
+	for _, b := range m.backends {
+		bTools, err := b.Toolset.Tools(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("getting tools from LSP backend %q: %w", b.LSP.handler.command, err)
+		}
+		for _, t := range bTools {
+			handlersByName[t.Name] = append(handlersByName[t.Name], lspRouteTarget{b.LSP, t.Handler})
+			if _, exists := seenTools[t.Name]; !exists {
+				seenTools[t.Name] = t
+				toolOrder = append(toolOrder, t.Name)
+			}
+		}
+	}
+
+	result := make([]tools.Tool, 0, len(toolOrder))
+	for _, name := range toolOrder {
+		t := seenTools[name]
+		handlers := handlersByName[name]
+		if name == ToolNameLSPWorkspace || name == ToolNameLSPWorkspaceSymbols {
+			t.Handler = broadcastLSP(handlers)
+		} else {
+			t.Handler = routeByFile(handlers)
+		}
+		result = append(result, t)
+	}
+	return result, nil
+}
+
+// routeByFile returns a handler that extracts the "file" field from the JSON
+// arguments and dispatches to the backend whose file-type filter matches.
+func routeByFile(handlers []lspRouteTarget) tools.ToolHandler {
+	return func(ctx context.Context, tc tools.ToolCall) (*tools.ToolCallResult, error) {
+		var args struct {
+			File string `json:"file"`
+		}
+		if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err != nil {
+			return tools.ResultError(fmt.Sprintf("failed to parse file argument: %s", err)), nil
+		}
+		if args.File == "" {
+			return tools.ResultError("file argument is required"), nil
+		}
+		for _, h := range handlers {
+			if h.lsp.HandlesFile(args.File) {
+				return h.handler(ctx, tc)
+			}
+		}
+		return tools.ResultError(fmt.Sprintf("no LSP server configured for file: %s", args.File)), nil
+	}
+}
+
+// broadcastLSP returns a handler that calls every backend best-effort and
+// merges the outputs. Individual backend failures are collected rather than
+// aborting the entire operation.
+func broadcastLSP(handlers []lspRouteTarget) tools.ToolHandler {
+	return func(ctx context.Context, tc tools.ToolCall) (*tools.ToolCallResult, error) {
+		var sections []string
+		var errs []error
+		for _, h := range handlers {
+			result, err := h.handler(ctx, tc)
+			if err != nil {
+				errs = append(errs, fmt.Errorf("backend %s: %w", h.lsp.handler.command, err))
+				continue
+			}
+			if result.IsError {
+				sections = append(sections, fmt.Sprintf("[LSP %s] Error: %s", h.lsp.handler.command, result.Output))
+			} else if result.Output != "" {
+				sections = append(sections, result.Output)
+			}
+		}
+		if len(sections) == 0 && len(errs) > 0 {
+			return nil, errors.Join(errs...)
+		}
+		if len(sections) == 0 {
+			return tools.ResultSuccess("No results"), nil
+		}
+		return tools.ResultSuccess(strings.Join(sections, "\n---\n")), nil
+	}
+}