[REMOTE-1820] Handle cancellation requests from dispatch server (#83)

## Summary

This is the worker-side half of REMOTE-1820: self-hosted worker cancellation support.

The companion `warp-server` change now forwards `task_cancellation` control messages to the worker. This PR makes `oz-agent-worker` consume those messages, cancel the matching active task context, and report the task back as terminal `CANCELLED` rather than as a normal failure.

End to end:
- Adds the `task_cancellation` WebSocket message and a shared `TaskStateCancelled` wire value.
- Tracks each active task's context and cancel function so the worker can cancel only the requested task.
- Cancels matching active tasks when a `task_cancellation` control message arrives; cancellations for inactive tasks are logged and ignored.
- Reports cancelled tasks through the existing `task_completed` message shape with `task_state=CANCELLED`, so server-side lifecycle reconciliation can distinguish cancellation from failure.
- Preserves the normal success/failure paths for non-cancelled tasks.
- Records cancellation lifecycle events and marks cancelled executions separately from failed executions.
- Handles context cancellation cleanly in the direct backend and uses bounded Docker cleanup so worker shutdown does not block indefinitely.

## Code overview
- `internal/types/messages.go` — adds `task_cancellation`, `TaskState`, and optional `task_state` fields on terminal worker messages.
- `internal/worker/worker.go` — handles cancellation messages, tracks active task contexts, emits terminal `CANCELLED`, and keeps shutdown cancellation on the same lifecycle path.
- `internal/worker/direct.go` — classifies context cancellation as task cancellation instead of a generic backend failure.
- `internal/worker/docker.go` — bounds container cleanup during cancellation/shutdown.
- `internal/worker/worker_test.go` — covers cancellation message handling, terminal `CANCELLED` reporting, and shutdown behavior.

## Validation
- `go test ./internal/metrics ./internal/worker`
- `go test ./...`
- `helm template oz-agent-worker charts/oz-agent-worker --set image.tag=test --set worker.workerId=test-worker --set warp.apiKeySecret.create=true --set warp.apiKeySecret.value=dummy --set metrics.enabled=true`
- `git diff --check`

Co-Authored-By: Oz <oz-agent@warp.dev>

Author: vkodithala

internal/types/messages.go

@@ -9,12 +9,13 @@ import (
 type MessageType string
 
 const (
-	MessageTypeTaskAssignment MessageType = "task_assignment"
-	MessageTypeTaskClaimed    MessageType = "task_claimed"
-	MessageTypeTaskCompleted  MessageType = "task_completed"
-	MessageTypeTaskFailed     MessageType = "task_failed"
-	MessageTypeTaskRejected   MessageType = "task_rejected"
-	MessageTypeHeartbeat      MessageType = "heartbeat"
+	MessageTypeTaskAssignment   MessageType = "task_assignment"
+	MessageTypeTaskClaimed      MessageType = "task_claimed"
+	MessageTypeTaskCompleted    MessageType = "task_completed"
+	MessageTypeTaskFailed       MessageType = "task_failed"
+	MessageTypeTaskRejected     MessageType = "task_rejected"
+	MessageTypeTaskCancellation MessageType = "task_cancellation"
+	MessageTypeHeartbeat        MessageType = "heartbeat"
 )
 
 // WebSocketMessage is the base structure for all WebSocket messages
@@ -58,14 +59,16 @@ type TaskClaimedMessage struct {
 
 // TaskCompletedMessage tells the server to end the active run execution after a successful agent process exit.
 type TaskCompletedMessage struct {
-	TaskID  string `json:"task_id"`
-	Message string `json:"message"`
+	TaskID    string     `json:"task_id"`
+	Message   string     `json:"message"`
+	TaskState *TaskState `json:"task_state,omitempty"`
 }
 
 // TaskFailedMessage is sent from worker to server if task launch fails
 type TaskFailedMessage struct {
-	TaskID  string `json:"task_id"`
-	Message string `json:"message"`
+	TaskID    string     `json:"task_id"`
+	Message   string     `json:"message"`
+	TaskState *TaskState `json:"task_state,omitempty"`
 }
 
 // TaskRejectedMessage is sent from worker to server when the worker cannot accept the task
@@ -75,6 +78,18 @@ type TaskRejectedMessage struct {
 	Reason string `json:"reason"`
 }
 
+// TaskCancellationMessage is sent from server to worker to cancel an active task.
+type TaskCancellationMessage struct {
+	TaskID string `json:"task_id"`
+}
+
+// TaskState is the serialized terminal task state accepted by warp-server.
+type TaskState string
+
+const (
+	TaskStateCancelled TaskState = "CANCELLED"
+)
+
 type TaskDefinition struct {
 	Prompt string `json:"prompt"`
 }

internal/worker/direct.go

@@ -202,6 +202,9 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) err
 
 		log.Infof(ctx, "Running setup command: %s", b.config.SetupCommand)
 		if err := b.runCommand(ctx, b.config.SetupCommand, workspaceDir, setupEnv); err != nil {
+			if ctx.Err() != nil {
+				return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonTaskCancelled, ctx.Err())
+			}
 			return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonSetupCommand, fmt.Errorf("setup command failed: %w", err))
 		}
 	}
@@ -232,6 +235,9 @@ func (b *DirectBackend) ExecuteTask(ctx context.Context, params *TaskParams) err
 	log.Debugf(ctx, "Command: %s %s", b.ozPath, strings.Join(params.BaseArgs, " "))
 
 	if err := cmd.Run(); err != nil {
+		if ctx.Err() != nil {
+			return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonTaskCancelled, ctx.Err())
+		}
 		return newBackendFailure(metrics.TaskFailurePhaseBackend, metrics.TaskFailureReasonAgentInvocation, fmt.Errorf("oz agent exited with error: %w", err))
 	}
 

internal/worker/docker.go

@@ -153,7 +153,9 @@ func (b *DockerBackend) ExecuteTask(ctx context.Context, params *TaskParams) err
 
 	defer func() {
 		if containerID != "" && !b.config.NoCleanup {
-			if _, removeErr := dockerClient.ContainerRemove(ctx, containerID, client.ContainerRemoveOptions{Force: true}); removeErr != nil {
+			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), BackendShutdownTimeout)
+			defer cleanupCancel()
+			if _, removeErr := dockerClient.ContainerRemove(cleanupCtx, containerID, client.ContainerRemoveOptions{Force: true}); removeErr != nil {
 				log.Debugf(ctx, "Container %s already removed or removal failed: %v", containerID, removeErr)
 			}
 		}

internal/worker/worker.go

@@ -61,11 +61,23 @@ type Worker struct {
 	reconnectDelay time.Duration
 	lastHeartbeat  time.Time
 	sendChan       chan []byte
-	activeTasks    map[string]context.CancelFunc
+	activeTasks    map[string]activeTask
 	tasksMutex     sync.Mutex
 	backend        Backend
 	taskSemaphore  *semaphore.Weighted // nil when unlimited
 }
+type taskCancellationSource string
+
+const (
+	taskCancellationSourceUser     taskCancellationSource = "user"
+	taskCancellationSourceShutdown taskCancellationSource = "shutdown"
+)
+
+type activeTask struct {
+	ctx                context.Context
+	cancel             context.CancelFunc
+	cancellationSource taskCancellationSource
+}
 
 func New(ctx context.Context, config Config) (*Worker, error) {
 	workerCtx, cancel := context.WithCancel(ctx)
@@ -111,7 +123,7 @@ func New(ctx context.Context, config Config) (*Worker, error) {
 		cancel:         cancel,
 		reconnectDelay: InitialReconnectDelay,
 		sendChan:       make(chan []byte, 256),
-		activeTasks:    make(map[string]context.CancelFunc),
+		activeTasks:    make(map[string]activeTask),
 		backend:        backend,
 		taskSemaphore:  taskSemaphore,
 	}, nil
@@ -324,11 +336,40 @@ func (w *Worker) handleMessage(message []byte) {
 		}
 		w.handleTaskAssignment(&assignment)
 
+	case types.MessageTypeTaskCancellation:
+		var cancellation types.TaskCancellationMessage
+		if err := json.Unmarshal(msg.Data, &cancellation); err != nil {
+			log.Errorf(w.ctx, "Failed to unmarshal task cancellation: %v", err)
+			return
+		}
+		w.handleTaskCancellation(&cancellation)
+
 	default:
 		log.Warnf(w.ctx, "Unknown message type: %s", msg.Type)
 	}
 }
 
+func (w *Worker) handleTaskCancellation(cancellation *types.TaskCancellationMessage) {
+	w.tasksMutex.Lock()
+	task, ok := w.activeTasks[cancellation.TaskID]
+	if ok && task.cancellationSource == "" {
+		task.cancellationSource = taskCancellationSourceUser
+		w.activeTasks[cancellation.TaskID] = task
+	}
+	w.tasksMutex.Unlock()
+	if !ok {
+		log.Warnf(w.ctx, "Received cancellation for inactive task: taskID=%s", cancellation.TaskID)
+		return
+	}
+
+	log.Infof(w.ctx, "Cancelling task from server request: taskID=%s", cancellation.TaskID)
+	metrics.AddTaskEvent(task.ctx, "task.cancellation_requested",
+		attribute.String("source", "server"),
+		attribute.String("task.id", cancellation.TaskID),
+	)
+	task.cancel()
+}
+
 func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) {
 	receivedAt := time.Now()
 	log.Infof(w.ctx, "Received task assignment: taskID=%s, title=%s", assignment.TaskID, assignment.Task.Title)
@@ -381,7 +422,10 @@ func (w *Worker) handleTaskAssignment(assignment *types.TaskAssignmentMessage) {
 	taskCtx, taskCancel := context.WithCancel(executionCtx)
 
 	w.tasksMutex.Lock()
-	w.activeTasks[assignment.TaskID] = taskCancel
+	w.activeTasks[assignment.TaskID] = activeTask{
+		ctx:    taskCtx,
+		cancel: taskCancel,
+	}
 	w.tasksMutex.Unlock()
 	go w.executeTask(taskCtx, taskCancel, span, assignment, receivedAt)
 }
@@ -512,6 +556,19 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc,
 
 	err := w.backend.ExecuteTask(ctx, params)
 	if err != nil {
+		if ctx.Err() == context.Canceled && w.cancellationSource(taskID) == taskCancellationSourceUser {
+			result = metrics.TaskResultCancelled
+			metrics.AddTaskEvent(ctx, "task.cancelled",
+				attribute.String("source", string(taskCancellationSourceUser)),
+			)
+			span.SetStatus(codes.Ok, "task cancelled by user request")
+			log.Infof(ctx, "Task execution cancelled by user request: taskID=%s", taskID)
+			if statusErr := w.sendTaskCancelled(taskID, "Task cancelled by user request."); statusErr != nil {
+				log.Errorf(ctx, "Failed to send task cancelled message: %v", statusErr)
+			}
+			return
+		}
+
 		result = metrics.TaskResultFailed
 		phase, reason := taskFailureLabels(err)
 		metrics.RecordTaskFailure(phase, reason)
@@ -537,6 +594,16 @@ func (w *Worker) executeTask(ctx context.Context, taskCancel context.CancelFunc,
 	}
 }
 
+func (w *Worker) cancellationSource(taskID string) taskCancellationSource {
+	w.tasksMutex.Lock()
+	defer w.tasksMutex.Unlock()
+
+	task, ok := w.activeTasks[taskID]
+	if !ok {
+		return ""
+	}
+	return task.cancellationSource
+}
 func (w *Worker) sendTaskClaimed(taskID string) error {
 	claimed := types.TaskClaimedMessage{
 		TaskID:   taskID,
@@ -561,6 +628,32 @@ func (w *Worker) sendTaskClaimed(taskID string) error {
 	return w.sendMessage(msgBytes)
 }
 
+func (w *Worker) sendTaskCancelled(taskID, message string) error {
+	taskState := types.TaskStateCancelled
+	completedMsg := types.TaskCompletedMessage{
+		TaskID:    taskID,
+		Message:   message,
+		TaskState: &taskState,
+	}
+
+	data, err := json.Marshal(completedMsg)
+	if err != nil {
+		return fmt.Errorf("failed to marshal task cancelled message: %w", err)
+	}
+
+	msg := types.WebSocketMessage{
+		Type: types.MessageTypeTaskCompleted,
+		Data: data,
+	}
+
+	msgBytes, err := json.Marshal(msg)
+	if err != nil {
+		return fmt.Errorf("failed to marshal websocket message: %w", err)
+	}
+
+	return w.sendMessage(msgBytes)
+}
+
 func (w *Worker) sendTaskRejected(taskID, reason string) error {
 	rejectedMsg := types.TaskRejectedMessage{
 		TaskID: taskID,
@@ -654,9 +747,17 @@ func (w *Worker) Shutdown() {
 		log.Infof(w.ctx, "Preserving %d active tasks during worker shutdown", activeTaskCount)
 	} else if activeTaskCount > 0 {
 		log.Infof(w.ctx, "Cancelling %d active tasks", activeTaskCount)
-		for taskID, cancel := range w.activeTasks {
+		for taskID, task := range w.activeTasks {
+			if task.cancellationSource == "" {
+				task.cancellationSource = taskCancellationSourceShutdown
+				w.activeTasks[taskID] = task
+			}
 			log.Debugf(w.ctx, "Cancelling task: %s", taskID)
-			cancel()
+			metrics.AddTaskEvent(task.ctx, "task.cancellation_requested",
+				attribute.String("source", "signal"),
+				attribute.String("task.id", taskID),
+			)
+			task.cancel()
 		}
 	}
 	w.tasksMutex.Unlock()