Commit 6dc2c1b

NgoQuocViet2001  ·  2026-08-16 22:23:45 -0400 EDT
parent 9575587 
feat(attach): apply labels as the session is created

Adds --labels "key=value ..." to attach, in the same form set takes.

A caller that creates a session and then labels it makes two round trips,
and if it dies in between the session is left running with no labels, so
nothing can attribute it afterwards. This closes that window.

Labels are validated before ensureSession, so a rejected label does not
leave a session behind that the caller never asked for. The validation
loop moved out of labelSet into assertLabels; labelSet still runs it.

Closes #200
2 files changed,  +156, -18
M src/completions.zig
M src/main.zig
M src/completions.zig
+2, -0
......@@ -160,6 +160,7 @@ const fish_completions =
160160 \\complete -c zmx -n "__fish_is_nth_token 2; and __fish_seen_subcommand_from c completions" -a 'bash zsh fish nu' -d Shell
161161 \\
162162 \\# Subcommand flags
163+ \\complete -c zmx -n "__fish_seen_subcommand_from a attach" -l labels -d 'Apply "key=value ..." labels as the session is created' -r
163164 \\complete -c zmx -n "__fish_seen_subcommand_from r run" -s d -d 'Detach from the calling terminal; use `wait` to track its status'
164165 \\complete -c zmx -n "__fish_seen_subcommand_from r run" -l fish -d 'Required when the session runs fish shell'
165166 \\complete -c zmx -n "__fish_seen_subcommand_from l list" -l short -d 'Short output'
......@@ -180,6 +181,7 @@ const nu_completions =
180181 \\
181182 \\export extern "zmx attach" [
182183 \\ name: string@"nu-complete zmx sessions"
184+ \\ --labels: string
183185 \\ ...rest: string
184186 \\]
185187 \\
M src/main.zig
+154, -18
......@@ -119,27 +119,37 @@ pub fn main(init: std.process.Init) !void {
119119 defer gpa.free(sesh);
120120 return history(gpa, io, &cfg, sesh, format);
121121 } else if (std.mem.eql(u8, cmd, "attach") or std.mem.eql(u8, cmd, "a")) {
122- const session_name = args.next() orelse "";
123- if (std.mem.eql(u8, session_name, "--help") or std.mem.eql(u8, session_name, "-h")) {
124- return help(io);
122+ var attach_args: std.ArrayList([]const u8) = .empty;
123+ defer attach_args.deinit(gpa);
124+ while (args.next()) |arg| {
125+ try attach_args.append(gpa, arg);
125126 }
126127
127- var command_args: std.ArrayList([]const u8) = .empty;
128- defer command_args.deinit(gpa);
129- while (args.next()) |arg| {
130- try command_args.append(gpa, arg);
128+ const parsed = parseAttachArgs(attach_args.items);
129+ if (parsed.want_help) {
130+ return help(io);
131131 }
132+ if (parsed.missing_labels_value) {
133+ var buf: [4096]u8 = undefined;
134+ var w = std.Io.File.stderr().writer(io, &buf);
135+ w.interface.print("error: --labels requires \"key=value ...\"\n", .{}) catch {};
136+ w.interface.flush() catch {};
137+ std.process.exit(1);
138+ }
139+ // Before ensureSession, so a rejected label does not leave a session
140+ // behind that the caller never asked for.
141+ if (parsed.labels) |kvs| assertLabels(io, kvs);
132142
133143 var command: ?[][]const u8 = null;
134- if (command_args.items.len > 0) {
135- command = command_args.items;
144+ if (parsed.command_start < attach_args.items.len) {
145+ command = attach_args.items[parsed.command_start..];
136146 }
137147
138148 var cwd_buf: [std.fs.max_path_bytes]u8 = undefined;
139149 const cwd_len = std.process.currentPath(io, &cwd_buf) catch 0;
140150 const cwd = cwd_buf[0..cwd_len];
141151
142- const sesh = try socket.getSeshName(gpa, session_name);
152+ const sesh = try socket.getSeshName(gpa, parsed.session_name);
143153 defer gpa.free(sesh);
144154 const socket_path = socket.getSocketPath(gpa, cfg.socket_dir, sesh) catch |err| switch (err) {
145155 error.NameTooLong => return socket.printSessionNameTooLong(io, sesh, cfg.socket_dir),
......@@ -150,7 +160,7 @@ pub fn main(init: std.process.Init) !void {
150160 daemon.setCwd(cwd);
151161 daemon.shell = shell_env;
152162 std.log.info("socket path={s}", .{daemon.socket_path});
153- return attach(gpa, io, &daemon);
163+ return attach(gpa, io, &daemon, parsed.labels);
154164 } else if (std.mem.eql(u8, cmd, "run") or std.mem.eql(u8, cmd, "r")) {
155165 const session_name = args.next() orelse "";
156166 if (std.mem.eql(u8, session_name, "--help") or std.mem.eql(u8, session_name, "-h")) {
......@@ -411,7 +421,7 @@ fn help(io: std.Io) !void {
411421 \\Usage: zmx <command> [args...]
412422 \\
413423 \\Commands:
414- \\ [a]ttach <name> [command...] Attach to session, creating if needed
424+ \\ [a]ttach [--labels kv] <name> [command...] Attach to session, creating if needed
415425 \\ [r]un <name> [-d] [command...] Send command without attaching
416426 \\ [s]end <name> <text...> Send raw input to session PTY
417427 \\ [p]rint <name> <text...> Inject text into session display
......@@ -433,9 +443,14 @@ fn help(io: std.Io) !void {
433443 \\ This will spawn a login $SHELL with a PTY. You can provide a
434444 \\ command instead of creating a shell.
435445 \\
446+ \\ --labels applies labels as the session is created, in the same form
447+ \\ `zmx set` takes. A caller that creates and then labels in two steps
448+ \\ leaves an unlabelled session behind if it dies between them.
449+ \\
436450 \\ Examples:
437451 \\ zmx attach dev
438452 \\ zmx attach dev vim
453+ \\ zmx attach --labels "project=api role=worker" build
439454 \\
440455 \\History:
441456 \\ This should generally be used with `tail` to print the last lines
......@@ -535,7 +550,8 @@ fn help(io: std.Io) !void {
535550 \\ ZMX_DIR Socket directory (priority 1)
536551 \\ XDG_RUNTIME_DIR Socket directory (priority 2)
537552 \\ TMPDIR Socket directory (priority 3)
538- \\ ZMX_SESSION Session name (injected automatically)
553+ \\ ZMX_SESSION Session name (injected automatically; makes attach
554+ \\ switch this session)
539555 \\ ZMX_SESSION_PREFIX Prefix added to all session names
540556 \\ ZMX_DIR_MODE Sets mode for socket and log directories (octal, defaults to 0750)
541557 \\ ZMX_LOG_MODE Sets mode for log files (octal, defaults to 0640)
......@@ -1076,9 +1092,10 @@ fn labelGet(alloc: std.mem.Allocator, io: std.Io, cfg: *Cfg, session_name: []con
10761092 try stdout.interface.flush();
10771093 }
10781094
1079-fn labelSet(alloc: std.mem.Allocator, io: std.Io, cfg: *Cfg, session_name: []const u8, labels: []const u8) !void {
1080- std.log.info("label set session={s}", .{session_name});
1081-
1095+/// Rejects a malformed label set before the caller acts on it. Exits rather
1096+/// than returning, so `attach --labels` can check its labels before a session
1097+/// exists to be left behind.
1098+fn assertLabels(io: std.Io, labels: []const u8) void {
10821099 var kvs = label.LabelIterator.init(labels);
10831100 while (kvs.next()) |kv| {
10841101 label.assertLabel(kv.key, kv.value) catch |err| {
......@@ -1103,6 +1120,12 @@ fn labelSet(alloc: std.mem.Allocator, io: std.Io, cfg: *Cfg, session_name: []con
11031120 std.process.exit(1);
11041121 };
11051122 }
1123+}
1124+
1125+fn labelSet(alloc: std.mem.Allocator, io: std.Io, cfg: *Cfg, session_name: []const u8, labels: []const u8) !void {
1126+ std.log.info("label set session={s}", .{session_name});
1127+
1128+ assertLabels(io, labels);
11061129
11071130 const socket_path = socket.getSocketPath(alloc, cfg.socket_dir, session_name) catch |err| switch (err) {
11081131 error.NameTooLong => return socket.printSessionNameTooLong(io, session_name, cfg.socket_dir),
......@@ -1287,7 +1310,55 @@ fn switchSesh(gpa: std.mem.Allocator, io: std.Io, daemon: *Daemon, current_sesh:
12871310 };
12881311 }
12891312
1290-fn attach(gpa: std.mem.Allocator, io: std.Io, daemon: *Daemon) !void {
1313+const AttachArgs = struct {
1314+ /// Session name, or "" when the caller did not name one.
1315+ session_name: []const u8 = "",
1316+ /// Index of the first word of the session command.
1317+ command_start: usize = 0,
1318+ /// `--labels "k=v ..."`: labels to apply once the session exists, in the
1319+ /// same space-separated form `zmx set` takes.
1320+ labels: ?[]const u8 = null,
1321+ want_help: bool = false,
1322+ /// `--labels` was given with nothing to apply.
1323+ missing_labels_value: bool = false,
1324+};
1325+
1326+/// Parses the arguments that follow `zmx attach`. Flags are only recognized
1327+/// before the session name, so everything after it stays part of the command
1328+/// handed to the session.
1329+fn parseAttachArgs(argv: []const []const u8) AttachArgs {
1330+ const labels_flag = "--labels";
1331+ var parsed: AttachArgs = .{};
1332+ var i: usize = 0;
1333+ while (i < argv.len) : (i += 1) {
1334+ const arg = argv[i];
1335+ if (std.mem.eql(u8, arg, "--help") or std.mem.eql(u8, arg, "-h")) {
1336+ parsed.want_help = true;
1337+ return parsed;
1338+ }
1339+ if (std.mem.startsWith(u8, arg, labels_flag ++ "=")) {
1340+ parsed.labels = arg[labels_flag.len + 1 ..];
1341+ continue;
1342+ }
1343+ if (std.mem.eql(u8, arg, labels_flag)) {
1344+ if (i + 1 >= argv.len) {
1345+ parsed.missing_labels_value = true;
1346+ parsed.command_start = argv.len;
1347+ return parsed;
1348+ }
1349+ i += 1;
1350+ parsed.labels = argv[i];
1351+ continue;
1352+ }
1353+ parsed.session_name = arg;
1354+ i += 1;
1355+ break;
1356+ }
1357+ parsed.command_start = i;
1358+ return parsed;
1359+}
1360+
1361+fn attach(gpa: std.mem.Allocator, io: std.Io, daemon: *Daemon, labels: ?[]const u8) !void {
12911362 const sesh = socket.getSeshNameFromEnv();
12921363 if (sesh.len > 0) {
12931364 return switchSesh(gpa, io, daemon, sesh);
......@@ -1296,6 +1367,14 @@ fn attach(gpa: std.mem.Allocator, io: std.Io, daemon: *Daemon) !void {
12961367 const is_daemon_proc = try daemon.ensureSession(io);
12971368 if (is_daemon_proc) return;
12981369
1370+ // The session exists now, so labels land before the client takes over the
1371+ // terminal. Doing it here rather than in a follow-up `zmx set` keeps a
1372+ // supervisor from leaving an unlabelled session behind if it dies in
1373+ // between the two calls.
1374+ if (labels) |kvs| {
1375+ try labelSet(gpa, io, daemon.cfg, daemon.session_name, kvs);
1376+ }
1377+
12991378 const client_sock = try socket.sessionConnect(daemon.socket_path);
13001379 std.log.info("attached session={s}", .{daemon.session_name});
13011380 // This is typically used with tcsetattr() to modify terminal settings.
......@@ -1374,7 +1453,7 @@ fn attach(gpa: std.mem.Allocator, io: std.Io, daemon: *Daemon) !void {
13741453 std.log.info("switching to new session cwd={s}", .{switch_cwd});
13751454 target_daemon.setCwd(switch_cwd);
13761455 target_daemon.shell = daemon.shell;
1377- return attach(gpa, io, &target_daemon);
1456+ return attach(gpa, io, &target_daemon, null);
13781457 }
13791458 },
13801459 }
......@@ -1615,3 +1694,60 @@ fn run(gpa: std.mem.Allocator, io: std.Io, daemon: *Daemon, detached: bool, comm
16151694 const exit_code = try tail(gpa, fds, detached, true);
16161695 lib_posix.exit(exit_code);
16171696 }
1697+
1698+test "parseAttachArgs reads a bare session name" {
1699+ const parsed = parseAttachArgs(&.{"dev"});
1700+ try std.testing.expectEqualStrings("dev", parsed.session_name);
1701+ try std.testing.expect(!parsed.want_help);
1702+ try std.testing.expectEqual(@as(usize, 1), parsed.command_start);
1703+}
1704+
1705+test "parseAttachArgs keeps the session command intact" {
1706+ const argv: []const []const u8 = &.{ "dev", "vim", "-n" };
1707+ const parsed = parseAttachArgs(argv);
1708+ try std.testing.expectEqualStrings("dev", parsed.session_name);
1709+ // -n after the session name belongs to the command, not to zmx.
1710+ try std.testing.expectEqualSlices([]const u8, argv[1..], argv[parsed.command_start..]);
1711+}
1712+
1713+test "parseAttachArgs reports help before the session name" {
1714+ try std.testing.expect(parseAttachArgs(&.{"--help"}).want_help);
1715+ try std.testing.expect(parseAttachArgs(&.{"-h"}).want_help);
1716+ try std.testing.expect(parseAttachArgs(&.{ "--labels", "a=1", "--help" }).want_help);
1717+ // Once a session name is read, -h belongs to the command.
1718+ try std.testing.expect(!parseAttachArgs(&.{ "dev", "-h" }).want_help);
1719+}
1720+
1721+test "parseAttachArgs reads --labels in both forms" {
1722+ for ([_][]const []const u8{
1723+ &.{ "--labels", "a=1 b=2", "build" },
1724+ &.{ "--labels=a=1 b=2", "build" },
1725+ }) |argv| {
1726+ const parsed = parseAttachArgs(argv);
1727+ try std.testing.expectEqualStrings("a=1 b=2", parsed.labels.?);
1728+ try std.testing.expectEqualStrings("build", parsed.session_name);
1729+ try std.testing.expect(!parsed.missing_labels_value);
1730+ try std.testing.expectEqual(argv.len, parsed.command_start);
1731+ }
1732+}
1733+
1734+test "parseAttachArgs combines --labels with a session command" {
1735+ const argv: []const []const u8 = &.{ "--labels", "a=1", "build", "make", "--labels" };
1736+ const parsed = parseAttachArgs(argv);
1737+ try std.testing.expectEqualStrings("a=1", parsed.labels.?);
1738+ try std.testing.expectEqualStrings("build", parsed.session_name);
1739+ // --labels after the session name belongs to the command.
1740+ try std.testing.expectEqualSlices([]const u8, argv[3..], argv[parsed.command_start..]);
1741+}
1742+
1743+test "parseAttachArgs reports --labels with no value" {
1744+ const parsed = parseAttachArgs(&.{"--labels"});
1745+ try std.testing.expect(parsed.missing_labels_value);
1746+ try std.testing.expectEqual(@as(?[]const u8, null), parsed.labels);
1747+}
1748+
1749+test "parseAttachArgs leaves labels unset when the flag is absent" {
1750+ const parsed = parseAttachArgs(&.{"dev"});
1751+ try std.testing.expectEqual(@as(?[]const u8, null), parsed.labels);
1752+ try std.testing.expect(!parsed.missing_labels_value);
1753+}