Skip to content

Commit 9333ef4

Browse files
ysyneuclaude
andcommitted
feat(ws): handle shutdown message for graceful destroy
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>
1 parent c23cd4d commit 9333ef4

2 files changed

Lines changed: 24 additions & 0 deletions

File tree

protocol/messages.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,11 @@ const (
2525
MessageTypeTaskRequest MessageType = "task.request"
2626
MessageTypeTaskCancel MessageType = "task.cancel"
2727
MessageTypeMCPCall MessageType = "mcp.call"
28+
// Shutdown is Safari telling the runner to drop everything and exit —
29+
// used when CloudManager tears down a cloud sandbox so the runner's
30+
// WebSocket closes immediately instead of waiting for the 5s force
31+
// unregister timeout.
32+
MessageTypeShutdown MessageType = "shutdown"
2833
)
2934

3035
// Message is the base WebSocket message structure.

ws/handler.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,8 @@ func (h *Handler) Handle(ctx context.Context, msg *protocol.Message) error {
8080
return h.handleTaskCancel(ctx, msg)
8181
case protocol.MessageTypeMCPCall:
8282
return h.handleMCPCall(ctx, msg)
83+
case protocol.MessageTypeShutdown:
84+
return h.handleShutdown()
8385
default:
8486
slog.Warn("unknown message type",
8587
"type", msg.Type,
@@ -88,6 +90,23 @@ func (h *Handler) Handle(ctx context.Context, msg *protocol.Message) error {
8890
}
8991
}
9092

93+
// handleShutdown cancels running tasks and closes the WebSocket client.
94+
// Triggered by Safari when it's tearing down the sandbox (cloud sandbox
95+
// recycle, eviction, explicit destroy). Close() sets c.closed so
96+
// RunWithReconnect returns nil instead of attempting to reconnect.
97+
func (h *Handler) handleShutdown() error {
98+
slog.Info("shutdown message received, cancelling tasks and closing connection")
99+
h.CancelAllTasks()
100+
// Brief window for in-flight tool results to flush before the socket
101+
// closes. Safari's CloudManager.disconnectRunner waits 5s before forcing
102+
// the unregister — stay well inside that.
103+
_ = h.WaitForTasks(2 * time.Second)
104+
if h.client != nil {
105+
_ = h.client.Close()
106+
}
107+
return nil
108+
}
109+
91110
func (h *Handler) handleTaskRequest(ctx context.Context, msg *protocol.Message) error {
92111
var req protocol.TaskRequestPayload
93112
if err := json.Unmarshal(msg.Payload, &req); err != nil {

0 commit comments

Comments
 (0)