Better tool registry

cmd/wasm/runtime_wasm.go

@@ -125,7 +125,7 @@ func buildRuntime(ctx context.Context, cfg *latest.Config, env environment.Provi
 			case ts.Type == "filesystem" || ts.Type == "":
 				// Default toolset type is filesystem.
 				if ts.Type == "filesystem" {
-					fsTool := filesystem.NewFilesystemTool("/")
+					fsTool := filesystem.New("/")
 					opts = append(opts, agent.WithToolSets(fsTool))
 				}
 			case ts.Remote.URL != "":

examples/golibrary/builtintool/main.go

@@ -47,7 +47,7 @@ func run(ctx context.Context) error {
 				"root",
 				"You are an expert hacker",
 				agent.WithModel(llm),
-				agent.WithToolSets(shell.NewShellTool(os.Environ(), &config.RuntimeConfig{Config: config.Config{WorkingDir: "/tmp"}})),
+				agent.WithToolSets(shell.New(os.Environ(), &config.RuntimeConfig{Config: config.Config{WorkingDir: "/tmp"}})),
 			),
 		),
 	)

examples/golibrary/multi/main.go

@@ -49,7 +49,7 @@ func run(ctx context.Context) error {
 		"You are a human, with feelings and emotions.",
 		agent.WithModel(llm),
 		agent.WithSubAgents(child),
-		agent.WithToolSets(transfertask.NewTransferTaskTool()),
+		agent.WithToolSets(transfertask.New()),
 	)
 	rt, err := runtime.New(team.New(team.WithAgents(root, child)))
 	if err != nil {

pkg/acp/filesystem.go

@@ -32,7 +32,7 @@ func getSessionID(ctx context.Context) (string, bool) {
 // FilesystemToolset wraps a standard Tool and overrides read_file, write_file,
 // and edit_file to use the ACP connection for file operations
 type FilesystemToolset struct {
-	*filesystem.Tool
+	*filesystem.ToolSet
 
 	agent      *Agent
 	workingDir string
@@ -43,15 +43,15 @@ var _ tools.ToolSet = (*FilesystemToolset)(nil)
 // NewFilesystemToolset creates a new ACP-specific filesystem toolset
 func NewFilesystemToolset(agent *Agent, workingDir string, opts ...filesystem.Opt) *FilesystemToolset {
 	return &FilesystemToolset{
-		Tool:       filesystem.NewFilesystemTool(workingDir, opts...),
+		ToolSet:    filesystem.New(workingDir, opts...),
 		agent:      agent,
 		workingDir: workingDir,
 	}
 }
 
 // Tools returns the tool definitions with ACP-specific overrides
 func (t *FilesystemToolset) Tools(ctx context.Context) ([]tools.Tool, error) {
-	baseTools, err := t.Tool.Tools(ctx)
+	baseTools, err := t.ToolSet.Tools(ctx)
 	if err != nil {
 		return nil, err
 	}

pkg/acp/registry.go

@@ -11,10 +11,20 @@ import (
 )
 
 // createToolsetRegistry creates a custom toolset registry with ACP-specific filesystem toolset
-func createToolsetRegistry(agent *Agent) *teamloader.ToolsetRegistry {
-	registry := teamloader.NewDefaultToolsetRegistry()
+func createToolsetRegistry(agent *Agent) teamloader.ToolsetRegistry {
+	return &acpToolsetRegistry{
+		agent:    agent,
+		registry: teamloader.NewDefaultToolsetRegistry(),
+	}
+}
+
+type acpToolsetRegistry struct {
+	agent    *Agent
+	registry teamloader.ToolsetRegistry
+}
 
-	registry.Register("filesystem", func(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+func (r *acpToolsetRegistry) CreateTool(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, agentName string) (tools.ToolSet, error) {
+	if toolset.Type == "filesystem" {
 		wd := runConfig.WorkingDir
 		if wd == "" {
 			var err error
@@ -24,8 +34,8 @@ func createToolsetRegistry(agent *Agent) *teamloader.ToolsetRegistry {
 			}
 		}
 
-		return NewFilesystemToolset(agent, wd), nil
-	})
+		return NewFilesystemToolset(r.agent, wd), nil
+	}
 
-	return registry
+	return r.registry.CreateTool(ctx, toolset, parentDir, runConfig, agentName)
 }

pkg/app/app_test.go

@@ -60,7 +60,7 @@ func (m *mockRuntime) SessionStore() session.Store { return m.store }
 func (m *mockRuntime) Summarize(ctx context.Context, sess *session.Session, additionalPrompt string, events runtime.EventSink) {
 }
 func (m *mockRuntime) PermissionsInfo() *runtime.PermissionsInfo { return nil }
-func (m *mockRuntime) CurrentAgentSkillsToolset() *skillstool.Toolset {
+func (m *mockRuntime) CurrentAgentSkillsToolset() *skillstool.ToolSet {
 	return nil
 }
 

pkg/cli/runner_test.go

@@ -51,7 +51,7 @@ func (m *mockRuntime) ResumeElicitation(_ context.Context, action tools.Elicitat
 func (m *mockRuntime) SessionStore() session.Store                                            { return nil }
 func (m *mockRuntime) Summarize(context.Context, *session.Session, string, runtime.EventSink) {}
 func (m *mockRuntime) PermissionsInfo() *runtime.PermissionsInfo                              { return nil }
-func (m *mockRuntime) CurrentAgentSkillsToolset() *skillstool.Toolset                         { return nil }
+func (m *mockRuntime) CurrentAgentSkillsToolset() *skillstool.ToolSet                         { return nil }
 func (m *mockRuntime) CurrentMCPPrompts(context.Context) map[string]mcptools.PromptInfo {
 	return nil
 }

pkg/runtime/commands_test.go

@@ -54,7 +54,7 @@ func (m *mockRuntime) SessionStore() session.Store { return nil }
 func (m *mockRuntime) Summarize(context.Context, *session.Session, string, EventSink) {
 }
 func (m *mockRuntime) PermissionsInfo() *PermissionsInfo { return nil }
-func (m *mockRuntime) CurrentAgentSkillsToolset() *skillstool.Toolset {
+func (m *mockRuntime) CurrentAgentSkillsToolset() *skillstool.ToolSet {
 	return nil
 }
 

pkg/runtime/loop.go

@@ -841,7 +841,7 @@ func (r *LocalRuntime) configureToolsetHandlers(a *agent.Agent, events EventSink
 		// channel; a blocking send after the channel is closed would
 		// crash, and a blocking send when the consumer has gone away
 		// would deadlock.
-		if ragTool, ok := tools.As[*builtinrag.Tool](toolset); ok {
+		if ragTool, ok := tools.As[*builtinrag.ToolSet](toolset); ok {
 			ragTool.SetEventCallback(ragEventForwarder(ragTool.Name(), r, nonBlocking(events).Emit))
 		}
 	}

pkg/runtime/model_picker.go

@@ -15,14 +15,14 @@ import (
 
 // findModelPickerTool returns the Tool from the current agent's
 // toolsets, or nil if the agent has no model_picker configured.
-func (r *LocalRuntime) findModelPickerTool() *modelpicker.Tool {
+func (r *LocalRuntime) findModelPickerTool() *modelpicker.ToolSet {
 	currentName := r.CurrentAgentName()
 	a, err := r.team.Agent(currentName)
 	if err != nil {
 		return nil
 	}
 	for _, ts := range a.ToolSets() {
-		if mpt, ok := tools.As[*modelpicker.Tool](ts); ok {
+		if mpt, ok := tools.As[*modelpicker.ToolSet](ts); ok {
 			return mpt
 		}
 	}

pkg/runtime/remote_runtime.go

@@ -567,7 +567,7 @@ func (r *RemoteRuntime) ResetStartupInfo() {
 }
 
 // CurrentAgentSkillsToolset returns nil for remote runtimes since skills are managed server-side.
-func (r *RemoteRuntime) CurrentAgentSkillsToolset() *skills.Toolset {
+func (r *RemoteRuntime) CurrentAgentSkillsToolset() *skills.ToolSet {
 	return nil
 }
 

pkg/runtime/runtime.go

@@ -83,7 +83,7 @@ type Runtime interface {
 	PermissionsInfo() *PermissionsInfo
 
 	// CurrentAgentSkillsToolset returns the skills toolset for the current agent, or nil if skills are not enabled.
-	CurrentAgentSkillsToolset() *skills.Toolset
+	CurrentAgentSkillsToolset() *skills.ToolSet
 
 	// CurrentMCPPrompts returns MCP prompts available from the current agent's toolsets.
 	// Returns an empty map if no MCP prompts are available.
@@ -779,13 +779,13 @@ func (r *LocalRuntime) resolveSessionAgent(sess *session.Session) *agent.Agent {
 }
 
 // CurrentAgentSkillsToolset returns the skills toolset for the current agent, or nil if not enabled.
-func (r *LocalRuntime) CurrentAgentSkillsToolset() *skills.Toolset {
+func (r *LocalRuntime) CurrentAgentSkillsToolset() *skills.ToolSet {
 	a := r.CurrentAgent()
 	if a == nil {
 		return nil
 	}
 	for _, ts := range a.ToolSets() {
-		if st, ok := tools.As[*skills.Toolset](ts); ok {
+		if st, ok := tools.As[*skills.ToolSet](ts); ok {
 			return st
 		}
 	}

pkg/runtime/skill_runner.go

@@ -21,7 +21,7 @@ import (
 // its final response is returned as the tool result.
 //
 // All skill-specific business rules (lookup, fork-mode validation, content
-// expansion) live in (*skills.Toolset).PrepareForkSubSession; this
+// expansion) live in (*skills.ToolSet).PrepareForkSubSession; this
 // handler keeps only the runtime-private orchestration that runForwarding
 // can't generalise — namely the optional model override that applies for
 // the sub-session's lifetime.

pkg/session/session_test.go

@@ -10,10 +10,18 @@ import (
 
 	"github.com/docker/docker-agent/pkg/agent"
 	"github.com/docker/docker-agent/pkg/chat"
+	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/tools"
 	"github.com/docker/docker-agent/pkg/tools/builtin/todo"
 )
 
+func todoToolSet(t *testing.T) tools.ToolSet {
+	t.Helper()
+	toolSet, err := todo.CreateToolSet(latest.Toolset{})
+	require.NoError(t, err)
+	return toolSet
+}
+
 func TestTrimMessagesWithToolCalls(t *testing.T) {
 	messages := []chat.Message{
 		{
@@ -173,7 +181,7 @@ func TestGetMessages_Instructions(t *testing.T) {
 }
 
 func TestGetMessages_CacheControl(t *testing.T) {
-	testAgent := agent.New("root", "instructions", agent.WithToolSets(&todo.Tool{}))
+	testAgent := agent.New("root", "instructions", agent.WithToolSets(todoToolSet(t)))
 
 	s := New()
 	messages := s.GetMessages(testAgent)
@@ -197,7 +205,7 @@ func TestGetMessages_CacheControlWithSummary(t *testing.T) {
 	//     buildContextSpecificSystemMessages caching behavior.
 	//   - Summary and conversation messages are not cache-controlled.
 	testAgent := agent.New("root", "instructions",
-		agent.WithToolSets(&todo.Tool{}),
+		agent.WithToolSets(todoToolSet(t)),
 	)
 
 	s := New()