feat: don't use context for executor (#622)

Author: ShubyM

pkg/agent/conversation.go

@@ -277,13 +277,10 @@ func (s *Agent) Init(ctx context.Context) error {
 	s.workDir = workDir
 
 	// Register tools with executor if none registered yet
-	// We need to preserve existing tools (e.g. custom tools) while ensuring we have a fresh map
-	// for this agent instance to avoid polluting the global default tools.
-	existingTools := s.Tools.AllTools()
-	s.Tools.Init()
-	for _, tool := range existingTools {
-		s.Tools.RegisterTool(tool)
-	}
+	// We clone existing tools (e.g. custom tools) to ensure we have a fresh map
+	// This avoids polluting the global default tools and ensures thread safety.
+	s.Tools = s.Tools.CloneWithExecutor(s.executor)
+
 	s.Tools.RegisterTool(tools.NewBashTool(s.executor))
 	s.Tools.RegisterTool(tools.NewKubectlTool(s.executor))
 
@@ -913,6 +910,9 @@ func (c *Agent) NewSession() (string, error) {
 		c.executor = sb
 		klog.Info("Created new sandbox for new session", "name", sandboxName)
 
+		// Re-bind all tools to the new executor
+		c.Tools = c.Tools.CloneWithExecutor(c.executor)
+
 		c.Tools.RegisterTool(tools.NewBashTool(c.executor))
 		c.Tools.RegisterTool(tools.NewKubectlTool(c.executor))
 		c.sessionMu.Unlock()

pkg/tools/custom_tool.go

@@ -36,7 +36,8 @@ type CustomToolConfig struct {
 
 // CustomTool implements the Tool interface for external commands.
 type CustomTool struct {
-	config CustomToolConfig
+	config   CustomToolConfig
+	executor sandbox.Executor
 }
 
 // NewCustomTool creates a new CustomTool instance.
@@ -139,11 +140,9 @@ func (t *CustomTool) Run(ctx context.Context, args map[string]any) (any, error)
 	workDir := ctx.Value(WorkDirKey).(string)
 	env := os.Environ()
 
-	// Get executor from context or default to local
-	var executor sandbox.Executor
-	if v := ctx.Value(ExecutorKey); v != nil {
-		executor = v.(sandbox.Executor)
-	} else {
+	// Use the injected executor, or fallback to local if not set (e.g. for global instance)
+	executor := t.executor
+	if executor == nil {
 		executor = sandbox.NewLocalExecutor()
 	}
 
@@ -159,3 +158,12 @@ func (t *CustomTool) CheckModifiesResource(args map[string]any) string {
 	// For custom tools, we'll conservatively use "unknown" since we can't
 	return "unknown"
 }
+
+// CloneWithExecutor creates a copy of the CustomTool with the given executor.
+// This is used to create a session-specific instance of the tool.
+func (t *CustomTool) CloneWithExecutor(executor sandbox.Executor) *CustomTool {
+	return &CustomTool{
+		config:   t.config,
+		executor: executor,
+	}
+}

pkg/tools/custom_tool_test.go

@@ -15,7 +15,11 @@
 package tools
 
 import (
+	"context"
+	"strings"
 	"testing"
+
+	"github.com/GoogleCloudPlatform/kubectl-ai/pkg/sandbox"
 )
 
 func TestCustomTool_AddCommandPrefix(t *testing.T) {
@@ -111,3 +115,64 @@ func TestCustomTool_AddCommandPrefix(t *testing.T) {
 		})
 	}
 }
+
+// MockExecutor implements sandbox.Executor for testing
+type MockExecutor struct {
+	CapturedCommand string
+	CapturedEnv     []string
+	CapturedWorkDir string
+}
+
+func (m *MockExecutor) Execute(ctx context.Context, command string, env []string, workDir string) (*sandbox.ExecResult, error) {
+	m.CapturedCommand = command
+	m.CapturedEnv = env
+	m.CapturedWorkDir = workDir
+	return &sandbox.ExecResult{Stdout: "executed"}, nil
+}
+
+func (m *MockExecutor) Close(ctx context.Context) error {
+	return nil
+}
+
+func TestCustomTool_CloneWithExecutor(t *testing.T) {
+	config := CustomToolConfig{
+		Name:    "test-tool",
+		Command: "echo",
+	}
+
+	tool, err := NewCustomTool(config)
+	if err != nil {
+		t.Fatalf("failed to create tool: %v", err)
+	}
+
+	mockExec := &MockExecutor{}
+	clonedTool := tool.CloneWithExecutor(mockExec)
+
+	ctx := context.WithValue(context.Background(), WorkDirKey, "/tmp")
+	args := map[string]any{
+		"command": "hello",
+	}
+
+	result, err := clonedTool.Run(ctx, args)
+	if err != nil {
+		t.Fatalf("tool run failed: %v", err)
+	}
+
+	execResult, ok := result.(*sandbox.ExecResult)
+	if !ok {
+		t.Fatalf("expected *sandbox.ExecResult, got %T", result)
+	}
+	if execResult.Stdout != "executed" {
+		t.Errorf("expected Stdout 'executed', got %q", execResult.Stdout)
+	}
+
+	if mockExec.CapturedCommand != "echo hello" {
+		t.Errorf("expected command 'echo hello', got %q", mockExec.CapturedCommand)
+	}
+	if !strings.Contains(strings.Join(mockExec.CapturedEnv, "\n"), "PATH") {
+		t.Errorf("expected captured environment to contain PATH")
+	}
+	if mockExec.CapturedWorkDir != "/tmp" {
+		t.Errorf("expected workdir '/tmp', got %q", mockExec.CapturedWorkDir)
+	}
+}

pkg/tools/tools.go

@@ -98,6 +98,25 @@ func (t *Tools) RegisterTool(tool Tool) {
 	t.tools[name] = tool
 }
 
+// CloneWithExecutor creates a shallow copy of the Tools collection,
+// but clones any tools that need a session-specific executor (like CustomTool).
+func (t *Tools) CloneWithExecutor(executor sandbox.Executor) Tools {
+	newTools := Tools{
+		tools: make(map[string]Tool),
+	}
+
+	for name, tool := range t.tools {
+		// If it's a CustomTool, we need to clone it with the session-specific executor
+		if ct, ok := tool.(*CustomTool); ok {
+			newTools.tools[name] = ct.CloneWithExecutor(executor)
+		} else {
+			// For other tools (like MCP tools), we reuse the existing instance
+			newTools.tools[name] = tool
+		}
+	}
+	return newTools
+}
+
 type ToolCall struct {
 	tool      Tool
 	name      string