iop/apps/edge/internal/service/single_request_cleanup_test.go

package service

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	iop "iop/proto/gen/iop"
)

type countingWorkspaceLifecycle struct {
	openCount    atomic.Int32
	toolCount    atomic.Int32
	cleanupCount atomic.Int32
	cleanupStart chan struct{}
	cleanupGate  chan struct{}
	cleanupErr   error
	startOnce    sync.Once
}

func newCountingWorkspaceLifecycle(block bool, cleanupErr error) *countingWorkspaceLifecycle {
	runtime := &countingWorkspaceLifecycle{cleanupStart: make(chan struct{}), cleanupErr: cleanupErr}
	if block {
		runtime.cleanupGate = make(chan struct{})
	}
	return runtime
}

func (r *countingWorkspaceLifecycle) workspaceOpen(_ context.Context, _ *SingleRequestWorkspaceBinding, req *iop.WorkspaceOpenRequest) (*iop.WorkspaceOpenResponse, error) {
	r.openCount.Add(1)
	return &iop.WorkspaceOpenResponse{RequestId: req.GetRequestId(), WorkspaceRef: req.GetWorkspaceRef(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}, nil
}

func (r *countingWorkspaceLifecycle) workspaceTool(_ context.Context, _ *SingleRequestWorkspaceBinding, req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
	r.toolCount.Add(1)
	return &iop.WorkspaceToolResponse{
		RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(),
		Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, Content: []byte("workspace result"),
	}, nil
}

func (r *countingWorkspaceLifecycle) CleanupWorkspace(ctx context.Context, _ *SingleRequestWorkspaceBinding, _ string) error {
	r.cleanupCount.Add(1)
	r.startOnce.Do(func() { close(r.cleanupStart) })
	if r.cleanupGate != nil {
		select {
		case <-r.cleanupGate:
		case <-ctx.Done():
			return ErrSingleRequestWorkspaceCleanup
		}
	}
	return r.cleanupErr
}

func cleanupTestRequest(t *testing.T) SingleRequestRequest {
	t.Helper()
	binding := createTestBinding(t)
	binding.Workspace = &SingleRequestWorkspaceBinding{
		Ref: "workspace-ref-123", NodeID: "node-cleanup", ConnectionGeneration: 7,
		OperationIDs: []string{"read"},
		Limits:       SingleRequestWorkspaceLimits{MaxReadBytes: 1024},
	}
	return SingleRequestRequest{RequestID: "request-cleanup", Binding: binding, Prompt: "complete work"}
}

func startCleanupExecution(t *testing.T, runtime *countingWorkspaceLifecycle) SingleRequestExecution {
	t.Helper()
	executor := newScriptedInternalToolExecutor(InternalWorkspaceToolCall{
		ToolCallID: "tool-read", Name: InternalWorkspaceToolRead,
		Arguments: []byte(`{"relative_path":"README.md"}`),
	})
	handle, err := startSingleRequestWithToolLoop(context.Background(), executor, executor, runtime, cleanupTestRequest(t))
	if err != nil {
		t.Fatal(err)
	}
	return handle
}

func waitCleanupStarted(t *testing.T, runtime *countingWorkspaceLifecycle) {
	t.Helper()
	select {
	case <-runtime.cleanupStart:
	case <-time.After(2 * time.Second):
		t.Fatal("workspace cleanup did not start")
	}
}

func TestSingleRequestCleanupPrecedesSuccessfulTerminal(t *testing.T) {
	runtime := newCountingWorkspaceLifecycle(true, nil)
	handle := startCleanupExecution(t, runtime)
	waitCleanupStarted(t, runtime)
	if handle.State() != SingleRequestStateFinalizing {
		t.Fatalf("state during cleanup = %s", handle.State())
	}
	internal := handle.(*singleRequestHandle)
	internal.mu.Lock()
	cleanupComplete := internal.cleanupComplete
	internal.mu.Unlock()
	if cleanupComplete {
		t.Fatal("cleanup completed before its lifecycle gate was released")
	}
	if err := handle.AcknowledgeTerminal(true); !errors.Is(err, ErrSingleRequestInvalidState) {
		t.Fatalf("early acknowledgement = %v", err)
	}
	close(runtime.cleanupGate)
	waitForSingleRequestCleanup(t, handle)
	if err := handle.AcknowledgeTerminal(true); err != nil {
		t.Fatal(err)
	}
	result, err := waitForExecution(t, handle)
	if err != nil || result.Output != "private tools completed" {
		t.Fatalf("Wait = (%q, %v)", result.Output, err)
	}
	if runtime.openCount.Load() != 1 || runtime.toolCount.Load() != 1 || runtime.cleanupCount.Load() != 1 {
		t.Fatalf("open/tool/cleanup = %d/%d/%d", runtime.openCount.Load(), runtime.toolCount.Load(), runtime.cleanupCount.Load())
	}
}

func TestSingleRequestCleanupFailureFailsClosed(t *testing.T) {
	runtime := newCountingWorkspaceLifecycle(false, errors.New("raw node cleanup detail"))
	handle := startCleanupExecution(t, runtime)
	_, err := waitForExecution(t, handle)
	if !errors.Is(err, ErrSingleRequestWorkspaceCleanup) || handle.State() != SingleRequestStateFailed {
		t.Fatalf("state/error = %s/%v", handle.State(), err)
	}
	if err != nil && err.Error() != ErrSingleRequestWorkspaceCleanup.Error() {
		t.Fatalf("cleanup leaked implementation detail: %q", err)
	}
	if runtime.cleanupCount.Load() != 1 {
		t.Fatalf("cleanup calls = %d", runtime.cleanupCount.Load())
	}
}

func TestSingleRequestCleanupTerminalRacesExactlyOnce(t *testing.T) {
	for iteration := 0; iteration < 20; iteration++ {
		runtime := newCountingWorkspaceLifecycle(true, nil)
		handle := startCleanupExecution(t, runtime)
		waitCleanupStarted(t, runtime)
		var group sync.WaitGroup
		group.Add(4)
		go func() { defer group.Done(); handle.Cancel() }()
		go func() { defer group.Done(); _ = handle.AcknowledgeTerminal(false) }()
		go func() {
			defer group.Done()
			_ = handle.SubmitEnvelope(SingleRequestEnvelope{RequestID: "request-cleanup", Sequence: 99, Stage: SingleRequestStateFailed, Err: ErrSingleRequestFailed})
		}()
		go func() { defer group.Done(); handle.Cancel() }()
		group.Wait()
		close(runtime.cleanupGate)
		_, _ = waitForExecution(t, handle)
		if runtime.cleanupCount.Load() != 1 {
			t.Fatalf("iteration %d cleanup calls = %d", iteration, runtime.cleanupCount.Load())
		}
	}
}

func TestSingleRequestCleanupPreservesCancelAndWriteFailureCategory(t *testing.T) {
	t.Run("cancel", func(t *testing.T) {
		runtime := newCountingWorkspaceLifecycle(true, ErrSingleRequestWorkspaceCleanup)
		handle := startCleanupExecution(t, runtime)
		waitCleanupStarted(t, runtime)
		handle.Cancel()
		close(runtime.cleanupGate)
		_, err := waitForExecution(t, handle)
		if !errors.Is(err, ErrSingleRequestCancelled) || !errors.Is(err, ErrSingleRequestWorkspaceCleanup) || handle.State() != SingleRequestStateCancelled {
			t.Fatalf("cancel state/error = %s/%v", handle.State(), err)
		}
	})

	t.Run("terminal write failure", func(t *testing.T) {
		runtime := newCountingWorkspaceLifecycle(false, nil)
		handle := startCleanupExecution(t, runtime)
		waitForSingleRequestCleanup(t, handle)
		if err := handle.AcknowledgeTerminal(false); err != nil {
			t.Fatal(err)
		}
		_, err := waitForExecution(t, handle)
		if !errors.Is(err, ErrSingleRequestFailed) || handle.State() != SingleRequestStateFailed {
			t.Fatalf("write failure state/error = %s/%v", handle.State(), err)
		}
		if runtime.cleanupCount.Load() != 1 {
			t.Fatalf("write failure cleanup calls = %d", runtime.cleanupCount.Load())
		}
	})
}

func TestSingleRequestCleanupSkipsUnopenedWorkspace(t *testing.T) {
	runtime := newCountingWorkspaceLifecycle(false, nil)
	executor := &channelFakeExecutor{fn: func(_ context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
		return submitToFinalizing(req, ctrl, &SingleRequestResult{Output: "no workspace"})
	}}
	handle, err := startSingleRequestWithToolLoop(context.Background(), executor, nil, runtime, cleanupTestRequest(t))
	if err != nil {
		t.Fatal(err)
	}
	waitForState(t, handle, SingleRequestStateFinalizing)
	if err := handle.AcknowledgeTerminal(true); err != nil {
		t.Fatal(err)
	}
	if _, err := waitForExecution(t, handle); err != nil {
		t.Fatal(err)
	}
	if runtime.cleanupCount.Load() != 0 || runtime.openCount.Load() != 0 {
		t.Fatalf("unopened workspace lifecycle = open %d cleanup %d", runtime.openCount.Load(), runtime.cleanupCount.Load())
	}
}