Cleanup after MCP Server runs

Signed-off-by: David Gageot <david.gageot@docker.com>

Author: dgageot

pkg/teamloader/teamloader.go

@@ -256,17 +256,17 @@ func getToolsForAgent(a *latest.AgentConfig, parentDir string, sharedTools map[s
 			serverName := strings.TrimPrefix(toolset.Ref, "docker:")
 			args := []string{"mcp", "gateway", "run", "--servers=" + serverName}
 
+			var cleanUp func() error
 			if toolset.Config != nil {
-				serverConfig := map[string]any{
-					serverName: toolset.Config,
-				}
-
 				file, err := os.CreateTemp("", "mcp-config-*.yaml")
 				if err != nil {
 					return nil, fmt.Errorf("failed to create temp file: %w", err)
 				}
-				// TODO(dga): Delete the temp file after use.
+				cleanUp = func() error { return os.Remove(file.Name()) }
 
+				serverConfig := map[string]any{
+					serverName: toolset.Config,
+				}
 				if err := yaml.NewEncoder(file).Encode(serverConfig); err != nil {
 					return nil, fmt.Errorf("failed to write config to temp file: %w", err)
 				}
@@ -279,10 +279,10 @@ func getToolsForAgent(a *latest.AgentConfig, parentDir string, sharedTools map[s
 			args = append(args, "--catalog=https://desktop.docker.com/mcp/catalog/v2/catalog.yaml")
 
 			// TODO(dga): If the server's docker image had the right annotations, we could run it directly with `docker run` or with the MCP gateway as a go library.
-			t = append(t, mcp.NewToolsetCommand("docker", args, os.Environ(), toolset.Tools))
+			t = append(t, mcp.NewToolsetCommand("docker", args, os.Environ(), toolset.Tools, cleanUp))
 
 		case toolset.Type == "mcp" && toolset.Command != "":
-			t = append(t, mcp.NewToolsetCommand(toolset.Command, toolset.Args, os.Environ(), toolset.Tools))
+			t = append(t, mcp.NewToolsetCommand(toolset.Command, toolset.Args, os.Environ(), toolset.Tools, nil))
 
 		case toolset.Type == "mcp" && toolset.Remote.URL != "":
 			// Expand headers.

pkg/tools/mcp/toolset.go

@@ -13,18 +13,20 @@ import (
 type Toolset struct {
 	c          *Client
 	toolFilter []string
+	cleanup    func() error
 }
 
 // Make sure the MCP Toolset always implements _our_ ToolSet interface
 var _ tools.ToolSet = (*Toolset)(nil)
 
 // NewToolsetCommand creates a new MCP toolset from a command.
-func NewToolsetCommand(command string, args, env, toolFilter []string) *Toolset {
+func NewToolsetCommand(command string, args, env, toolFilter []string, cleanup func() error) *Toolset {
 	slog.Debug("Creating MCP toolset", "command", command, "args", args, "toolFilter", toolFilter)
 
 	return &Toolset{
 		c:          NewStdioClient(command, args, env),
 		toolFilter: toolFilter,
+		cleanup:    cleanup,
 	}
 }
 
@@ -86,6 +88,12 @@ func (t *Toolset) Stop() error {
 		slog.Error("Failed to stop MCP toolset", "error", err)
 		return err
 	}
+	if t.cleanup != nil {
+		if err := t.cleanup(); err != nil {
+			slog.Error("Failed to cleanup MCP toolset", "error", err)
+			return err
+		}
+	}
 	slog.Debug("Stopped MCP toolset successfully")
 	return nil
 }