diff --git a/dot_config/sway/caffeine.sh b/dot_config/sway/caffeine.sh
index cfec65e..2793910 100755
--- a/dot_config/sway/caffeine.sh
+++ b/dot_config/sway/caffeine.sh
@@ -11,10 +11,17 @@
 # drifted out of sync with sway/config — toggling caffeine off would
 # silently change your lock command and timeout values.)
 #
-# start-swayidle.sh also handles env restoration: when invoked from
-# waybar's on-click, WAYLAND_DISPLAY isn't set, and swayidle silently
-# exits without it (which means idle lock is silently dead — the
-# worst failure mode because the bar still shows "caffeine off").
+# At top: source restore-wayland-env.sh to inherit WAYLAND_DISPLAY (and
+# friends) from a running Wayland client. waybar's on-click context
+# strips these, so without this:
+#   - notify-send writes to D-Bus (exits 0) but mako can't display
+#     (no Wayland socket to render on)
+#   - the click appears to silently do nothing
+# This is the SAME bug start-swayidle.sh had for swayidle itself —
+# fixed in both places via the shared helper.
+
+# shellcheck source=restore-wayland-env.sh
+source "$HOME/.config/sway/restore-wayland-env.sh"
 
 FLAG="/tmp/caffeine-inhibit"
 
diff --git a/dot_config/sway/restore-wayland-env.sh b/dot_config/sway/restore-wayland-env.sh
new file mode 100755
index 0000000..4c27b5d
--- /dev/null
+++ b/dot_config/sway/restore-wayland-env.sh
@@ -0,0 +1,34 @@
+#!/bin/bash
+# restore-wayland-env.sh — probe /proc for a running Wayland client and
+# inherit its WAYLAND_DISPLAY / DISPLAY / DBUS_SESSION_BUS_ADDRESS /
+# XDG_RUNTIME_DIR. Used by caffeine.sh and start-swayidle.sh so on-click
+# scripts run from waybar's context (which doesn't inherit the Sway env)
+# can still talk to the compositor.
+#
+# Why probe /proc? waybar's on-click handler forks the script in waybar's
+# own process context. That context has WAYLAND_DISPLAY set (waybar is a
+# Wayland client), but other Sway context (XDG_CURRENT_DESKTOP, sway IPC
+# socket, etc.) isn't passed to children. We need WAYLAND_DISPLAY for
+# swaymsg/grim/wl-copy/notify-send, so the env-restoration here covers the
+# "Wayland client tools" subset. Swaymsg-specific env (I3SOCK / SWAYSOCK)
+# isn't included because swaymsg falls back to discovery via the
+# /run/user/UID/sway-ipc.*.sock path.
+#
+# Idempotent: if WAYLAND_DISPLAY and DISPLAY are already set, does nothing.
+# Source-able: `source "$HOME/.config/sway/restore-wayland-env.sh"`.
+
+if [[ -n "$WAYLAND_DISPLAY" || -n "$DISPLAY" ]]; then
+    return 0 2>/dev/null || exit 0
+fi
+
+# Find a running Wayland client. mako, swaybar, foot, wofi — anything
+# bound to wayland-N. We grab the first one with WAYLAND_DISPLAY set.
+for pid in /proc/[0-9]*; do
+    [[ -r "$pid/environ" ]] || continue
+    candidate=$(tr '\0' '\n' < "$pid/environ" 2>/dev/null | grep -E '^(WAYLAND_DISPLAY|DISPLAY|DBUS_SESSION_BUS_ADDRESS|XDG_RUNTIME_DIR)=' || true)
+    if [[ -n "$candidate" ]]; then
+        # shellcheck disable=SC2086
+        export $candidate
+        return 0 2>/dev/null || exit 0
+    fi
+done
diff --git a/dot_config/sway/start-swayidle.sh b/dot_config/sway/start-swayidle.sh
index 48d914d..d6a0959 100755
--- a/dot_config/sway/start-swayidle.sh
+++ b/dot_config/sway/start-swayidle.sh
@@ -16,8 +16,8 @@
 #   - sway at session start, with the full Wayland/Sway env (correct)
 #   - caffeine.sh from waybar's on-click, with NO Wayland env (broken — see below)
 # When run from caffeine.sh, swayidle needs WAYLAND_DISPLAY to talk to the
-# compositor. We probe /proc for the env of an already-running Wayland client
-# (mako, swaybar, foot, etc.) and inherit WAYLAND_DISPLAY / DISPLAY /
+# compositor. The shared helper restore-wayland-env.sh probes /proc for
+# an already-running Wayland client and inherits WAYLAND_DISPLAY / DISPLAY /
 # DBUS_SESSION_BUS_ADDRESS / XDG_RUNTIME_DIR from it. This is the same
 # trick `dbus-update-activation-environment` uses internally, just per-process.
 #
@@ -26,30 +26,9 @@
 # caffeine.sh, the script is backgrounded (&) and this function returns.
 set -e
 
-# --- 1. Restore Wayland env if missing -----------------------------------
-# The bug we're fixing: clicking the caffeine button runs this script
-# from waybar's on-click context, which doesn't inherit sway's env.
-# swayidle and lock-fancy.sh both call swaymsg/grim which need
-# WAYLAND_DISPLAY — without it, swayidle silently exits and the user
-# has NO idle lock at all (the worst possible failure mode: looks normal
-# in the bar, but auto-lock is silently dead).
-if [[ -z "$WAYLAND_DISPLAY" && -z "$SWAYSOCK" ]]; then
-    # Find any running Wayland client and steal its env. mako (notification
-    # daemon) is always present and always has the right env.
-    for pid in /proc/[0-9]*; do
-        [[ -r "$pid/environ" ]] || continue
-        candidate=$(tr '\0' '\n' < "$pid/environ" 2>/dev/null | grep -E '^(WAYLAND_DISPLAY|DBUS_SESSION_BUS_ADDRESS|XDG_RUNTIME_DIR|DISPLAY)=' || true)
-        if [[ -n "$candidate" ]]; then
-            export $candidate
-            break
-        fi
-    done
-fi
+# shellcheck source=restore-wayland-env.sh
+source "$HOME/.config/sway/restore-wayland-env.sh"
 
-# --- 2. Start swayidle ----------------------------------------------------
-# If invoked from sway's `exec` line (config), swayidle replaces this shell
-# and runs forever. If invoked from caffeine.sh (backgrounded with &), it
-# runs as a child of the parent shell and caffeine.sh returns immediately.
 exec swayidle -w \
     timeout 300 "$HOME/.config/sway/lock-fancy.sh" \
     timeout 600 'swaymsg "output * power off"' \