Merge pull request #1548 from dgageot/wrappers

Work on Tool Wrappers

Author: rumpl

pkg/agent/agent.go

@@ -15,27 +15,28 @@ import (
 
 // Agent represents an AI agent
 type Agent struct {
-	name               string
-	description        string
-	welcomeMessage     string
-	instruction        string
-	toolsets           []*StartableToolSet
-	models             []provider.Provider
-	modelOverrides     atomic.Pointer[[]provider.Provider] // Optional model override(s) set at runtime (supports alloy)
-	subAgents          []*Agent
-	handoffs           []*Agent
-	parents            []*Agent
-	addDate            bool
-	addEnvironmentInfo bool
-	maxIterations      int
-	numHistoryItems    int
-	addPromptFiles     []string
-	tools              []tools.Tool
-	commands           types.Commands
-	pendingWarnings    []string
-	skillsEnabled      bool
-	hooks              *latest.HooksConfig
-	thinkingConfigured bool // true if thinking_budget was explicitly set in config
+	name                    string
+	description             string
+	welcomeMessage          string
+	instruction             string
+	toolsets                []*tools.StartableToolSet
+	models                  []provider.Provider
+	modelOverrides          atomic.Pointer[[]provider.Provider] // Optional model override(s) set at runtime (supports alloy)
+	subAgents               []*Agent
+	handoffs                []*Agent
+	parents                 []*Agent
+	addDate                 bool
+	addEnvironmentInfo      bool
+	addDescriptionParameter bool
+	maxIterations           int
+	numHistoryItems         int
+	addPromptFiles          []string
+	tools                   []tools.Tool
+	commands                types.Commands
+	pendingWarnings         []string
+	skillsEnabled           bool
+	hooks                   *latest.HooksConfig
+	thinkingConfigured      bool // true if thinking_budget was explicitly set in config
 }
 
 // New creates a new agent
@@ -203,6 +204,10 @@ func (a *Agent) Tools(ctx context.Context) ([]tools.Tool, error) {
 
 	agentTools = append(agentTools, a.tools...)
 
+	if a.addDescriptionParameter {
+		agentTools = tools.AddDescriptionParameter(agentTools)
+	}
+
 	return agentTools, nil
 }
 

pkg/agent/agent_test.go

@@ -14,12 +14,17 @@ import (
 )
 
 type stubToolSet struct {
-	tools.BaseToolSet
 	startErr error
 	tools    []tools.Tool
 	listErr  error
 }
 
+// Verify interface compliance
+var (
+	_ tools.ToolSet   = (*stubToolSet)(nil)
+	_ tools.Startable = (*stubToolSet)(nil)
+)
+
 func newStubToolSet(startErr error, toolsList []tools.Tool, listErr error) tools.ToolSet {
 	return &stubToolSet{
 		startErr: startErr,
@@ -29,6 +34,7 @@ func newStubToolSet(startErr error, toolsList []tools.Tool, listErr error) tools
 }
 
 func (s *stubToolSet) Start(context.Context) error { return s.startErr }
+func (s *stubToolSet) Stop(context.Context) error  { return nil }
 func (s *stubToolSet) Tools(context.Context) ([]tools.Tool, error) {
 	if s.listErr != nil {
 		return nil, s.listErr

pkg/agent/opts.go

@@ -1,9 +1,6 @@
 package agent
 
 import (
-	"context"
-	"sync"
-
 	"github.com/docker/cagent/pkg/config/latest"
 	"github.com/docker/cagent/pkg/config/types"
 	"github.com/docker/cagent/pkg/model/provider"
@@ -19,11 +16,9 @@ func WithInstruction(instruction string) Opt {
 }
 
 func WithToolSets(toolSet ...tools.ToolSet) Opt {
-	var startableToolSet []*StartableToolSet
+	var startableToolSet []*tools.StartableToolSet
 	for _, ts := range toolSet {
-		startableToolSet = append(startableToolSet, &StartableToolSet{
-			ToolSet: ts,
-		})
+		startableToolSet = append(startableToolSet, tools.NewStartable(ts))
 	}
 
 	return func(a *Agent) {
@@ -88,6 +83,12 @@ func WithAddEnvironmentInfo(addEnvironmentInfo bool) Opt {
 	}
 }
 
+func WithAddDescriptionParameter(addDescriptionParameter bool) Opt {
+	return func(a *Agent) {
+		a.addDescriptionParameter = addDescriptionParameter
+	}
+}
+
 func WithAddPromptFiles(addPromptFiles []string) Opt {
 	return func(a *Agent) {
 		a.addPromptFiles = addPromptFiles
@@ -139,36 +140,3 @@ func WithThinkingConfigured(configured bool) Opt {
 		a.thinkingConfigured = configured
 	}
 }
-
-type StartableToolSet struct {
-	tools.ToolSet
-
-	mu      sync.Mutex
-	started bool
-}
-
-// IsStarted returns whether the toolset has been successfully started.
-func (s *StartableToolSet) IsStarted() bool {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	return s.started
-}
-
-// Start starts the toolset.
-//
-// It provides single-flight semantics: concurrent callers block until this start
-// attempt completes. If this attempt fails, a future call will retry.
-func (s *StartableToolSet) Start(ctx context.Context) error {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	if s.started {
-		return nil
-	}
-
-	err := s.ToolSet.Start(ctx)
-	if err == nil {
-		s.started = true
-	}
-	return err
-}

pkg/agent/opts_test.go

@@ -12,23 +12,30 @@ import (
 )
 
 type flakyStartToolset struct {
-	tools.BaseToolSet
 	calls atomic.Int64
 }
 
+// Verify interface compliance
+var (
+	_ tools.ToolSet   = (*flakyStartToolset)(nil)
+	_ tools.Startable = (*flakyStartToolset)(nil)
+)
+
 func (f *flakyStartToolset) Start(context.Context) error {
 	if f.calls.Add(1) == 1 {
 		return errors.New("no events channel available for elicitation")
 	}
 	return nil
 }
 
+func (f *flakyStartToolset) Stop(context.Context) error { return nil }
+
 func (f *flakyStartToolset) Tools(context.Context) ([]tools.Tool, error) { return nil, nil }
 
 func TestStartableToolSet_RetriesAfterFailure(t *testing.T) {
 	ctx := t.Context()
 	inner := &flakyStartToolset{}
-	ts := &StartableToolSet{ToolSet: inner}
+	ts := tools.NewStartable(inner)
 
 	err := ts.Start(ctx)
 	require.Error(t, err)

pkg/app/app.go

@@ -188,7 +188,7 @@ func (a *App) ExecuteMCPPrompt(ctx context.Context, promptName string, arguments
 	}
 
 	for _, toolset := range currentAgent.ToolSets() {
-		if mcpToolset := runtime.UnwrapMCPToolset(toolset); mcpToolset != nil {
+		if mcpToolset, ok := tools.As[*mcptools.Toolset](toolset); ok {
 			result, err := mcpToolset.GetPrompt(ctx, promptName, arguments)
 			if err == nil {
 				// Convert the MCP result to a string format suitable for the editor

pkg/runtime/runtime.go

@@ -38,23 +38,6 @@ import (
 	mcptools "github.com/docker/cagent/pkg/tools/mcp"
 )
 
-// UnwrapMCPToolset extracts an MCP toolset from a potentially wrapped StartableToolSet.
-// Returns the MCP toolset if found, or nil if the toolset is not an MCP toolset.
-func UnwrapMCPToolset(toolset tools.ToolSet) *mcptools.Toolset {
-	var innerToolset tools.ToolSet
-	if startableTS, ok := toolset.(*agent.StartableToolSet); ok {
-		innerToolset = startableTS.ToolSet
-	} else {
-		innerToolset = toolset
-	}
-
-	if mcpToolset, ok := innerToolset.(*mcptools.Toolset); ok {
-		return mcpToolset
-	}
-
-	return nil
-}
-
 type ResumeType string
 
 // ElicitationResult represents the result of an elicitation request
@@ -459,7 +442,7 @@ func (r *LocalRuntime) CurrentMCPPrompts(ctx context.Context) map[string]mcptool
 
 	// Iterate through all toolsets of the current agent
 	for _, toolset := range currentAgent.ToolSets() {
-		if mcpToolset := UnwrapMCPToolset(toolset); mcpToolset != nil {
+		if mcpToolset, ok := tools.As[*mcptools.Toolset](toolset); ok {
 			slog.Debug("Found MCP toolset", "toolset", mcpToolset)
 			// Discover prompts from this MCP toolset
 			mcpPrompts := r.discoverMCPPrompts(ctx, mcpToolset)
@@ -639,7 +622,7 @@ func (r *LocalRuntime) emitToolsProgressively(ctx context.Context, a *agent.Agen
 		isLast := i == totalToolsets-1
 
 		// Start the toolset if needed
-		if startable, ok := toolset.(*agent.StartableToolSet); ok {
+		if startable, ok := toolset.(*tools.StartableToolSet); ok {
 			if !startable.IsStarted() {
 				if err := startable.Start(ctx); err != nil {
 					slog.Warn("Toolset start failed; skipping", "agent", a.Name(), "toolset", fmt.Sprintf("%T", startable.ToolSet), "error", err)
@@ -1013,11 +996,11 @@ func (r *LocalRuntime) getTools(ctx context.Context, a *agent.Agent, sessionSpan
 // configureToolsetHandlers sets up elicitation and OAuth handlers for all toolsets of an agent.
 func (r *LocalRuntime) configureToolsetHandlers(a *agent.Agent, events chan Event) {
 	for _, toolset := range a.ToolSets() {
-		toolset.SetElicitationHandler(r.elicitationHandler)
-		toolset.SetOAuthSuccessHandler(func() {
-			events <- Authorization(tools.ElicitationActionAccept, r.currentAgent)
-		})
-		toolset.SetManagedOAuth(r.managedOAuth)
+		tools.ConfigureHandlers(toolset,
+			r.elicitationHandler,
+			func() { events <- Authorization(tools.ElicitationActionAccept, r.currentAgent) },
+			r.managedOAuth,
+		)
 	}
 }
 

pkg/runtime/runtime_test.go

@@ -29,12 +29,17 @@ import (
 )
 
 type stubToolSet struct {
-	tools.BaseToolSet
 	startErr error
 	tools    []tools.Tool
 	listErr  error
 }
 
+// Verify interface compliance
+var (
+	_ tools.ToolSet   = (*stubToolSet)(nil)
+	_ tools.Startable = (*stubToolSet)(nil)
+)
+
 func newStubToolSet(startErr error, toolsList []tools.Tool, listErr error) tools.ToolSet {
 	return &stubToolSet{
 		startErr: startErr,
@@ -44,6 +49,7 @@ func newStubToolSet(startErr error, toolsList []tools.Tool, listErr error) tools
 }
 
 func (s *stubToolSet) Start(context.Context) error { return s.startErr }
+func (s *stubToolSet) Stop(context.Context) error  { return nil }
 func (s *stubToolSet) Tools(context.Context) ([]tools.Tool, error) {
 	if s.listErr != nil {
 		return nil, s.listErr

pkg/session/session.go

@@ -11,6 +11,7 @@ import (
 	"github.com/docker/cagent/pkg/agent"
 	"github.com/docker/cagent/pkg/chat"
 	"github.com/docker/cagent/pkg/skills"
+	"github.com/docker/cagent/pkg/tools"
 )
 
 const (
@@ -482,10 +483,10 @@ func buildInvariantSystemMessages(a *agent.Agent) []chat.Message {
 	}
 
 	for _, toolSet := range a.ToolSets() {
-		if toolSet.Instructions() != "" {
+		if instructions := tools.GetInstructions(toolSet); instructions != "" {
 			messages = append(messages, chat.Message{
 				Role:    chat.MessageRoleSystem,
-				Content: toolSet.Instructions(),
+				Content: instructions,
 			})
 		}
 	}

pkg/teamloader/filter.go

@@ -42,6 +42,14 @@ type filterTools struct {
 	exclude   bool
 }
 
+// Verify interface compliance
+var _ tools.Instructable = (*filterTools)(nil)
+
+// Instructions implements tools.Instructable by delegating to the inner toolset.
+func (f *filterTools) Instructions() string {
+	return tools.GetInstructions(f.ToolSet)
+}
+
 func (f *filterTools) Tools(ctx context.Context) ([]tools.Tool, error) {
 	allTools, err := f.ToolSet.Tools(ctx)
 	if err != nil {

pkg/teamloader/filter_test.go

@@ -119,3 +119,51 @@ func TestWithToolsFilter_CaseSensitive(t *testing.T) {
 	require.Len(t, result, 1)
 	assert.Equal(t, "tool1", result[0].Name)
 }
+
+type instructableToolSet struct {
+	mockToolSet
+	instructions string
+}
+
+func (i *instructableToolSet) Instructions() string {
+	return i.instructions
+}
+
+func TestWithToolsFilter_InstructablePassthrough(t *testing.T) {
+	// Test that filtering preserves instructions from inner toolset
+	inner := &instructableToolSet{
+		mockToolSet: mockToolSet{
+			toolsFunc: func(context.Context) ([]tools.Tool, error) {
+				return []tools.Tool{{Name: "tool1"}, {Name: "tool2"}}, nil
+			},
+		},
+		instructions: "Test instructions for the toolset",
+	}
+
+	wrapped := WithToolsFilter(inner, "tool1")
+
+	// Verify instructions are preserved through the filter wrapper
+	instructions := tools.GetInstructions(wrapped)
+	assert.Equal(t, "Test instructions for the toolset", instructions)
+
+	// Verify filtering still works
+	result, err := wrapped.Tools(t.Context())
+	require.NoError(t, err)
+	require.Len(t, result, 1)
+	assert.Equal(t, "tool1", result[0].Name)
+}
+
+func TestWithToolsFilter_NonInstructableInner(t *testing.T) {
+	// Test that filter works with toolsets that don't implement Instructable
+	inner := &mockToolSet{
+		toolsFunc: func(context.Context) ([]tools.Tool, error) {
+			return []tools.Tool{{Name: "tool1"}}, nil
+		},
+	}
+
+	wrapped := WithToolsFilter(inner, "tool1")
+
+	// Verify instructions are empty for non-instructable inner toolset
+	instructions := tools.GetInstructions(wrapped)
+	assert.Empty(t, instructions)
+}