iop/apps/node/internal/workspace/cleanup_test.go

package workspace

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"golang.org/x/sys/unix"

	iop "iop/proto/gen/iop"
)

func TestWorkspaceCleanupArtifactsDuplicateRaceAndIsolation(t *testing.T) {
	root := t.TempDir()
	runtime, err := NewRuntime([]*iop.WorkspaceConfig{testWorkspaceConfig(root)}, "darwin", nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = runtime.Close() })
	for _, requestID := range []string{"request-a", "request-b"} {
		if _, err := runtime.Open(testRequestAuthority(requestID)); err != nil {
			t.Fatal(err)
		}
	}
	if err := runtime.WriteInternalArtifact("request-a", "plan.md", []byte("plan")); err != nil {
		t.Fatal(err)
	}
	if err := runtime.WriteInternalArtifact("request-a", "nested/review.md", []byte("review")); err != nil {
		t.Fatal(err)
	}
	if err := runtime.WriteInternalArtifact("request-b", "plan.md", []byte("foreign request")); err != nil {
		t.Fatal(err)
	}
	userResult := filepath.Join(root, "result.txt")
	if err := os.WriteFile(userResult, []byte("preserve"), 0o600); err != nil {
		t.Fatal(err)
	}

	const callers = 24
	results := make(chan CleanupResult, callers)
	var group sync.WaitGroup
	for range callers {
		group.Add(1)
		go func() {
			defer group.Done()
			results <- runtime.Cleanup(context.Background(), "request-a")
		}()
	}
	group.Wait()
	close(results)
	var first *CleanupResult
	for result := range results {
		if first == nil {
			copy := result
			first = &copy
		}
		if result != *first {
			t.Fatalf("cleanup callers observed different results: first=%+v got=%+v", *first, result)
		}
	}
	if first == nil || first.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS || first.CleanedArtifacts != 4 {
		t.Fatalf("cleanup result = %+v", first)
	}
	if duplicate := runtime.Cleanup(context.Background(), "request-a"); duplicate != *first {
		t.Fatalf("duplicate cleanup = %+v, want %+v", duplicate, *first)
	}
	if _, err := runtime.Open(testRequestAuthority("request-a")); err != ErrRequestConflict {
		t.Fatalf("completed request identity reopened: %v", err)
	}
	if _, err := os.Lstat(requestArtifactRoot(root, "request-a")); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("request-a tree remains: %v", err)
	}
	if data, err := os.ReadFile(filepath.Join(requestArtifactRoot(root, "request-b"), "plan.md")); err != nil || string(data) != "foreign request" {
		t.Fatalf("request-b artifact = %q, %v", data, err)
	}
	if data, err := os.ReadFile(userResult); err != nil || string(data) != "preserve" {
		t.Fatalf("user result = %q, %v", data, err)
	}
}

func TestWorkspaceCleanupCancelsActiveProcessGroup(t *testing.T) {
	root := t.TempDir()
	runtime := newCommandRuntime(t, root, 64)
	openCommandRequest(t, runtime, "request-process", 64)
	userResult := filepath.Join(root, "user-result.txt")
	if err := os.WriteFile(userResult, []byte("preserve"), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := runtime.WriteInternalArtifact("request-process", "plan.md", []byte("plan")); err != nil {
		t.Fatal(err)
	}
	pidFile := filepath.Join(root, "cleanup-child.pid")
	input := commandInput("request-process", "tool-process", "group")
	input.Environment["IOP_CHILD_PID_FILE"] = pidFile
	resultCh := make(chan Result, 1)
	go func() { resultCh <- runtime.ExecuteCommand(context.Background(), input) }()
	waitForFile(t, pidFile)
	payload, err := os.ReadFile(pidFile)
	if err != nil {
		t.Fatal(err)
	}
	pid, err := strconv.Atoi(strings.TrimSpace(string(payload)))
	if err != nil {
		t.Fatal(err)
	}

	cleanup := runtime.Cleanup(context.Background(), "request-process")
	if cleanup.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS || cleanup.CleanedProcesses != 1 || cleanup.CleanedArtifacts != 2 {
		t.Fatalf("cleanup = %+v", cleanup)
	}
	if result := <-resultCh; result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED {
		t.Fatalf("command result = %+v", result)
	}
	deadline := time.Now().Add(2 * time.Second)
	for processExists(pid) && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if processExists(pid) {
		t.Fatalf("descendant process %d survived cleanup", pid)
	}
	if _, err := os.Lstat(requestArtifactRoot(root, "request-process")); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("request artifacts remain: %v", err)
	}
	if data, err := os.ReadFile(userResult); err != nil || string(data) != "preserve" {
		t.Fatalf("cleanup changed user result: %q, %v", data, err)
	}
}

func TestWorkspaceCleanupTimeoutIsBoundedAndCached(t *testing.T) {
	runtime, root := openedRuntime(t)
	if err := runtime.WriteInternalArtifact("request-1", "plan.md", []byte("preserve on timeout")); err != nil {
		t.Fatal(err)
	}
	req, err := runtime.Request("request-1")
	if err != nil {
		t.Fatal(err)
	}
	execution := newCommandExecution()
	key := commandKey{requestID: "request-1", toolCallID: "tool-stuck"}
	if !runtime.registerCommand(req, key, execution) {
		t.Fatal("failed to install deterministic stuck command")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	started := time.Now()
	result := runtime.Cleanup(ctx, "request-1")
	if result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_TIMEOUT || result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_TIMEOUT || result.CleanedProcesses != 1 {
		t.Fatalf("cleanup timeout = %+v", result)
	}
	if time.Since(started) > time.Second {
		t.Fatalf("cleanup exceeded bound: %s", time.Since(started))
	}
	if duplicate := runtime.Cleanup(context.Background(), "request-1"); duplicate != result {
		t.Fatalf("cached timeout = %+v, want %+v", duplicate, result)
	}
	if _, err := os.Stat(filepath.Join(requestArtifactRoot(root, "request-1"), "plan.md")); err != nil {
		t.Fatalf("timed-out cleanup removed artifact: %v", err)
	}
	execution.finish()
	runtime.commandsMu.Lock()
	delete(runtime.activeCommands, key)
	runtime.commandsMu.Unlock()
}

func TestWorkspaceCleanupRefusesUnownedAndUnsafeEntries(t *testing.T) {
	tests := map[string]func(*testing.T, *Runtime, string, string){
		"unowned entry": func(t *testing.T, _ *Runtime, root, requestID string) {
			if err := os.WriteFile(filepath.Join(requestArtifactRoot(root, requestID), "injected.txt"), []byte("unowned"), 0o600); err != nil {
				t.Fatal(err)
			}
		},
		"symlink": func(t *testing.T, _ *Runtime, root, requestID string) {
			if err := os.Symlink(filepath.Join(root, "outside"), filepath.Join(requestArtifactRoot(root, requestID), "link")); err != nil {
				t.Fatal(err)
			}
		},
		"identity replacement": func(t *testing.T, runtime *Runtime, root, requestID string) {
			if err := runtime.WriteInternalArtifact(requestID, "plan.md", []byte("owned")); err != nil {
				t.Fatal(err)
			}
			target := filepath.Join(requestArtifactRoot(root, requestID), "plan.md")
			if err := os.Remove(target); err != nil {
				t.Fatal(err)
			}
			if err := os.WriteFile(target, []byte("replacement"), 0o600); err != nil {
				t.Fatal(err)
			}
		},
		"special file": func(t *testing.T, _ *Runtime, root, requestID string) {
			if err := unix.Mkfifo(filepath.Join(requestArtifactRoot(root, requestID), "pipe"), 0o600); err != nil {
				t.Fatal(err)
			}
		},
		"mount device boundary": func(t *testing.T, runtime *Runtime, _ string, requestID string) {
			req, err := runtime.Request(requestID)
			if err != nil {
				t.Fatal(err)
			}
			req.mu.Lock()
			rootArtifact := req.artifacts[req.internalPrefix]
			rootArtifact.device++
			req.artifacts[req.internalPrefix] = rootArtifact
			req.mu.Unlock()
		},
	}
	for name, inject := range tests {
		t.Run(name, func(t *testing.T) {
			root := t.TempDir()
			runtime, err := NewRuntime([]*iop.WorkspaceConfig{testWorkspaceConfig(root)}, "darwin", nil)
			if err != nil {
				t.Fatal(err)
			}
			t.Cleanup(func() { _ = runtime.Close() })
			requestID := "request-refuse"
			if _, err := runtime.Open(testRequestAuthority(requestID)); err != nil {
				t.Fatal(err)
			}
			userResult := filepath.Join(root, "user-result.txt")
			if err := os.WriteFile(userResult, []byte("preserve"), 0o600); err != nil {
				t.Fatal(err)
			}
			inject(t, runtime, root, requestID)
			result := runtime.Cleanup(context.Background(), requestID)
			if result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_ERROR || result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INTERNAL {
				t.Fatalf("cleanup = %+v", result)
			}
			if _, err := os.Lstat(requestArtifactRoot(root, requestID)); err != nil {
				t.Fatalf("suspect request tree was removed: %v", err)
			}
			if data, err := os.ReadFile(userResult); err != nil || string(data) != "preserve" {
				t.Fatalf("failed cleanup changed user result: %q, %v", data, err)
			}
		})
	}
}

func TestWorkspaceCleanupRejectsInvalidOrPreexistingIdentity(t *testing.T) {
	root := t.TempDir()
	runtime, err := NewRuntime([]*iop.WorkspaceConfig{testWorkspaceConfig(root)}, "darwin", nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = runtime.Close() })
	if result := runtime.Cleanup(context.Background(), "../foreign"); result.Code != iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_INVALID_REQUEST {
		t.Fatalf("invalid cleanup = %+v", result)
	}
	preexisting := requestArtifactRoot(root, "request-existing")
	if err := os.MkdirAll(preexisting, 0o700); err != nil {
		t.Fatal(err)
	}
	if _, err := runtime.Open(testRequestAuthority("request-existing")); err != ErrInvalidRequest {
		t.Fatalf("preexisting request namespace was admitted: %v", err)
	}
	if _, err := os.Stat(preexisting); err != nil {
		t.Fatalf("preexisting namespace was changed: %v", err)
	}
}

func TestWorkspaceCleanupRuntimeCloseUsesSamePrimitive(t *testing.T) {
	root := t.TempDir()
	runtime, err := NewRuntime([]*iop.WorkspaceConfig{testWorkspaceConfig(root)}, "darwin", nil)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := runtime.Open(testRequestAuthority("request-close")); err != nil {
		t.Fatal(err)
	}
	userResult := filepath.Join(root, "user-result.txt")
	if err := os.WriteFile(userResult, []byte("preserve"), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := runtime.WriteInternalArtifact("request-close", "review.md", []byte("review")); err != nil {
		t.Fatal(err)
	}
	if err := runtime.Close(); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Lstat(requestArtifactRoot(root, "request-close")); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("runtime close left request artifacts: %v", err)
	}
	if data, err := os.ReadFile(userResult); err != nil || string(data) != "preserve" {
		t.Fatalf("runtime close changed user result: %q, %v", data, err)
	}
}

func TestWorkspaceCleanupCompletedCacheIsBounded(t *testing.T) {
	root := t.TempDir()
	runtime, err := NewRuntime([]*iop.WorkspaceConfig{testWorkspaceConfig(root)}, "darwin", nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = runtime.Close() })
	for index := 0; index < completedCleanupLimit+32; index++ {
		requestID := "request-cache-" + strconv.Itoa(index)
		if _, err := runtime.Open(testRequestAuthority(requestID)); err != nil {
			t.Fatal(err)
		}
		if result := runtime.Cleanup(context.Background(), requestID); result.Status != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS {
			t.Fatalf("cleanup %d = %+v", index, result)
		}
	}
	runtime.cleanupMu.Lock()
	calls, order := len(runtime.cleanupCalls), len(runtime.cleanupOrder)
	runtime.cleanupMu.Unlock()
	if calls != completedCleanupLimit || order != completedCleanupLimit {
		t.Fatalf("completed cleanup cache = calls %d order %d, want %d", calls, order, completedCleanupLimit)
	}
}

func requestArtifactRoot(root, requestID string) string {
	return filepath.Join(root, ".iop", "job", requestID)
}