From 920e8c725f7d6014a8ab889f5df787cbce85fd73 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Sat, 5 Sep 2026 04:37:19 -0700
Subject: [PATCH] fix(vmm/cli): look where a VMM registers, not only where it
 usually does

A VMM registers itself in `$XDG_RUNTIME_DIR/dstack-vmm`, falling back to
`/run/user/<uid>/dstack-vmm` only when that variable is unset. `vmm ls` and
`vmm switch` read `/run/user/*/dstack-vmm` and nothing else, so a session whose
`XDG_RUNTIME_DIR` points anywhere else has a VMM that is running, reachable,
and invisible to its own CLI -- reported as "No running VMM instances found",
which is the one answer that is certainly wrong.

Scanning `/run/user` stays: it is what lets an operator see every user's
instances, and it covers the usual case where the variable holds exactly that
path. This VMM's own directory is now added by name as well, deduplicated by
`realpath` so the usual case still lists it once.

The diagnostic printed when nothing is found named `/run/user/*/dstack-vmm`
whether or not that is where it looked. It now names the directories it read,
or -- when none exist yet -- the ones it would have.
---
 dstack/vmm/src/vmm-cli.py | 59 +++++++++++++++++++++++++++++++--------
 1 file changed, 48 insertions(+), 11 deletions(-)

diff --git a/dstack/vmm/src/vmm-cli.py b/dstack/vmm/src/vmm-cli.py
index 1a533434f..58bad2243 100755
--- a/dstack/vmm/src/vmm-cli.py
+++ b/dstack/vmm/src/vmm-cli.py
@@ -36,29 +36,68 @@
 
 
 # VMM discovery directories
-# Each user's instances are in $XDG_RUNTIME_DIR/dstack-vmm (typically /run/user/<uid>/dstack-vmm).
-# CLI scans all users' directories so operators can see every instance on the host.
+# A VMM registers in $XDG_RUNTIME_DIR/dstack-vmm, falling back to
+# /run/user/<uid>/dstack-vmm when that variable is unset.
+# Scanning /run/user is what lets an operator see every user's instances, and
+# covers the usual case where the variable holds exactly that path. It does not
+# cover a session that points XDG_RUNTIME_DIR somewhere else, so this VMM's own
+# directory is added by name rather than assumed to be under /run/user.
 def _get_discovery_dirs() -> List[Tuple[str, Optional[str]]]:
     """Return list of (discovery_dir, username) tuples."""
     import pwd
 
     dirs = []
+    seen = set()
+
+    def add(path: str, username: Optional[str]):
+        if not os.path.isdir(path):
+            return
+        key = os.path.realpath(path)
+        if key in seen:
+            return
+        seen.add(key)
+        dirs.append((path, username))
+
     run_user = "/run/user"
     if os.path.isdir(run_user):
         try:
             for uid_str in os.listdir(run_user):
                 candidate = os.path.join(run_user, uid_str, "dstack-vmm")
-                if os.path.isdir(candidate):
-                    try:
-                        username = pwd.getpwuid(int(uid_str)).pw_name
-                    except (KeyError, ValueError):
-                        username = f"uid:{uid_str}"
-                    dirs.append((candidate, username))
+                try:
+                    username = pwd.getpwuid(int(uid_str)).pw_name
+                except (KeyError, ValueError):
+                    username = f"uid:{uid_str}"
+                add(candidate, username)
         except PermissionError:
             pass
+
+    xdg = os.environ.get("XDG_RUNTIME_DIR")
+    if xdg:
+        try:
+            username = pwd.getpwuid(os.getuid()).pw_name
+        except KeyError:
+            username = f"uid:{os.getuid()}"
+        add(os.path.join(xdg, "dstack-vmm"), username)
+
     return dirs
 
 
+def _discovery_dirs_scanned() -> List[str]:
+    """Where a listing looked, for a listing that found nothing.
+
+    An empty result means no directory exists yet, so naming the ones that
+    would have been read beats naming none of them.
+    """
+    found = [d for d, _ in _get_discovery_dirs()]
+    if found:
+        return found
+    candidates = ["/run/user/*/dstack-vmm"]
+    xdg = os.environ.get("XDG_RUNTIME_DIR")
+    if xdg:
+        candidates.append(os.path.join(xdg, "dstack-vmm"))
+    return candidates
+
+
 def load_config() -> Dict[str, Any]:
     """Load configuration from the default config file.
 
@@ -183,9 +222,7 @@ def cmd_ls_vmm(args):
 
     if not instances:
         print("No running VMM instances found.")
-        print(
-            f"  (scanned: {', '.join(d for d, _ in _get_discovery_dirs()) or '/run/user/*/dstack-vmm'})"
-        )
+        print(f"  (scanned: {', '.join(_discovery_dirs_scanned())})")
         return
 
     if getattr(args, "json", False):