Merge pull request #559 from dgageot/improve-tools-code

dgageot · web-flow · commit 26330cc0e157 · 2025-10-20T15:36:13.000+02:00
Improve tools code
diff --git a/pkg/creator/agent.go b/pkg/creator/agent.go
@@ -27,6 +27,8 @@ import (
 var agentBuilderInstructions string
 
 type fsToolset struct {
+	tools.ElicitationTool
+
 	inner                    tools.ToolSet
 	originalWriteFileHandler tools.ToolHandler
 	path                     string
@@ -44,14 +46,6 @@ func (f *fsToolset) Stop(ctx context.Context) error {
 	return f.inner.Stop(ctx)
 }
 
-func (f *fsToolset) SetElicitationHandler(tools.ElicitationHandler) {
-	// No-op, this tool does not use elicitation
-}
-
-func (f *fsToolset) SetOAuthSuccessHandler(func()) {
-	// No-op, this tool does not use OAuth
-}
-
 func (f *fsToolset) Tools(ctx context.Context) ([]tools.Tool, error) {
 	innerTools, err := f.inner.Tools(ctx)
 	if err != nil {
diff --git a/pkg/teamloader/teamloader.go b/pkg/teamloader/teamloader.go
@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/docker/cagent/pkg/agent"
-	"github.com/docker/cagent/pkg/codemode"
 	"github.com/docker/cagent/pkg/config"
 	latest "github.com/docker/cagent/pkg/config/v2"
 	"github.com/docker/cagent/pkg/environment"
@@ -24,6 +23,7 @@ import (
 	"github.com/docker/cagent/pkg/team"
 	"github.com/docker/cagent/pkg/tools"
 	"github.com/docker/cagent/pkg/tools/builtin"
+	"github.com/docker/cagent/pkg/tools/codemode"
 	"github.com/docker/cagent/pkg/tools/mcp"
 )
 
diff --git a/pkg/tools/builtin/builtiin.go b/pkg/tools/builtin/builtiin.go
diff --git a/pkg/tools/builtin/fetch.go b/pkg/tools/builtin/fetch.go
@@ -20,7 +20,7 @@ import (
 const userAgent = "cagent/1.0"
 
 type FetchTool struct {
-	elicitationTool
+	tools.ElicitationTool
 	handler *fetchHandler
 }
 
diff --git a/pkg/tools/builtin/filesystem.go b/pkg/tools/builtin/filesystem.go
@@ -24,7 +24,7 @@ type PostEditConfig struct {
 }
 
 type FilesystemTool struct {
-	elicitationTool
+	tools.ElicitationTool
 
 	allowedDirectories []string
 	allowedTools       []string
diff --git a/pkg/tools/builtin/memory.go b/pkg/tools/builtin/memory.go
@@ -12,7 +12,7 @@ import (
 )
 
 type MemoryTool struct {
-	elicitationTool
+	tools.ElicitationTool
 	manager memorymanager.Manager
 }
 
diff --git a/pkg/tools/builtin/script_shell.go b/pkg/tools/builtin/script_shell.go
@@ -14,7 +14,7 @@ import (
 )
 
 type ScriptShellTool struct {
-	elicitationTool
+	tools.ElicitationTool
 	shellTools map[string]latest.ScriptShellToolConfig
 	env        []string
 }
diff --git a/pkg/tools/builtin/shell.go b/pkg/tools/builtin/shell.go
@@ -15,7 +15,7 @@ import (
 )
 
 type ShellTool struct {
-	elicitationTool
+	tools.ElicitationTool
 	handler *shellHandler
 }
 
diff --git a/pkg/tools/builtin/think.go b/pkg/tools/builtin/think.go
@@ -10,7 +10,7 @@ import (
 )
 
 type ThinkTool struct {
-	elicitationTool
+	tools.ElicitationTool
 	handler *thinkHandler
 }
 
diff --git a/pkg/tools/builtin/todo.go b/pkg/tools/builtin/todo.go
@@ -12,7 +12,7 @@ import (
 )
 
 type TodoTool struct {
-	elicitationTool
+	tools.ElicitationTool
 	handler *todoHandler
 }
 
diff --git a/pkg/tools/builtin/transfertask.go b/pkg/tools/builtin/transfertask.go
@@ -7,7 +7,7 @@ import (
 )
 
 type TransferTaskTool struct {
-	elicitationTool
+	tools.ElicitationTool
 }
 
 // Make sure Transfer Tool implements the ToolSet Interface
diff --git a/pkg/tools/codemode/codemode.go b/pkg/tools/codemode/codemode.go
@@ -35,6 +35,7 @@ func Wrap(toolsets ...tools.ToolSet) tools.ToolSet {
 }
 
 type codeModeTool struct {
+	tools.ElicitationTool
 	toolsets []tools.ToolSet
 }
 
@@ -128,11 +129,3 @@ func (c *codeModeTool) Stop(ctx context.Context) error {
 
 	return errors.Join(errs...)
 }
-
-func (c *codeModeTool) SetElicitationHandler(tools.ElicitationHandler) {
-	// No-op, this tool does not use elicitation
-}
-
-func (c *codeModeTool) SetOAuthSuccessHandler(func()) {
-	// No-op, this tool does not use OAuth
-}
diff --git a/pkg/tools/codemode/codemode_test.go b/pkg/tools/codemode/codemode_test.go
@@ -165,6 +165,8 @@ func TestCodeModeTool_CallEcho(t *testing.T) {
 }
 
 type testToolSet struct {
+	tools.ElicitationTool
+
 	tools []tools.Tool
 	start int
 	stop  int
@@ -187,7 +189,3 @@ func (t *testToolSet) Stop(context.Context) error {
 	t.stop++
 	return nil
 }
-
-func (t *testToolSet) SetElicitationHandler(tools.ElicitationHandler) {}
-
-func (t *testToolSet) SetOAuthSuccessHandler(func()) {}
diff --git a/pkg/tools/codemode/console.go b/pkg/tools/codemode/console.go
diff --git a/pkg/tools/codemode/console_test.go b/pkg/tools/codemode/console_test.go
diff --git a/pkg/tools/codemode/exec.go b/pkg/tools/codemode/exec.go
diff --git a/pkg/tools/codemode/exec_test.go b/pkg/tools/codemode/exec_test.go
diff --git a/pkg/tools/codemode/functions.go b/pkg/tools/codemode/functions.go
diff --git a/pkg/tools/codemode/functions_test.go b/pkg/tools/codemode/functions_test.go
diff --git a/pkg/tools/mcp/gateway.go b/pkg/tools/mcp/gateway.go
@@ -17,6 +17,8 @@ import (
 )
 
 type GatewayToolset struct {
+	tools.ElicitationTool
+
 	mcpServerName string
 	config        any
 	toolFilter    []string
@@ -113,14 +115,6 @@ func (t *GatewayToolset) Stop(ctx context.Context) error {
 	return errors.Join(stopErr, cleanUpSecretsErr, cleanUpConfigErr)
 }
 
-func (t *GatewayToolset) SetElicitationHandler(tools.ElicitationHandler) {
-	// TODO: implement elicitations for the gateway
-}
-
-func (t *GatewayToolset) SetOAuthSuccessHandler(func()) {
-	// No-op, as the gateway does not support OAuth
-}
-
 func writeSecretsToFile(ctx context.Context, mcpServerName string, secrets []gateway.Secret, envProvider environment.Provider) (string, error) {
 	var secretValues []string
 	for _, secret := range secrets {
diff --git a/pkg/tools/tools.go b/pkg/tools/tools.go
@@ -6,7 +6,7 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 )
 
-type ToolHandler = func(ctx context.Context, toolCall ToolCall) (*ToolCallResult, error)
+type ToolHandler func(ctx context.Context, toolCall ToolCall) (*ToolCallResult, error)
 
 type ToolCall struct {
 	ID       string       `json:"id,omitempty"`
@@ -48,6 +48,16 @@ type ElicitationResult struct {
 // This allows the runtime to handle elicitation requests and propagate them to its own client
 type ElicitationHandler func(ctx context.Context, req *mcp.ElicitParams) (ElicitationResult, error)
 
+type ElicitationTool struct{}
+
+func (t *ElicitationTool) SetElicitationHandler(ElicitationHandler) {
+	// No-op, this tool does not use elicitation
+}
+
+func (t *ElicitationTool) SetOAuthSuccessHandler(func()) {
+	// No-op, this tool does not use OAuth
+}
+
 type ToolSet interface {
 	Tools(ctx context.Context) ([]Tool, error)
 	Instructions() string

Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ import (`
`20`	`20`	`const userAgent = "cagent/1.0"`
`21`	`21`
`22`	`22`	`type FetchTool struct {`
`23`		`- elicitationTool`
	`23`	`+ tools.ElicitationTool`
`24`	`24`	`handler *fetchHandler`
`25`	`25`	`}`
`26`	`26`
Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@ type PostEditConfig struct {`
`24`	`24`	`}`
`25`	`25`
`26`	`26`	`type FilesystemTool struct {`
`27`		`- elicitationTool`
	`27`	`+ tools.ElicitationTool`
`28`	`28`
`29`	`29`	`allowedDirectories []string`
`30`	`30`	`allowedTools []string`
Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,7 @@ import (`
`12`	`12`	`)`
`13`	`13`
`14`	`14`	`type MemoryTool struct {`
`15`		`- elicitationTool`
	`15`	`+ tools.ElicitationTool`
`16`	`16`	`manager memorymanager.Manager`
`17`	`17`	`}`
`18`	`18`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ import (`
`14`	`14`	`)`
`15`	`15`
`16`	`16`	`type ScriptShellTool struct {`
`17`		`- elicitationTool`
	`17`	`+ tools.ElicitationTool`
`18`	`18`	`shellTools map[string]latest.ScriptShellToolConfig`
`19`	`19`	`env []string`
`20`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ import (`
`15`	`15`	`)`
`16`	`16`
`17`	`17`	`type ShellTool struct {`
`18`		`- elicitationTool`
	`18`	`+ tools.ElicitationTool`
`19`	`19`	`handler *shellHandler`
`20`	`20`	`}`
`21`	`21`
Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@ import (`
`10`	`10`	`)`
`11`	`11`
`12`	`12`	`type ThinkTool struct {`
`13`		`- elicitationTool`
	`13`	`+ tools.ElicitationTool`
`14`	`14`	`handler *thinkHandler`
`15`	`15`	`}`
`16`	`16`