refactor: enhance sandbox process termination, update state directory path, and expand manager and host sandbox test coverage.

cmd/picoclaw/internal/onboard/command.go

@@ -6,6 +6,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
+//go:generate rm -rf workspace
 //go:generate cp -r ../../../../workspace .
 //go:embed workspace
 var embeddedFiles embed.FS

pkg/agent/sandbox/container.go

@@ -356,7 +356,7 @@ func (c *ContainerSandbox) ensureContainer(ctx context.Context) error {
 			_ = removeRegistryEntry(regPath, c.cfg.ContainerName)
 			return c.createAndStart(ctx)
 		}
-		// LOGIC-1: Container is actively running; recreating it now would disrupt
+		// Container is actively running; recreating it now would disrupt
 		// in-flight work. Log a warning so operators can detect configuration drift.
 		// The container will be recreated on the next cold start or prune cycle.
 		logger.WarnCF(
@@ -468,7 +468,7 @@ func (c *ContainerSandbox) registryPath() string {
 }
 
 func (c *ContainerSandbox) sandboxStateDir() string {
-	return filepath.Join(infra.ResolveHomeDir(), "sandboxes")
+	return filepath.Join(infra.ResolveHomeDir(), "sandbox")
 }
 
 func (c *ContainerSandbox) stopAndRemoveContainer(ctx context.Context, containerName string) error {

pkg/agent/sandbox/container_test.go

@@ -325,7 +325,7 @@ func TestContainerSandbox_RegistryPath_UsesSandboxStateDir(t *testing.T) {
 		Workspace:     "/tmp/ws",
 		WorkspaceRoot: "/tmp/sbx",
 	})
-	want := filepath.Join(home, ".picoclaw", "sandboxes", "containers.json")
+	want := filepath.Join(home, ".picoclaw", "sandbox", "containers.json")
 	if got := sb.registryPath(); got != want {
 		t.Fatalf("registryPath = %q, want %q", got, want)
 	}
@@ -335,7 +335,7 @@ func TestContainerSandbox_RegistryPath_UsesPicoClawHomeOverride(t *testing.T) {
 	picoHome := t.TempDir()
 	t.Setenv("PICOCLAW_HOME", picoHome)
 	sb := NewContainerSandbox(ContainerSandboxConfig{})
-	want := filepath.Join(picoHome, "sandboxes", "containers.json")
+	want := filepath.Join(picoHome, "sandbox", "containers.json")
 	if got := sb.registryPath(); got != want {
 		t.Fatalf("registryPath = %q, want %q", got, want)
 	}
@@ -649,3 +649,11 @@ func TestSyncAgentWorkspace_SyncsSkillsDirectory(t *testing.T) {
 		t.Fatalf("skill file content mismatch. got: %s", string(content))
 	}
 }
+
+func TestContainerSandbox_Resolve(t *testing.T) {
+	c := NewContainerSandbox(ContainerSandboxConfig{})
+	sb, err := c.Resolve(context.Background())
+	if err != nil || sb != c {
+		t.Fatal("expected Resolve to return self")
+	}
+}

pkg/agent/sandbox/host.go

@@ -114,6 +114,9 @@ func (h *HostSandbox) ExecStream(
 	}
 
 	prepareCommandForTermination(cmd)
+	cmd.Cancel = func() error {
+		return terminateProcessTree(cmd)
+	}
 
 	if err := cmd.Start(); err != nil {
 		return nil, err
@@ -168,11 +171,9 @@ func (h *HostSandbox) ExecStream(
 
 	waitErr := cmd.Wait()
 	if streamErr != nil {
-		_ = terminateProcessTree(cmd)
 		return nil, streamErr
 	}
 	if cmdCtx.Err() != nil {
-		_ = terminateProcessTree(cmd)
 		return nil, cmdCtx.Err()
 	}
 

pkg/agent/sandbox/host_process_unix_test.go

@@ -0,0 +1,65 @@
+//go:build !windows
+
+package sandbox
+
+import (
+	"os/exec"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestPrepareCommandForTermination(t *testing.T) {
+	// Should not panic on nil
+	prepareCommandForTermination(nil)
+
+	cmd := exec.Command("echo", "test")
+	prepareCommandForTermination(cmd)
+
+	if cmd.SysProcAttr == nil {
+		t.Fatal("expected SysProcAttr to be initialized")
+	}
+	if !cmd.SysProcAttr.Setpgid {
+		t.Fatal("expected Setpgid to be true")
+	}
+}
+
+func TestTerminateProcessTree(t *testing.T) {
+	// Should not panic on nil cmd or nil process
+	if err := terminateProcessTree(nil); err != nil {
+		t.Fatalf("expected nil error for nil cmd, got: %v", err)
+	}
+
+	cmdUnstarted := exec.Command("echo", "test")
+	if err := terminateProcessTree(cmdUnstarted); err != nil {
+		t.Fatalf("expected nil error for unstarted cmd, got: %v", err)
+	}
+
+	// Start a real dummy process to test killing
+	cmd := exec.Command("sleep", "1")
+	prepareCommandForTermination(cmd)
+	if err := cmd.Start(); err != nil {
+		t.Fatalf("failed to start cmd: %v", err)
+	}
+
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- cmd.Wait()
+	}()
+
+	if err := terminateProcessTree(cmd); err != nil {
+		t.Fatalf("terminateProcessTree failed: %v", err)
+	}
+
+	// Verify the process is dead by waiting for Wait() to return
+	require.Eventually(t, func() bool {
+		select {
+		case err := <-errChan:
+			// Process died (killed), err should not be nil
+			return err != nil
+		default:
+			return false
+		}
+	}, 2*time.Second, 50*time.Millisecond, "expected process Wait to finish after termination")
+}

pkg/agent/sandbox/host_test.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"testing"
+	"time"
 )
 
 func TestHostSandbox_StartStopFs(t *testing.T) {
@@ -392,3 +393,40 @@ func TestValidatePathErrors(t *testing.T) {
 		t.Fatalf("expected error when ancestor is file")
 	}
 }
+
+func TestHostSandbox_ExecStream_Cancellation(t *testing.T) {
+	root := t.TempDir()
+	sb := NewHostSandbox(root, false)
+
+	// Test graceful timeout/cancel.
+	// We'll run a bash script that sleeps indefinitely in the foreground.
+	// We expect the command to be terminated and streamErr or ctx.Err() returned.
+	req := ExecRequest{
+		Command:   "sleep 3600",
+		TimeoutMs: 100, // Very short timeout
+	}
+
+	start := time.Now()
+	res, err := sb.Exec(context.Background(), req)
+	elapsed := time.Since(start)
+
+	if err == nil {
+		t.Fatalf("expected timeout error for sleep command, got nil. Res: %v", res)
+	}
+	if elapsed > 5*time.Second {
+		t.Fatalf("command failed to time out reasonably quickly: %v", elapsed)
+	}
+}
+
+func TestHostSandbox_ExecStream_NoCommand(t *testing.T) {
+	root := t.TempDir()
+	sb := NewHostSandbox(root, false)
+
+	req := ExecRequest{
+		Command: "",
+	}
+	_, err := sb.Exec(context.Background(), req)
+	if err == nil || err.Error() != "empty command" {
+		t.Fatalf("expected 'empty command', got: %v", err)
+	}
+}

pkg/agent/sandbox/manager.go

@@ -287,7 +287,7 @@ func (m *scopedSandboxManager) pruneOnce(ctx context.Context) error {
 	if m.pruneIdleHours <= 0 && m.pruneMaxAgeDays <= 0 {
 		return nil
 	}
-	regPath := filepath.Join(infra.ResolveHomeDir(), "sandboxes", defaultSandboxRegistryFile)
+	regPath := filepath.Join(infra.ResolveHomeDir(), "sandbox", defaultSandboxRegistryFile)
 	registryMu.Lock()
 	data, err := loadRegistry(regPath)
 	registryMu.Unlock()
@@ -315,12 +315,17 @@ func (m *scopedSandboxManager) pruneOnce(ctx context.Context) error {
 		if !shouldPruneEntry(pruneCfg, now, entry) {
 			continue
 		}
+
+		// Best-effort cleanup: we attempt to prune or remove all eligible entries.
+		// If one fails, we record the error and continue to the next to prevent
+		// a single bad container from halting the entire garbage collection process.
 		if sb, ok := byContainer[entry.ContainerName]; ok {
 			if err := sb.Prune(ctx); err != nil && firstErr == nil {
 				firstErr = err
 			}
 			continue
 		}
+
 		if err := stopAndRemoveContainerByName(ctx, entry.ContainerName); err != nil && firstErr == nil {
 			firstErr = err
 		}

pkg/agent/sandbox/manager_test.go

@@ -2,8 +2,10 @@ package sandbox
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -95,7 +97,7 @@ func TestScopedSandboxManager_PruneLoopLifecycle(t *testing.T) {
 func TestScopedSandboxManager_PruneOnceLoadRegistryError(t *testing.T) {
 	home := t.TempDir()
 	t.Setenv("HOME", home)
-	stateDir := filepath.Join(home, ".picoclaw", "sandboxes")
+	stateDir := filepath.Join(home, ".picoclaw", "sandbox")
 	if err := os.MkdirAll(stateDir, 0o755); err != nil {
 		t.Fatalf("mkdir state dir: %v", err)
 	}
@@ -139,3 +141,293 @@ func TestScopedSandboxManager_ShouldSandbox_NonMain(t *testing.T) {
 		t.Fatal("expected non-main session to use sandbox path")
 	}
 }
+
+func TestHostOnlyManager(t *testing.T) {
+	// hostOnlyManager just delegates to its inner host sandbox.
+	workspace := t.TempDir()
+	host := NewHostSandbox(workspace, false)
+	mgr := &hostOnlyManager{host: host}
+
+	ctx := context.Background()
+	if err := mgr.Start(ctx); err != nil {
+		t.Fatalf("Start() returned error: %v", err)
+	}
+	if err := mgr.Prune(ctx); err != nil {
+		t.Fatalf("Prune() returned error: %v", err)
+	}
+
+	sb, err := mgr.Resolve(ctx)
+	if err != nil {
+		t.Fatalf("Resolve() returned error: %v", err)
+	}
+	if sb == nil {
+		t.Fatal("Resolve() returned nil sandbox")
+	}
+
+	if fs := mgr.Fs(); fs == nil {
+		t.Fatal("Fs() returned nil")
+	}
+
+	gotWs := mgr.GetWorkspace(ctx)
+	if gotWs != host.GetWorkspace(ctx) {
+		t.Fatalf("GetWorkspace() = %q, want %q", gotWs, host.GetWorkspace(ctx))
+	}
+
+	// Exec and ExecStream should also just delegate without panic.
+	// Executing a simple command like "echo"
+	req := ExecRequest{Command: "echo", Args: []string{"test"}}
+	res, err := mgr.Exec(ctx, req)
+	if err != nil {
+		t.Fatalf("Exec() returned error: %v", err)
+	}
+	if res.ExitCode != 0 {
+		t.Fatalf("Exec() returned non-zero exit code: %d", res.ExitCode)
+	}
+
+	streamRes, streamErr := mgr.ExecStream(ctx, req, func(e ExecEvent) error { return nil })
+	if streamErr != nil {
+		t.Fatalf("ExecStream() returned error: %v", streamErr)
+	}
+	if streamRes.ExitCode != 0 {
+		t.Fatalf("ExecStream() returned non-zero exit code: %d", streamRes.ExitCode)
+	}
+}
+
+func TestUnavailableSandboxManager(t *testing.T) {
+	errReason := os.ErrPermission
+	mgr := NewUnavailableSandboxManager(errReason)
+	ctx := context.Background()
+
+	if err := mgr.Start(ctx); err != errReason {
+		t.Fatalf("Start() = %v, want %v", err, errReason)
+	}
+
+	// Prune is a no-op, shouldn't return error
+	if err := mgr.Prune(ctx); err != nil {
+		t.Fatalf("Prune() = %v, want nil", err)
+	}
+
+	_, err := mgr.Resolve(ctx)
+	if err == nil {
+		t.Fatal("Resolve() expected error, got nil")
+	}
+
+	if ws := mgr.GetWorkspace(ctx); ws != "" {
+		t.Fatalf("GetWorkspace() = %q, want empty", ws)
+	}
+
+	req := ExecRequest{Command: "ls"}
+	_, err = mgr.Exec(ctx, req)
+	if err == nil {
+		t.Fatal("Exec() expected error, got nil")
+	}
+
+	_, err = mgr.ExecStream(ctx, req, nil)
+	if err == nil {
+		t.Fatal("ExecStream() expected error, got nil")
+	}
+
+	fs := mgr.Fs()
+	if fs == nil {
+		t.Fatal("Fs() returned nil")
+	}
+
+	_, err = fs.ReadFile(ctx, "test.txt")
+	if err == nil {
+		t.Fatal("ReadFile() expected error, got nil")
+	}
+
+	err = fs.WriteFile(ctx, "test.txt", []byte("a"), false)
+	if err == nil {
+		t.Fatal("WriteFile() expected error, got nil")
+	}
+
+	_, err = fs.ReadDir(ctx, ".")
+	if err == nil {
+		t.Fatal("ReadDir() expected error, got nil")
+	}
+}
+
+func TestScopedSandboxManager_Delegates(t *testing.T) {
+	home := t.TempDir()
+	t.Setenv("HOME", home)
+	ws := filepath.Join(home, "default_ws")
+
+	m := &scopedSandboxManager{
+		mode:    config.SandboxModeNonMain,
+		agentID: "agent-1",
+		host:    NewHostSandbox(ws, false),
+		scoped:  map[string]Sandbox{},
+	}
+	_ = m.host.Start(context.Background())
+	m.fs = &managerFS{m: m}
+
+	// 1. When ShouldSandbox is false (Context is main session), it delegates to HostSandbox.
+	ctxMain := WithSessionKey(context.Background(), routing.BuildAgentMainSessionKey("agent-1"))
+
+	if got := m.GetWorkspace(ctxMain); got != ws {
+		t.Fatalf("GetWorkspace(main) = %q, want %q", got, ws)
+	}
+
+	req := ExecRequest{Command: "echo", Args: []string{"hello"}}
+	res, err := m.Exec(ctxMain, req)
+	if err != nil {
+		t.Fatalf("Exec(main) error: %v", err)
+	}
+	if res.ExitCode != 0 {
+		t.Fatalf("Exec(main) exit code: %d", res.ExitCode)
+	}
+
+	_, err = m.ExecStream(ctxMain, req, func(e ExecEvent) error { return nil })
+	if err != nil {
+		t.Fatalf("ExecStream(main) error: %v", err)
+	}
+
+	fs := m.Fs()
+	testFile := "test_delegate.txt"
+	err = fs.WriteFile(ctxMain, testFile, []byte("ok"), true)
+	if err != nil {
+		t.Fatalf("WriteFile(main) error: %v", err)
+	}
+	defer os.Remove(filepath.Join(ws, testFile))
+
+	data, err := fs.ReadFile(ctxMain, testFile)
+	if err != nil || string(data) != "ok" {
+		t.Fatalf("ReadFile(main) error: %v, data: %q", err, string(data))
+	}
+
+	entries, err := fs.ReadDir(ctxMain, ".")
+	if err != nil || len(entries) == 0 {
+		t.Fatalf("ReadDir(main) error: %v, len: %d", err, len(entries))
+	}
+
+	sb, err := m.Resolve(ctxMain)
+	if err != nil {
+		t.Fatalf("Resolve(main) error: %v", err)
+	}
+	if sb != m.host {
+		t.Fatal("Resolve(main) should return host sandbox")
+	}
+}
+
+func TestScopedSandboxManager_ContainerDelegates(t *testing.T) {
+	home := t.TempDir()
+	t.Setenv("HOME", home)
+	ws := filepath.Join(home, "container_ws")
+	os.MkdirAll(ws, 0o755)
+
+	mockContainer := NewHostSandbox(ws, false)
+	_ = mockContainer.Start(context.Background())
+
+	m := &scopedSandboxManager{
+		mode:    config.SandboxModeAll,
+		agentID: "agent-1",
+		scoped:  map[string]Sandbox{},
+	}
+	m.fs = &managerFS{m: m}
+
+	ctx := WithSessionKey(context.Background(), "test-session")
+	scopeKey := m.scopeKeyFromContext(ctx)
+
+	// Pre-inject the mock container to bypass actual docker creation
+	m.scoped[scopeKey] = mockContainer
+
+	// Now shouldSandbox(ctx) is true, so manager methods should delegate to mockContainer
+	if got := m.GetWorkspace(ctx); got != ws {
+		t.Fatalf("GetWorkspace(container) = %q, want %q", got, ws)
+	}
+
+	req := ExecRequest{Command: "echo", Args: []string{"hello"}}
+	res, err := m.Exec(ctx, req)
+	if err != nil {
+		t.Fatalf("Exec(container) error: %v", err)
+	}
+	if res.ExitCode != 0 {
+		t.Fatalf("Exec(container) exit code: %d", res.ExitCode)
+	}
+
+	_, err = m.ExecStream(ctx, req, func(e ExecEvent) error { return nil })
+	if err != nil {
+		t.Fatalf("ExecStream(container) error: %v", err)
+	}
+
+	fs := m.Fs()
+	testFile := "test_container_delegate.txt"
+	err = fs.WriteFile(ctx, testFile, []byte("ok"), true)
+	if err != nil {
+		t.Fatalf("WriteFile(container) error: %v", err)
+	}
+	defer os.Remove(filepath.Join(ws, testFile))
+
+	data, err := fs.ReadFile(ctx, testFile)
+	if err != nil || string(data) != "ok" {
+		t.Fatalf("ReadFile(container) error: %v, data: %q", err, string(data))
+	}
+
+	entries, err := fs.ReadDir(ctx, ".")
+	if err != nil || len(entries) == 0 {
+		t.Fatalf("ReadDir(container) error: %v, len: %d", err, len(entries))
+	}
+
+	sb, err := m.Resolve(ctx)
+	if err != nil {
+		t.Fatalf("Resolve(container) error: %v", err)
+	}
+	if sb != mockContainer {
+		t.Fatal("Resolve(container) should return the mock container sandbox")
+	}
+}
+
+func TestScopedSandboxManager_ContainerCreationError(t *testing.T) {
+	home := t.TempDir()
+	t.Setenv("HOME", home)
+
+	m := &scopedSandboxManager{
+		mode:          config.SandboxModeAll,
+		workspaceRoot: home,
+		dockerCfg:     config.AgentSandboxDockerConfig{Image: "non-existent-image-12345"},
+		scoped:        map[string]Sandbox{},
+	}
+
+	ctx := WithSessionKey(context.Background(), "error-session")
+
+	// Fast-path misses, falls through to buildScopedContainerSandbox and Start()
+	// Start() succeeds due to lazy evaluation, but Exec() should fail because
+	// Docker is either unavailable or the image is missing
+	sb, err := m.Resolve(ctx)
+	if err != nil {
+		t.Fatalf("expected Resolve() to succeed lazily, got error: %v", err)
+	}
+
+	res, err := sb.Exec(ctx, ExecRequest{Command: "echo"})
+	if err == nil && res.ExitCode == 0 {
+		t.Fatalf("expected Exec() to fail when container creation fails, but it succeeded: %#v", res)
+	}
+}
+
+func TestScopedSandboxManager_PruneOnceFull(t *testing.T) {
+	home := t.TempDir()
+	t.Setenv("HOME", home)
+	stateDir := filepath.Join(home, ".picoclaw", "sandbox")
+	_ = os.MkdirAll(stateDir, 0o755)
+
+	m := &scopedSandboxManager{ // this is correct initialization for the test
+		mode:            config.SandboxModeAll,
+		pruneIdleHours:  1,
+		pruneMaxAgeDays: 0,
+		scoped:          map[string]Sandbox{},
+	}
+
+	regPath := filepath.Join(stateDir, "containers.json")
+	oldTime := time.Now().Add(-2 * time.Hour).UnixMilli()
+
+	data := fmt.Sprintf(`{"entries": [{"container_name": "prune-me", "last_active_at": %d}]}`, oldTime)
+	_ = os.WriteFile(regPath, []byte(data), 0o644)
+
+	_ = m.pruneOnce(context.Background())
+
+	b, _ := os.ReadFile(regPath)
+	if strings.Contains(string(b), "prune-me") {
+		t.Fatalf("expected prune-me to be removed from registry, got %s", string(b))
+	}
+}

pkg/agent/sandbox/sandbox_test.go

@@ -0,0 +1,50 @@
+package sandbox
+
+import (
+	"context"
+	"testing"
+)
+
+func TestContextHelpers(t *testing.T) {
+	// Test SessionKey
+	ctx := context.Background()
+	if got := SessionKeyFromContext(ctx); got != "" {
+		t.Fatalf("expected empty session key, got %q", got)
+	}
+
+	ctx = WithSessionKey(ctx, "session-123")
+	if got := SessionKeyFromContext(ctx); got != "session-123" {
+		t.Fatalf("expected session-123, got %q", got)
+	}
+
+	// Test nil contexts
+	if got := SessionKeyFromContext(nil); got != "" { //nolint:staticcheck
+		t.Fatalf("SessionKeyFromContext(nil) = %q, want empty", got)
+	}
+	if got := FromContext(nil); got != nil { //nolint:staticcheck
+		t.Fatalf("FromContext(nil) = %v, want nil", got)
+	}
+	if got := managerFromContext(nil); got != nil { //nolint:staticcheck
+		t.Fatalf("managerFromContext(nil) = %v, want nil", got)
+	}
+
+	// Test Sandbox context
+	mockSb := &unavailableSandboxManager{}
+	ctx = WithSandbox(context.Background(), mockSb)
+	if got := FromContext(ctx); got != mockSb {
+		t.Fatalf("expected to retrieve mock sandbox from context")
+	}
+
+	// Test Manager context resolving
+	mockMgr := NewUnavailableSandboxManager(nil)
+	ctx = WithManager(context.Background(), mockMgr)
+
+	if got := managerFromContext(ctx); got != mockMgr {
+		t.Fatalf("expected to retrieve mock manager from context")
+	}
+
+	// FromContext with Manager only should attempt to Resolve (which returns error/nil here)
+	if got := FromContext(ctx); got != nil {
+		t.Fatalf("expected nil from FromContext when Resolve fails, got %v", got)
+	}
+}

scripts/build-sandbox.sh

@@ -15,6 +15,6 @@ if [ ! -f "${DOCKERFILE}" ]; then
     exit 1
 fi
 
-docker build -t "${IMAGE_NAME}" -f "${DOCKERFILE}" "${REPO_ROOT}"
+docker build --no-cache -t "${IMAGE_NAME}" -f "${DOCKERFILE}" "${REPO_ROOT}"
 
 echo "Successfully built ${IMAGE_NAME}"