refactor: remove DescriptionToolSet wrapper

Move description parameter handling from a wrapper struct to a simple
function called at tool retrieval time. This eliminates the need for
DescriptionToolSet to implement all capability interfaces (Startable,
Instructable, Elicitable, OAuthCapable) just to forward calls.

Changes:
- Add addDescriptionParameter field to Agent struct
- Add WithAddDescriptionParameter option
- Add tools.AddDescriptionParameter function to transform tool lists
- Remove DescriptionToolSet wrapper entirely
- Apply transformation in agent.Tools() when flag is enabled

Assisted-By: cagent

Author: dgageot

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           []*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
-	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/opts.go

@@ -83,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

pkg/teamloader/teamloader.go

@@ -149,6 +149,7 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c
 			agent.WithWelcomeMessage(expander.Expand(ctx, agentConfig.WelcomeMessage)),
 			agent.WithAddDate(agentConfig.AddDate),
 			agent.WithAddEnvironmentInfo(agentConfig.AddEnvironmentInfo),
+			agent.WithAddDescriptionParameter(agentConfig.AddDescriptionParameter),
 			agent.WithAddPromptFiles(agentConfig.AddPromptFiles),
 			agent.WithMaxIterations(agentConfig.MaxIterations),
 			agent.WithNumHistoryItems(agentConfig.NumHistoryItems),
@@ -344,11 +345,6 @@ func getToolsForAgent(ctx context.Context, a *latest.AgentConfig, parentDir stri
 			}
 		}
 
-		// Apply friendly tools wrapper if enabled
-		if a.AddDescriptionParameter {
-			wrapped = tools.NewDescriptionToolSet(wrapped)
-		}
-
 		toolSets = append(toolSets, wrapped)
 	}
 

pkg/tools/description.go

@@ -1,7 +1,6 @@
 package tools
 
 import (
-	"context"
 	"encoding/json"
 )
 
@@ -10,76 +9,18 @@ const (
 	DescriptionParam = "description"
 )
 
-// DescriptionToolSet wraps a ToolSet and adds a "description" parameter to all tools.
-// This allows the LLM to provide context about what it's doing with each tool call.
-type DescriptionToolSet struct {
-	inner ToolSet
-}
-
-// Verify interface compliance
-var (
-	_ ToolSet      = (*DescriptionToolSet)(nil)
-	_ Startable    = (*DescriptionToolSet)(nil)
-	_ Instructable = (*DescriptionToolSet)(nil)
-	_ Elicitable   = (*DescriptionToolSet)(nil)
-	_ OAuthCapable = (*DescriptionToolSet)(nil)
-)
-
-// NewDescriptionToolSet creates a new DescriptionToolSet wrapping the given ToolSet.
-func NewDescriptionToolSet(inner ToolSet) *DescriptionToolSet {
-	return &DescriptionToolSet{inner: inner}
-}
-
-func (f *DescriptionToolSet) Tools(ctx context.Context) ([]Tool, error) {
-	tools, err := f.inner.Tools(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	result := make([]Tool, len(tools))
-	for i, tool := range tools {
-		result[i] = f.addDescriptionParam(tool)
-	}
-	return result, nil
-}
-
-func (f *DescriptionToolSet) Instructions() string {
-	return GetInstructions(f.inner)
-}
-
-func (f *DescriptionToolSet) Start(ctx context.Context) error {
-	if s, ok := As[Startable](f.inner); ok {
-		return s.Start(ctx)
-	}
-	return nil
-}
-
-func (f *DescriptionToolSet) Stop(ctx context.Context) error {
-	if s, ok := As[Startable](f.inner); ok {
-		return s.Stop(ctx)
-	}
-	return nil
-}
-
-func (f *DescriptionToolSet) SetElicitationHandler(handler ElicitationHandler) {
-	if e, ok := As[Elicitable](f.inner); ok {
-		e.SetElicitationHandler(handler)
-	}
-}
-
-func (f *DescriptionToolSet) SetOAuthSuccessHandler(handler func()) {
-	if o, ok := As[OAuthCapable](f.inner); ok {
-		o.SetOAuthSuccessHandler(handler)
-	}
-}
-
-func (f *DescriptionToolSet) SetManagedOAuth(managed bool) {
-	if o, ok := As[OAuthCapable](f.inner); ok {
-		o.SetManagedOAuth(managed)
+// AddDescriptionParameter adds a "description" parameter to tools that have
+// AddDescriptionParameter set to true. This allows the LLM to provide context
+// about what it's doing with each tool call.
+func AddDescriptionParameter(toolList []Tool) []Tool {
+	result := make([]Tool, len(toolList))
+	for i, tool := range toolList {
+		result[i] = addDescriptionParam(tool)
 	}
+	return result
 }
 
-func (f *DescriptionToolSet) addDescriptionParam(tool Tool) Tool {
+func addDescriptionParam(tool Tool) Tool {
 	if !tool.AddDescriptionParameter {
 		return tool
 	}

pkg/tools/description_test.go

@@ -1,17 +1,16 @@
 package tools
 
 import (
-	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-type mockToolSet struct{}
+func TestAddDescriptionParameter_AddsDescriptionParameter(t *testing.T) {
+	t.Parallel()
 
-func (m *mockToolSet) Tools(_ context.Context) ([]Tool, error) {
-	return []Tool{
+	tools := []Tool{
 		{
 			Name:        "test_tool",
 			Description: "A test tool",
@@ -25,23 +24,14 @@ func (m *mockToolSet) Tools(_ context.Context) ([]Tool, error) {
 				},
 				"required": []string{"path"},
 			},
-			Handler: func(_ context.Context, _ ToolCall) (*ToolCallResult, error) {
-				return &ToolCallResult{Output: "ok"}, nil
-			},
 			AddDescriptionParameter: true,
 		},
-	}, nil
-}
-
-func TestDescriptionToolSet_AddsDescriptionParameter(t *testing.T) {
-	t.Parallel()
-	desc := NewDescriptionToolSet(&mockToolSet{})
+	}
 
-	tools, err := desc.Tools(t.Context())
-	require.NoError(t, err)
-	require.Len(t, tools, 1)
+	result := AddDescriptionParameter(tools)
+	require.Len(t, result, 1)
 
-	tool := tools[0]
+	tool := result[0]
 	schema, ok := tool.Parameters.(map[string]any)
 	require.True(t, ok)
 
@@ -54,14 +44,28 @@ func TestDescriptionToolSet_AddsDescriptionParameter(t *testing.T) {
 	assert.Contains(t, descProp["description"], "human-readable")
 }
 
-func TestDescriptionToolSet_PreservesOriginalParameters(t *testing.T) {
+func TestAddDescriptionParameter_PreservesOriginalParameters(t *testing.T) {
 	t.Parallel()
-	desc := NewDescriptionToolSet(&mockToolSet{})
 
-	tools, err := desc.Tools(t.Context())
-	require.NoError(t, err)
+	tools := []Tool{
+		{
+			Name: "test_tool",
+			Parameters: map[string]any{
+				"type": "object",
+				"properties": map[string]any{
+					"path": map[string]any{
+						"type":        "string",
+						"description": "The file path",
+					},
+				},
+			},
+			AddDescriptionParameter: true,
+		},
+	}
 
-	schema := tools[0].Parameters.(map[string]any)
+	result := AddDescriptionParameter(tools)
+
+	schema := result[0].Parameters.(map[string]any)
 	properties := schema["properties"].(map[string]any)
 
 	pathProp, ok := properties["path"].(map[string]any)
@@ -70,6 +74,30 @@ func TestDescriptionToolSet_PreservesOriginalParameters(t *testing.T) {
 	assert.Equal(t, "The file path", pathProp["description"])
 }
 
+func TestAddDescriptionParameter_SkipsToolsWithoutFlag(t *testing.T) {
+	t.Parallel()
+
+	tools := []Tool{
+		{
+			Name: "test_tool",
+			Parameters: map[string]any{
+				"type":       "object",
+				"properties": map[string]any{},
+			},
+			AddDescriptionParameter: false,
+		},
+	}
+
+	result := AddDescriptionParameter(tools)
+	require.Len(t, result, 1)
+
+	schema := result[0].Parameters.(map[string]any)
+	properties := schema["properties"].(map[string]any)
+
+	_, hasDesc := properties[DescriptionParam]
+	assert.False(t, hasDesc, "description parameter should not be added")
+}
+
 func TestExtractDescription(t *testing.T) {
 	t.Parallel()
 	tests := []struct {