Merge pull request #2157 from dgageot/board/we-don-t-need-filewritetracker-anymore-i-07f1f5ff

Remove unused fileWriteTracker from creator package

Author: dgageot

pkg/creator/agent.go

@@ -5,18 +5,14 @@ package creator
 import (
 	"context"
 	_ "embed"
-	"encoding/json"
 	"fmt"
 	"strings"
 
 	"github.com/goccy/go-yaml"
 
 	"github.com/docker/docker-agent/pkg/config"
-	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/teamloader"
-	"github.com/docker/docker-agent/pkg/tools"
-	"github.com/docker/docker-agent/pkg/tools/builtin"
 )
 
 //go:embed instructions.txt
@@ -46,14 +42,11 @@ func Agent(ctx context.Context, runConfig *config.RuntimeConfig, modelNameOverri
 		return nil, fmt.Errorf("building creator config: %w", err)
 	}
 
-	registry := createToolsetRegistry(runConfig.WorkingDir)
-
 	return teamloader.Load(
 		ctx,
 		config.NewBytesSource("creator", configYAML),
 		runConfig,
 		teamloader.WithModelOverrides([]string{modelNameOverride}),
-		teamloader.WithToolsetRegistry(registry),
 	)
 }
 
@@ -111,64 +104,3 @@ func buildCreatorConfigYAML(instructions string) ([]byte, error) {
 
 	return yaml.Marshal(fullConfig)
 }
-
-// createToolsetRegistry creates a custom toolset registry that wraps the filesystem
-// toolset to track file paths written by the agent.
-func createToolsetRegistry(workingDir string) *teamloader.ToolsetRegistry {
-	tracker := &fileWriteTracker{
-		ToolSet: builtin.NewFilesystemTool(workingDir),
-	}
-
-	registry := teamloader.NewDefaultToolsetRegistry()
-	registry.Register("filesystem", func(context.Context, latest.Toolset, string, *config.RuntimeConfig, string) (tools.ToolSet, error) {
-		return tracker, nil
-	})
-
-	return registry
-}
-
-// fileWriteTracker wraps a filesystem toolset to track files written by the agent.
-// This allows the creator to know what files were created during the session.
-type fileWriteTracker struct {
-	tools.ToolSet
-	originalWriteFileHandler tools.ToolHandler
-	path                     string
-}
-
-// Tools returns the available tools, wrapping the write_file tool to track paths.
-func (t *fileWriteTracker) Tools(ctx context.Context) ([]tools.Tool, error) {
-	innerTools, err := t.ToolSet.Tools(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	for i, tool := range innerTools {
-		if tool.Name == builtin.ToolNameWriteFile {
-			t.originalWriteFileHandler = tool.Handler
-			innerTools[i].Handler = t.trackWriteFile
-		}
-	}
-
-	return innerTools, nil
-}
-
-// trackWriteFile intercepts write_file calls to track the path being written.
-func (t *fileWriteTracker) trackWriteFile(ctx context.Context, toolCall tools.ToolCall) (*tools.ToolCallResult, error) {
-	var args struct {
-		Path    string `json:"path"`
-		Content string `json:"content"`
-	}
-	if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &args); err != nil {
-		return nil, fmt.Errorf("failed to parse write_file arguments: %w", err)
-	}
-
-	t.path = args.Path
-
-	return t.originalWriteFileHandler(ctx, toolCall)
-}
-
-// LastWrittenPath returns the path of the last file written by the agent.
-// Returns an empty string if no file has been written yet.
-func (t *fileWriteTracker) LastWrittenPath() string {
-	return t.path
-}

pkg/creator/agent_test.go

@@ -3,58 +3,13 @@ package creator
 import (
 	"testing"
 
-	"github.com/goccy/go-yaml"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/docker/docker-agent/pkg/config"
-	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/environment"
 )
 
-func TestAgentConfigYAML(t *testing.T) {
-	t.Parallel()
-
-	// Build the same structure as the buildCreatorConfigYAML function
-	agentToolsets := []map[string]any{
-		{"type": "shell"},
-		{"type": "filesystem"},
-	}
-
-	rootAgent := yaml.MapSlice{
-		{Key: "model", Value: "auto"},
-		{Key: "welcome_message", Value: "Hello! I'm here to create agents for you.\n\nCan you explain to me what the agent will be used for?"},
-		{Key: "instruction", Value: "Some test instructions"},
-		{Key: "toolsets", Value: agentToolsets},
-	}
-
-	agentsMapSlice := yaml.MapSlice{
-		{Key: "root", Value: rootAgent},
-	}
-
-	newAgentConfig := yaml.MapSlice{
-		{Key: "agents", Value: agentsMapSlice},
-	}
-
-	data, err := yaml.Marshal(newAgentConfig)
-	require.NoError(t, err)
-
-	t.Logf("YAML output:\n%s", string(data))
-
-	// Verify it can be loaded by the config loader
-	cfg, err := config.Load(t.Context(), config.NewBytesSource("test", data))
-	require.NoError(t, err)
-
-	// Verify the config has the expected structure
-	require.Len(t, cfg.Agents, 1)
-	assert.Equal(t, "root", cfg.Agents[0].Name)
-	assert.Equal(t, "auto", cfg.Agents[0].Model)
-	assert.Contains(t, cfg.Agents[0].WelcomeMessage, "Hello!")
-	require.Len(t, cfg.Agents[0].Toolsets, 2)
-	assert.Equal(t, "shell", cfg.Agents[0].Toolsets[0].Type)
-	assert.Equal(t, "filesystem", cfg.Agents[0].Toolsets[1].Type)
-}
-
 func TestBuildCreatorConfigYAML(t *testing.T) {
 	t.Parallel()
 
@@ -144,45 +99,6 @@ func TestAgent(t *testing.T) {
 	assert.Contains(t, toolNames, "write_file")
 }
 
-func TestFileWriteTracker(t *testing.T) {
-	t.Parallel()
-
-	ctx := t.Context()
-	runConfig := &config.RuntimeConfig{
-		Config: config.Config{
-			WorkingDir: t.TempDir(),
-		},
-	}
-
-	registry := createToolsetRegistry(runConfig.WorkingDir)
-	require.NotNil(t, registry)
-
-	// Create the toolset through the registry
-	toolset, err := registry.CreateTool(ctx, latest.Toolset{Type: "filesystem"}, runConfig.WorkingDir, runConfig, "test-agent")
-	require.NoError(t, err)
-	require.NotNil(t, toolset)
-
-	// Verify the toolset is a file write tracker
-	tracker, ok := toolset.(*fileWriteTracker)
-	require.True(t, ok, "expected fileWriteTracker, got %T", toolset)
-
-	// Initially, no path should be tracked
-	assert.Empty(t, tracker.LastWrittenPath())
-
-	// Get the tools and verify write_file is present
-	tools, err := tracker.Tools(ctx)
-	require.NoError(t, err)
-
-	var hasWriteFile bool
-	for _, tool := range tools {
-		if tool.Name == "write_file" {
-			hasWriteFile = true
-			break
-		}
-	}
-	assert.True(t, hasWriteFile, "write_file tool should be present")
-}
-
 func TestBuildCreatorConfigYAML_MultilineStrings(t *testing.T) {
 	t.Parallel()