From 9333ef4ff323163bdc0b20984df1f19960b3788c Mon Sep 17 00:00:00 2001
From: ysyneu <ysyneu@gmail.com>
Date: Sun, 12 Apr 2026 00:09:00 +0800
Subject: [PATCH] feat(ws): handle shutdown message for graceful destroy
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds a new protocol.MessageTypeShutdown and a matching Handler.handleShutdown
so Safari can cleanly tear down a cloud sandbox's runner without waiting
on the 5s force-unregister timeout.

Flow when CloudManager.disconnectRunner sends a shutdown message:
1. Handler cancels all in-flight tasks via CancelAllTasks.
2. Waits up to 2s for in-flight tool results to flush via WaitForTasks —
   well inside Safari's 5s force-unregister window.
3. Closes the WebSocket client, which sets c.closed so RunWithReconnect
   returns nil instead of entering the reconnect loop.

Before this change, every cloud sandbox destroy hit the 5s timeout because
the runner silently ignored the shutdown message and stayed in its
reconnect loop until the server-side unregister finally fired.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 protocol/messages.go |  5 +++++
 ws/handler.go        | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/protocol/messages.go b/protocol/messages.go
index 44f233a..065a341 100644
--- a/protocol/messages.go
+++ b/protocol/messages.go
@@ -25,6 +25,11 @@ const (
 	MessageTypeTaskRequest MessageType = "task.request"
 	MessageTypeTaskCancel  MessageType = "task.cancel"
 	MessageTypeMCPCall     MessageType = "mcp.call"
+	// Shutdown is Safari telling the runner to drop everything and exit —
+	// used when CloudManager tears down a cloud sandbox so the runner's
+	// WebSocket closes immediately instead of waiting for the 5s force
+	// unregister timeout.
+	MessageTypeShutdown MessageType = "shutdown"
 )
 
 // Message is the base WebSocket message structure.
diff --git a/ws/handler.go b/ws/handler.go
index 95426e5..af8c213 100644
--- a/ws/handler.go
+++ b/ws/handler.go
@@ -80,6 +80,8 @@ func (h *Handler) Handle(ctx context.Context, msg *protocol.Message) error {
 		return h.handleTaskCancel(ctx, msg)
 	case protocol.MessageTypeMCPCall:
 		return h.handleMCPCall(ctx, msg)
+	case protocol.MessageTypeShutdown:
+		return h.handleShutdown()
 	default:
 		slog.Warn("unknown message type",
 			"type", msg.Type,
@@ -88,6 +90,23 @@ func (h *Handler) Handle(ctx context.Context, msg *protocol.Message) error {
 	}
 }
 
+// handleShutdown cancels running tasks and closes the WebSocket client.
+// Triggered by Safari when it's tearing down the sandbox (cloud sandbox
+// recycle, eviction, explicit destroy). Close() sets c.closed so
+// RunWithReconnect returns nil instead of attempting to reconnect.
+func (h *Handler) handleShutdown() error {
+	slog.Info("shutdown message received, cancelling tasks and closing connection")
+	h.CancelAllTasks()
+	// Brief window for in-flight tool results to flush before the socket
+	// closes. Safari's CloudManager.disconnectRunner waits 5s before forcing
+	// the unregister — stay well inside that.
+	_ = h.WaitForTasks(2 * time.Second)
+	if h.client != nil {
+		_ = h.client.Close()
+	}
+	return nil
+}
+
 func (h *Handler) handleTaskRequest(ctx context.Context, msg *protocol.Message) error {
 	var req protocol.TaskRequestPayload
 	if err := json.Unmarshal(msg.Payload, &req); err != nil {