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 {