Scope worker cancellation and closure acknowledgements to executions

internal/types/messages.go

@@ -143,8 +143,9 @@ type TaskAssignmentMessage struct {
 
 // TaskClaimedMessage is sent from worker to server after successfully claiming a task
 type TaskClaimedMessage struct {
-	TaskID   string `json:"task_id"`
-	WorkerID string `json:"worker_id"`
+	TaskID      string `json:"task_id"`
+	WorkerID    string `json:"worker_id"`
+	ExecutionID string `json:"execution_id,omitempty"`
 }
 
 // TaskCompletedMessage tells the server to end the active run execution after a successful agent process exit.
@@ -178,7 +179,8 @@ type TaskRejectedMessage struct {
 
 // TaskCancellationMessage is sent from server to worker to cancel an active task.
 type TaskCancellationMessage struct {
-	TaskID string `json:"task_id"`
+	TaskID      string `json:"task_id"`
+	ExecutionID string `json:"execution_id,omitempty"`
 }
 
 // TaskState is the serialized terminal task state accepted by warp-server.

internal/worker/backend.go

@@ -178,9 +178,9 @@ type Backend interface {
 	// failures are surfaced as ExecuteOutcomeError with an error that is (or
 	// wraps) *TaskFailure.
 	ExecuteTask(ctx context.Context, params *TaskParams) ExecuteResult
-	// CancelTask makes a best-effort attempt to cancel a task. The worker
-	// invokes it for every task cancellation, alongside cancelling the
-	// ExecuteTask context.
+	// CancelTask requests cancellation after ExecuteTask returns. A nil error
+	// means the request was accepted, not that asynchronous resources stopped.
+	// Failed requests may be retried and must be execution-scoped and idempotent.
 	CancelTask(ctx context.Context, params *CancelParams) error
 	// PreservesTasksOnShutdown reports whether active task execution units can
 	// safely outlive the worker process during shutdown.

internal/worker/command_integration_test.go

@@ -94,7 +94,7 @@ func TestIntegrationCommandBackendDispatchSuppressesTerminalMessage(t *testing.T
 	waitFor(t, 5*time.Second, func() bool {
 		w.tasksMutex.Lock()
 		defer w.tasksMutex.Unlock()
-		task, ok := w.activeTasks["task-1"]
+		task, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "exec-1"}]
 		return ok && task.spawned
 	})
 

internal/worker/direct.go

@@ -169,15 +169,27 @@
 		return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("invalid task ID for workspace path: %w", err)))
 	}
 
-	// Determine working directory: shared target dir or per-task workspace.
+	// Determine working directory: shared target dir or isolated workspace.
 	var workspaceDir string
 	usingTargetDir := b.config.TargetDir != ""
 
 	if usingTargetDir {
 		workspaceDir = b.config.TargetDir
 	} else {
-		// Create per-task workspace directory.
+		// Legacy assignments use a per-task workspace.
 		workspaceDir = filepath.Join(b.config.WorkspaceRoot, taskID)
+		if params.ExecutionID != "" {
+			if err := validateTaskIDForPath(params.ExecutionID); err != nil {
+				return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("invalid execution ID for workspace path: %w", err)))
+			}
+			var err error
+			// Give each execution a fresh workspace so overlapping executions
+			// cannot share files or remove each other's workspace during cleanup.
+			workspaceDir, err = os.MkdirTemp(b.config.WorkspaceRoot, taskID+"-"+params.ExecutionID+"-")
+			if err != nil {
+				return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("failed to create execution workspace: %w", err)))
+			}
+		}
 		if err := os.MkdirAll(workspaceDir, 0700); err != nil {
 			return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonWorkspaceSetup, fmt.Errorf("failed to create workspace directory: %w", err)))
 		}

internal/worker/dispatch_test.go

@@ -37,7 +37,7 @@ func newDispatchWorker(backend Backend) *Worker {
 		ctx:         context.Background(),
 		config:      Config{},
 		outbound:    newOutboundQueue(4),
-		activeTasks: map[string]activeTask{"task-1": {cancel: func() {}, executionID: "exec-1"}},
+		activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "exec-1"}: {cancel: func() {}, executionID: "exec-1"}},
 		backend:     backend,
 	}
 }
@@ -61,7 +61,7 @@ func TestExecuteTaskDispatchedSuppressesTerminalMessage(t *testing.T) {
 		t.Fatalf("expected no terminal message after dispatch, got %q", msg.Type)
 	}
 	w.tasksMutex.Lock()
-	task, ok := w.activeTasks["task-1"]
+	task, ok := w.activeTasks[taskExecution{taskID: "task-1", executionID: "exec-1"}]
 	w.tasksMutex.Unlock()
 	if !ok {
 		t.Fatal("spawned task should remain in activeTasks")
@@ -104,7 +104,7 @@ func TestHandleTaskCancellationRoutesToBackendCancelTask(t *testing.T) {
 	w := &Worker{
 		ctx:         context.Background(),
 		outbound:    newOutboundQueue(1),
-		activeTasks: map[string]activeTask{"task-1": spawnedActiveTask("exec-1")},
+		activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "exec-1"}: spawnedActiveTask("exec-1")},
 		backend:     backend,
 	}
 
@@ -119,11 +119,9 @@ func TestHandleTaskCancellationRoutesToBackendCancelTask(t *testing.T) {
 		t.Fatal("CancelTask was not invoked for a spawned task")
 	}
 
-	w.tasksMutex.Lock()
-	_, ok := w.activeTasks["task-1"]
-	w.tasksMutex.Unlock()
-	if ok {
-		t.Error("spawned task should be removed after cancellation is routed")
+	w.cancellationWG.Wait()
+	if w.activeTaskCount() != 0 {
+		t.Error("accepted cancellation should remove spawned task tracking")
 	}
 }
 
@@ -134,7 +132,7 @@ func TestHandleTaskCancellationRunningTaskCancelsContextAndBackend(t *testing.T)
 	w := &Worker{
 		ctx:      context.Background(),
 		outbound: newOutboundQueue(1),
-		activeTasks: map[string]activeTask{"task-1": {
+		activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "exec-1"}: {
 			ctx:         taskCtx,
 			cancel:      taskCancel,
 			executionID: "exec-1",
@@ -156,29 +154,24 @@ func TestHandleTaskCancellationRunningTaskCancelsContextAndBackend(t *testing.T)
 		t.Fatal("CancelTask was not invoked for a running task")
 	}
 
-	// The entry stays until executeTask's deferred cleanup removes it.
-	w.tasksMutex.Lock()
-	task, ok := w.activeTasks["task-1"]
-	w.tasksMutex.Unlock()
-	if !ok {
-		t.Fatal("running task should remain in activeTasks until executeTask returns")
-	}
-	if task.cancellationSource != taskCancellationSourceUser {
-		t.Fatalf("cancellation source = %q, want %q", task.cancellationSource, taskCancellationSourceUser)
+	w.cancellationWG.Wait()
+	if w.activeTaskCount() != 0 {
+		t.Fatal("accepted cancellation should remove task tracking")
 	}
 }
 
 func TestHandleTaskCancellationSpawnedNoopCancelEmitsNoMessage(t *testing.T) {
 	w := &Worker{
 		ctx:         context.Background(),
 		outbound:    newOutboundQueue(1),
-		activeTasks: map[string]activeTask{"task-1": spawnedActiveTask("exec-1")},
+		activeTasks: map[taskExecution]activeTask{{taskID: "task-1", executionID: "exec-1"}: spawnedActiveTask("exec-1")},
 		backend:     &dispatchBackend{},
 	}
 
 	// Must not panic and must not emit any task status message when the
 	// backend's CancelTask is a no-op.
 	w.handleTaskCancellation(&types.TaskCancellationMessage{TaskID: "task-1"})
+	w.cancellationWG.Wait()
 
 	if len(w.outbound.messages) != 0 {
 		t.Fatalf("expected no message for no-op-cancel spawned task, got %d", len(w.outbound.messages))

internal/worker/docker.go

@@ -67,6 +67,7 @@ type DockerBackend struct {
 	platform       string // Docker daemon platform (e.g., "linux/amd64" or "linux/arm64")
 	platformSpec   ocispec.Platform
 	sidecarVolumes sync.Map // volume name -> *sidecarVolumeState
+	taskContainers sync.Map // taskExecution -> container ID; retained after cleanup failures
 }
 
 // NewDockerBackend creates a new Docker backend, connecting to the Docker daemon.
@@ -181,24 +182,33 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe
 	}
 
 	doneContainerStart := params.SetupEvents.startPhase(ctx, SetupEventContainerStart)
+	key := taskExecution{params.TaskID, params.ExecutionID}
+	// Transport and server failures cannot prove that the daemon created no container.
+	b.taskContainers.Store(key, "")
 	resp, err := dockerClient.ContainerCreate(ctx, client.ContainerCreateOptions{
 		Config:     containerConfig,
 		HostConfig: hostConfig,
 	})
 	if err != nil {
+		if cerrdefs.IsInvalidArgument(err) || cerrdefs.IsNotFound(err) || cerrdefs.IsConflict(err) {
+			b.taskContainers.Delete(key)
+		}
 		doneContainerStart(true)
 		return executeError(newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonContainerCreate, fmt.Errorf("failed to create container: %w", err)))
 	}
 
 	containerID := resp.ID
+	b.taskContainers.Store(key, containerID)
 	log.Debugf(ctx, "Created Docker container: %s", containerID)
 
 	defer func() {
 		if containerID != "" && !b.config.NoCleanup {
 			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), BackendShutdownTimeout)
 			defer cleanupCancel()
-			if _, removeErr := dockerClient.ContainerRemove(cleanupCtx, containerID, client.ContainerRemoveOptions{Force: true}); removeErr != nil {
+			if _, removeErr := dockerClient.ContainerRemove(cleanupCtx, containerID, client.ContainerRemoveOptions{Force: true}); removeErr != nil && !cerrdefs.IsNotFound(removeErr) {
 				log.Debugf(ctx, "Container %s already removed or removal failed: %v", containerID, removeErr)
+			} else {
+				b.taskContainers.Delete(key)
 			}
 		}
 	}()
@@ -221,6 +231,11 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) Exe
 		}
 	case status := <-waitResult.Result:
 		log.Debugf(ctx, "Container exited with status code: %d", status.StatusCode)
+		if b.config.NoCleanup && status.Error == nil {
+			// Preserve the exited container for debugging, but not its tracking.
+			// Later cancellation intentionally leaves this stopped container alone.
+			b.taskContainers.Delete(key)
+		}
 
 		if params.Logs == nil && (zerolog.GlobalLevel() <= zerolog.DebugLevel || status.StatusCode != 0) {
 			logOutput, logErr := b.getContainerLogs(ctx, dockerClient, containerID)
@@ -272,9 +287,22 @@ func dockerResourcesForShape(shape *types.InstanceShape) container.Resources {
 	return res
 }
 
-// CancelTask is a no-op: cancelling the ExecuteTask context fully stops a
-// Docker-backend task.
-func (b *DockerBackend) CancelTask(context.Context, *CancelParams) error { return nil }
+// CancelTask removes the captured execution container, including when normal cleanup failed.
+func (b *DockerBackend) CancelTask(ctx context.Context, params *CancelParams) error {
+	key := taskExecution{params.TaskID, params.ExecutionID}
+	containerID, ok := b.taskContainers.Load(key)
+	if !ok {
+		return nil
+	}
+	if containerID == "" {
+		return fmt.Errorf("container creation outcome is unknown for execution %s", params.ExecutionID)
+	}
+	if _, err := b.dockerClient.ContainerRemove(ctx, containerID.(string), client.ContainerRemoveOptions{Force: true}); err != nil && !cerrdefs.IsNotFound(err) {
+		return err
+	}
+	b.taskContainers.Delete(key)
+	return nil
+}
 
 // Shutdown closes the Docker client.
 func (b *DockerBackend) Shutdown(ctx context.Context) {