pi-multi-agent-skill 0.0.0-stage → 0.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,2054 @@
1
+ #!/usr/bin/env python3
2
+ # multi_agent.py - start pi agents as separate `pi --mode rpc` processes and
3
+ # manage each one as a named job.
4
+ #
5
+ # ../SKILL.md explains how an agent uses this script; ../references/ describes
6
+ # every command, option, status and state file.
7
+ #
8
+ # Stdlib only, Python 3.9+. Invoke as `python3 multi_agent.py ...` rather than
9
+ # relying on the exec bit.
10
+
11
+ import calendar
12
+ import fcntl
13
+ import hashlib
14
+ import json
15
+ import math
16
+ import os
17
+ import re
18
+ import secrets
19
+ import selectors
20
+ import shlex
21
+ import shutil
22
+ import signal
23
+ import socket
24
+ import subprocess
25
+ import sys
26
+ import threading
27
+ import time
28
+ import traceback
29
+
30
+ # Prompts, job state and agent output are UTF-8 whatever the caller's locale.
31
+ for _stream in (sys.stdin, sys.stdout, sys.stderr):
32
+ try:
33
+ _stream.reconfigure(encoding="utf-8", errors="replace")
34
+ except (AttributeError, ValueError, OSError):
35
+ pass
36
+
37
+ SCRIPT_PATH = os.path.abspath(__file__)
38
+ SKILL_DIR = os.path.dirname(os.path.dirname(SCRIPT_PATH))
39
+ BUILTIN_PROFILES_DIR = os.path.join(SKILL_DIR, "profiles")
40
+
41
+ RESULT_DELIMITER = "--- multi-agent ---"
42
+ DEFAULT_PROFILE = "general-purpose"
43
+ NAME_RE = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._-]{0,62}[A-Za-z0-9])?$")
44
+ DIALOG_METHODS = ("select", "confirm", "input", "editor")
45
+ STOP_SIGNALS = tuple(
46
+ getattr(signal, n) for n in ("SIGTERM", "SIGHUP", "SIGINT") if hasattr(signal, n)
47
+ )
48
+ # Removed from the agent's environment so the herdr extension inside the agent
49
+ # does not report state for the caller's pane.
50
+ REMOVED_ENV = ("HERDR_ENV", "HERDR_PANE_ID", "HERDR_SOCKET_PATH")
51
+ NOT_RUNNING = "job is not running; use resume"
52
+
53
+
54
+ def time_scale():
55
+ """The factor for the timing constants below; tests set it to make them short."""
56
+ value = os.environ.get("MULTI_AGENT_TIME_SCALE") or "1"
57
+ try:
58
+ scale = float(value)
59
+ except ValueError:
60
+ scale = 0.0
61
+ if not 0 < scale < float("inf"):
62
+ raise SystemExit(
63
+ f"multi-agent: error: MULTI_AGENT_TIME_SCALE needs a positive number, got: {value}"
64
+ )
65
+ return scale
66
+
67
+
68
+ TIME_SCALE = time_scale()
69
+ BACKGROUND_TIMEOUT_S = 3600
70
+ DEFAULT_MAX_DEPTH = 2
71
+ ACK_TIMEOUT_S = 10 * TIME_SCALE
72
+ SETTLE_TIMEOUT_S = 30 * TIME_SCALE
73
+ PI_EXIT_GRACE_S = 10 * TIME_SCALE
74
+ PI_TERM_GRACE_S = 5 * TIME_SCALE
75
+ # After the runner accepts `kill`, the command waits this long for it to
76
+ # exit: the runner's whole stop sequence for pi, plus a margin.
77
+ RUNNER_STOP_GRACE_S = PI_EXIT_GRACE_S + PI_TERM_GRACE_S + 5 * TIME_SCALE
78
+ # Starters hold the job's lock for a moment before they hand it to the
79
+ # runner, and readers hold it shared for a moment, so a command that wants it
80
+ # exclusively waits this long before it decides that a runner holds it.
81
+ LOCK_TIMEOUT_S = 5 * TIME_SCALE
82
+ # How often `wait` checks whether the job has ended.
83
+ WAIT_POLL_S = 0.5 * TIME_SCALE
84
+ # sun_path holds the socket path and its trailing NUL.
85
+ SUN_PATH_MAX = 104 if sys.platform == "darwin" or "bsd" in sys.platform else 108
86
+ RECENT_S = 24 * 3600
87
+
88
+ EXIT_OK = 0
89
+ EXIT_FAIL = 1
90
+ EXIT_STILL_RUNNING = 2
91
+ EXIT_USAGE = 64
92
+
93
+ MULTI_AGENT_BLOCK = """<multi-agent>
94
+ You are an agent named {name}, started by another agent to work on one
95
+ task. Whoever reads your result sees only your final message, so make it a
96
+ complete, self-contained answer to the task. No person is available to
97
+ answer questions. If a tool call is blocked because it needs approval, do
98
+ not retry it; say in your final message what you could not do and why.
99
+ </multi-agent>
100
+ """
101
+
102
+ USAGE = """\
103
+ Usage:
104
+ multi_agent.py run [options] (PROMPT | -)
105
+ multi_agent.py resume NAME [--background] [--timeout S] (PROMPT | -)
106
+ multi_agent.py wait NAME [--timeout S]
107
+ multi_agent.py result NAME
108
+ multi_agent.py steer NAME (MESSAGE | -)
109
+ multi_agent.py abort NAME
110
+ multi_agent.py kill (NAME | --all)
111
+ multi_agent.py status NAME
112
+ multi_agent.py peek NAME [-n N]
113
+ multi_agent.py list [--all]
114
+ multi_agent.py profiles
115
+ multi_agent.py rm NAME
116
+
117
+ Options for run:
118
+ --profile P Profile name, matched case-insensitively (default general-purpose)
119
+ --name N Job name (default <profile>-<4 hex digits>)
120
+ --background Detach the runner and return immediately
121
+ --timeout S Wall-clock limit in seconds (default 3600 in the background,
122
+ none in the foreground; 0 means none)
123
+ --fork Start from a copy of the caller's conversation ($PI_SESSION_FILE)
124
+ --model M Model pattern passed to `pi --model`
125
+ --thinking L Thinking level passed to `pi --thinking`
126
+ --cwd DIR Working directory for the agent (default: the current directory)
127
+
128
+ A PROMPT or MESSAGE of `-` is read from stdin. See references/commands.md for details."""
129
+
130
+
131
+ class CommandError(Exception):
132
+ def __init__(self, message, code=EXIT_FAIL):
133
+ Exception.__init__(self, message)
134
+ self.code = code
135
+
136
+
137
+ class NotRunning(CommandError):
138
+ def __init__(self):
139
+ CommandError.__init__(self, NOT_RUNNING)
140
+
141
+
142
+ def usage_error(message):
143
+ return CommandError(message, EXIT_USAGE)
144
+
145
+
146
+ def note(message):
147
+ sys.stderr.write(f"multi-agent: {message}\n")
148
+ sys.stderr.flush()
149
+
150
+
151
+ # --- small helpers ---------------------------------------------------------
152
+
153
+
154
+ def utcnow():
155
+ return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
156
+
157
+
158
+ def parse_utc(text):
159
+ try:
160
+ return calendar.timegm(time.strptime(text.strip(), "%Y-%m-%dT%H:%M:%SZ"))
161
+ except (ValueError, AttributeError):
162
+ return None
163
+
164
+
165
+ def read_text(path, default=None):
166
+ try:
167
+ with open(path, encoding="utf-8", errors="replace") as fh:
168
+ return fh.read()
169
+ except OSError:
170
+ return default
171
+
172
+
173
+ def read_json(path, default=None):
174
+ text = read_text(path)
175
+ if text is None:
176
+ return default
177
+ try:
178
+ return json.loads(text)
179
+ except ValueError:
180
+ return default
181
+
182
+
183
+ def read_int(path):
184
+ try:
185
+ return int((read_text(path) or "").strip())
186
+ except ValueError:
187
+ return None
188
+
189
+
190
+ def write_atomic(path, text):
191
+ """Write a state file through a temporary file and a rename."""
192
+ tmp = os.path.join(
193
+ os.path.dirname(path),
194
+ f".{os.path.basename(path)}.tmp-{os.getpid()}-{threading.get_ident()}",
195
+ )
196
+ with open(tmp, "w", encoding="utf-8") as fh:
197
+ fh.write(text)
198
+ os.replace(tmp, path)
199
+
200
+
201
+ def remove(path):
202
+ try:
203
+ os.remove(path)
204
+ except OSError:
205
+ pass
206
+
207
+
208
+ def wait_while(check, timeout):
209
+ """Poll check() until it is false; False if it is still true after timeout seconds."""
210
+ end = time.monotonic() + timeout
211
+ while check():
212
+ if time.monotonic() >= end:
213
+ return False
214
+ time.sleep(0.1)
215
+ return True
216
+
217
+
218
+ def format_age(seconds):
219
+ seconds = max(0, int(seconds))
220
+ if seconds < 60:
221
+ return f"{seconds}s"
222
+ if seconds < 3600:
223
+ return f"{seconds // 60}m"
224
+ if seconds < 86400:
225
+ return f"{seconds // 3600}h"
226
+ return f"{seconds // 86400}d"
227
+
228
+
229
+ def excerpt(text, limit=120):
230
+ text = " ".join(str(text).split())
231
+ return text if len(text) <= limit else text[: limit - 3] + "..."
232
+
233
+
234
+ # --- locations ---------------------------------------------------------------
235
+
236
+
237
+ def logical_cwd():
238
+ # Use $PWD when it names the current directory, so symlinked paths keep
239
+ # the name the caller used and the project slug stays stable.
240
+ pwd = os.environ.get("PWD")
241
+ if pwd and os.path.isabs(pwd):
242
+ try:
243
+ if os.path.samefile(pwd, "."):
244
+ return pwd
245
+ except OSError:
246
+ pass
247
+ return os.getcwd()
248
+
249
+
250
+ def project_slug(path):
251
+ return re.sub(r"[^A-Za-z0-9]", "-", path).lstrip("-")
252
+
253
+
254
+ def state_root():
255
+ base = os.environ.get("XDG_STATE_HOME") or os.path.join(
256
+ os.path.expanduser("~"), ".local", "state"
257
+ )
258
+ return os.path.abspath(os.path.join(base, "pi-multi-agent"))
259
+
260
+
261
+ def pi_agent_dir():
262
+ return os.path.expanduser(
263
+ os.environ.get("PI_CODING_AGENT_DIR") or os.path.join("~", ".pi", "agent")
264
+ )
265
+
266
+
267
+ def jobs_dir(project):
268
+ return os.path.join(state_root(), project_slug(project), "jobs")
269
+
270
+
271
+ def sessions_dir(project):
272
+ return os.path.join(pi_agent_dir(), "multi-agent-sessions", project_slug(project))
273
+
274
+
275
+ def pi_bin():
276
+ return os.environ.get("MULTI_AGENT_PI_BIN") or "pi"
277
+
278
+
279
+ def find_session_file(session_dir, name):
280
+ """pi names a session file `<timestamp>_<id>.jsonl`; the timestamp has no `_`."""
281
+ pattern = re.compile(r"^[^_]+_" + re.escape(name) + r"\.jsonl$")
282
+ try:
283
+ entries = sorted(os.listdir(session_dir))
284
+ except OSError:
285
+ return []
286
+ return [os.path.join(session_dir, e) for e in entries if pattern.match(e)]
287
+
288
+
289
+ # --- profiles ----------------------------------------------------------------
290
+
291
+
292
+ def split_items(text):
293
+ """Split a comma-separated list, keeping commas inside quotes."""
294
+ items, current, quote = [], "", None
295
+ for ch in text:
296
+ if quote:
297
+ current += ch
298
+ if ch == quote:
299
+ quote = None
300
+ elif ch in "\"'":
301
+ quote = ch
302
+ current += ch
303
+ elif ch == ",":
304
+ items.append(current)
305
+ current = ""
306
+ else:
307
+ current += ch
308
+ items.append(current)
309
+ return [item.strip() for item in items if item.strip()]
310
+
311
+
312
+ def parse_scalar(raw):
313
+ value = raw.strip()
314
+ if len(value) >= 2 and value[0] == value[-1] == '"':
315
+ try:
316
+ return json.loads(value)
317
+ except ValueError:
318
+ return value[1:-1]
319
+ if len(value) >= 2 and value[0] == value[-1] == "'":
320
+ return value[1:-1].replace("''", "'")
321
+ if value.startswith("[") and value.endswith("]"):
322
+ return [parse_scalar(item) for item in split_items(value[1:-1])]
323
+ if value.lower() == "true":
324
+ return True
325
+ if value.lower() == "false":
326
+ return False
327
+ return value
328
+
329
+
330
+ def parse_frontmatter(text):
331
+ """Return (fields, body) for a Markdown file with `---` frontmatter.
332
+
333
+ Supports `key: value` lines, quoted strings, true/false and `[a, b]`
334
+ lists. Indented lines (nested YAML) and comments are skipped.
335
+ """
336
+ lines = text.lstrip("").splitlines()
337
+ if not lines or lines[0].strip() != "---":
338
+ return {}, text.strip()
339
+ end = None
340
+ for i in range(1, len(lines)):
341
+ if lines[i].strip() == "---":
342
+ end = i
343
+ break
344
+ if end is None:
345
+ return {}, text.strip()
346
+ fields = {}
347
+ for line in lines[1:end]:
348
+ if not line.strip() or line.lstrip().startswith("#") or line[0] in " \t":
349
+ continue
350
+ key, sep, value = line.partition(":")
351
+ if sep and key.strip():
352
+ fields[key.strip()] = parse_scalar(value)
353
+ return fields, "\n".join(lines[end + 1 :]).strip()
354
+
355
+
356
+ def tool_value(value):
357
+ """Join a tools field into pi's comma-separated form; None when unset or empty."""
358
+ if value is None or value is True or value is False:
359
+ return None
360
+ items = value if isinstance(value, list) else split_items(str(value))
361
+ joined = ",".join(str(item).strip() for item in items if str(item).strip())
362
+ return joined or None
363
+
364
+
365
+ def text_field(fields, key):
366
+ value = fields.get(key)
367
+ if value is None or isinstance(value, (bool, list)):
368
+ return None
369
+ value = str(value).strip()
370
+ return value or None
371
+
372
+
373
+ def make_profile(name, fields, body, source, scope):
374
+ return {
375
+ "name": name,
376
+ "description": text_field(fields, "description") or name,
377
+ "model": text_field(fields, "model"),
378
+ "thinking": text_field(fields, "thinking"),
379
+ "tools": tool_value(fields.get("tools")),
380
+ "disallowed_tools": tool_value(fields.get("disallowed_tools")),
381
+ "prompt_mode": (text_field(fields, "prompt_mode") or "append").lower(),
382
+ "enabled": fields.get("enabled") is not False,
383
+ "body": body,
384
+ "source": source,
385
+ "scope": scope,
386
+ }
387
+
388
+
389
+ def profile_dirs(cwd):
390
+ return [
391
+ ("project", os.path.join(cwd, ".pi", "agents")),
392
+ ("user", os.path.join(pi_agent_dir(), "agents")),
393
+ ("builtin", BUILTIN_PROFILES_DIR),
394
+ ]
395
+
396
+
397
+ def load_profiles(cwd):
398
+ """Map each lowercased profile name to the profile from the highest directory.
399
+
400
+ A profile replaces lower ones of the same name as a whole; a disabled one
401
+ hides them. Disabled entries stay in the map with enabled False.
402
+ """
403
+ found = {}
404
+ for scope, directory in profile_dirs(cwd):
405
+ try:
406
+ entries = sorted(os.listdir(directory))
407
+ except OSError:
408
+ continue
409
+ for entry in entries:
410
+ path = os.path.join(directory, entry)
411
+ if not entry.endswith(".md") or not os.path.isfile(path):
412
+ continue
413
+ text = read_text(path)
414
+ if text is None:
415
+ continue
416
+ fields, body = parse_frontmatter(text)
417
+ name = text_field(fields, "name") or entry[: -len(".md")]
418
+ found.setdefault(
419
+ name.lower(), make_profile(name, fields, body, path, scope)
420
+ )
421
+ return found
422
+
423
+
424
+ def find_profile(cwd, name):
425
+ profiles = load_profiles(cwd)
426
+ profile = profiles.get(name.lower())
427
+ if profile is None or not profile["enabled"]:
428
+ available = sorted(p["name"] for p in profiles.values() if p["enabled"])
429
+ raise CommandError(
430
+ "unknown profile: {} (available: {})".format(
431
+ name, ", ".join(available) or "none"
432
+ )
433
+ )
434
+ if profile["prompt_mode"] not in ("append", "replace"):
435
+ raise CommandError(
436
+ "profile {} in {}: prompt_mode must be append or replace, got {}".format(
437
+ profile["name"], profile["source"], profile["prompt_mode"]
438
+ )
439
+ )
440
+ return profile
441
+
442
+
443
+ def caller_model():
444
+ model = os.environ.get("PI_MODEL") or ""
445
+ provider = os.environ.get("PI_PROVIDER") or ""
446
+ if model and provider:
447
+ return provider + "/" + model
448
+ return model or None
449
+
450
+
451
+ def resolve_model(option, profile):
452
+ return option or profile["model"] or caller_model()
453
+
454
+
455
+ def resolve_thinking(option, profile):
456
+ return option or profile["thinking"] or os.environ.get("PI_REASONING_LEVEL") or None
457
+
458
+
459
+ def profile_args(profile, model, thinking, prompt_file):
460
+ """pi arguments for a profile. In replace mode, writes the profile's prompt to prompt_file.
461
+
462
+ In append mode, the prompt reaches pi through append_prompt() instead.
463
+ """
464
+ args = []
465
+ if model:
466
+ args += ["--model", model]
467
+ if thinking:
468
+ args += ["--thinking", thinking]
469
+ if profile["tools"]:
470
+ args += ["--tools", profile["tools"]]
471
+ if profile["disallowed_tools"]:
472
+ args += ["-xt", profile["disallowed_tools"]]
473
+ if profile["prompt_mode"] == "replace":
474
+ write_atomic(prompt_file, profile["body"] + "\n")
475
+ args += ["--system-prompt", prompt_file, "-nc"]
476
+ return args
477
+
478
+
479
+ def append_prompt(profile, name):
480
+ """Text passed with --append-system-prompt: the profile's prompt in append
481
+ mode, then the multi-agent block."""
482
+ parts = []
483
+ if profile["prompt_mode"] == "append" and profile["body"]:
484
+ parts.append(profile["body"])
485
+ parts.append(MULTI_AGENT_BLOCK.format(name=name).strip())
486
+ return "\n\n".join(parts) + "\n"
487
+
488
+
489
+ # --- job state ---------------------------------------------------------------
490
+
491
+
492
+ def validate_name(name):
493
+ if not NAME_RE.match(name):
494
+ raise usage_error(
495
+ f"invalid job name: {name!r} (use letters, digits, '.', '_' and '-', start and end with a"
496
+ " letter or digit, at most 64 characters)"
497
+ )
498
+
499
+
500
+ def default_name(profile_name):
501
+ base = (
502
+ re.sub(r"[^a-z0-9._-]+", "-", profile_name.lower()).strip("._-")[:59] or "agent"
503
+ )
504
+ return f"{base}-{secrets.token_hex(2)}"
505
+
506
+
507
+ def job_path(job_dir, name):
508
+ return os.path.join(job_dir, name)
509
+
510
+
511
+ def state_dir_of(job_dir):
512
+ """The state root that holds job_dir (`<root>/<project>/jobs/<name>`)."""
513
+ return os.path.dirname(os.path.dirname(os.path.dirname(job_dir)))
514
+
515
+
516
+ def job_lock(job_dir, busy_message):
517
+ """Take the job's lock exclusively and return its fd.
518
+
519
+ The lock is an flock on the file `lock`. The runner holds it for its
520
+ whole life, so the job is running exactly while the lock is held (see
521
+ job_state). A starter takes it, then hands its fd to the runner it
522
+ starts; `rm` takes it to delete the job. Readers hold it shared for a
523
+ moment, so the lock is retried for LOCK_TIMEOUT_S before this fails
524
+ with busy_message.
525
+ """
526
+ name = os.path.basename(job_dir)
527
+ path = job_path(job_dir, "lock")
528
+ try:
529
+ fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o644)
530
+ except FileNotFoundError:
531
+ raise CommandError(f"job {name} was removed")
532
+ try:
533
+ end = time.monotonic() + LOCK_TIMEOUT_S
534
+ while True:
535
+ try:
536
+ fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
537
+ break
538
+ except BlockingIOError:
539
+ if time.monotonic() >= end:
540
+ raise CommandError(busy_message)
541
+ time.sleep(0.05)
542
+ except OSError as exc:
543
+ raise CommandError(
544
+ f"flock does not work on {state_dir_of(job_dir)} ({exc.strerror or exc}); set XDG_STATE_HOME to a local file system"
545
+ )
546
+ # `rm` may have removed the job while this command waited.
547
+ try:
548
+ if os.stat(path).st_ino != os.fstat(fd).st_ino:
549
+ raise CommandError(f"job {name} was removed")
550
+ except FileNotFoundError:
551
+ raise CommandError(f"job {name} was removed")
552
+ except BaseException:
553
+ os.close(fd)
554
+ raise
555
+ return fd
556
+
557
+
558
+ def job_state(job_dir):
559
+ """The job's status: `running` while its lock is held, else the end its runner recorded.
560
+
561
+ The files are read under a shared lock, so no run can start while they
562
+ are read, and the reading is a consistent snapshot. A run without an
563
+ exit code ended without its runner recording a result: `died`.
564
+ """
565
+ try:
566
+ fd = os.open(job_path(job_dir, "lock"), os.O_RDONLY)
567
+ except OSError:
568
+ fd = None
569
+ try:
570
+ if fd is not None:
571
+ try:
572
+ fcntl.flock(fd, fcntl.LOCK_SH | fcntl.LOCK_NB)
573
+ except BlockingIOError:
574
+ return "running"
575
+ except OSError:
576
+ # flock does not work here, so `run` never started a job here.
577
+ pass
578
+ if not os.path.exists(job_path(job_dir, "exit_code")):
579
+ return "died"
580
+ return (read_text(job_path(job_dir, "status"), "") or "").strip() or "unknown"
581
+ finally:
582
+ if fd is not None:
583
+ os.close(fd)
584
+
585
+
586
+ def socket_path(job_dir):
587
+ """The path of the job's control socket.
588
+
589
+ It is computed from the job directory, so every command finds it without
590
+ a file that maps names to sockets. The hash keeps its length fixed:
591
+ the state root plus about 30 bytes.
592
+ """
593
+ digest = hashlib.sha256(os.fsencode(job_dir)).hexdigest()[:16]
594
+ return os.path.join(state_dir_of(job_dir), "sockets", digest + ".sock")
595
+
596
+
597
+ def check_socket_path(job_dir):
598
+ if len(os.fsencode(socket_path(job_dir))) >= SUN_PATH_MAX:
599
+ raise CommandError(
600
+ "the socket path for this job is too long; set XDG_STATE_HOME to a shorter path"
601
+ )
602
+
603
+
604
+ def listen_control(job_dir):
605
+ """Bind and listen on the job's control socket; called with the job's lock held.
606
+
607
+ A leftover socket file belongs to a runner that is gone, since the lock
608
+ is free, so it is replaced.
609
+ """
610
+ check_socket_path(job_dir)
611
+ path = socket_path(job_dir)
612
+ os.makedirs(os.path.dirname(path), mode=0o700, exist_ok=True)
613
+ remove(path)
614
+ server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
615
+ try:
616
+ server.bind(path)
617
+ server.listen(64)
618
+ except OSError as exc:
619
+ server.close()
620
+ raise CommandError(f"could not create the control socket {path}: {exc}")
621
+ return server
622
+
623
+
624
+ def find_job(name):
625
+ """Look a job up in the current project, then in all projects."""
626
+ validate_name(name)
627
+ direct = os.path.join(jobs_dir(logical_cwd()), name)
628
+ if os.path.isdir(direct):
629
+ return direct
630
+ matches = []
631
+ try:
632
+ projects = sorted(os.listdir(state_root()))
633
+ except OSError:
634
+ projects = []
635
+ for project in projects:
636
+ candidate = os.path.join(state_root(), project, "jobs", name)
637
+ if os.path.isdir(candidate):
638
+ matches.append(candidate)
639
+ if len(matches) == 1:
640
+ return matches[0]
641
+ if matches:
642
+ raise CommandError(
643
+ "job {} exists in several projects; run the command from the project's directory:\n {}".format(
644
+ name, "\n ".join(matches)
645
+ )
646
+ )
647
+ raise CommandError(f"unknown job: {name}")
648
+
649
+
650
+ def model_label(job_dir, job):
651
+ used = (read_text(job_path(job_dir, "model"), "") or "").strip()
652
+ return used or job.get("model") or "(pi default)"
653
+
654
+
655
+ def duration_s(job_dir):
656
+ started = parse_utc(read_text(job_path(job_dir, "started_at"), "") or "")
657
+ ended = parse_utc(read_text(job_path(job_dir, "ended_at"), "") or "")
658
+ if started is None:
659
+ return None
660
+ return int((ended if ended is not None else time.time()) - started)
661
+
662
+
663
+ def format_result(job_dir, status):
664
+ job = read_json(job_path(job_dir, "job.json"), {})
665
+ output = (read_text(job_path(job_dir, "output.txt"), "") or "").rstrip("\n")
666
+ duration = duration_s(job_dir)
667
+ return "{}\n\n{}\nname: {}\nstatus: {}\nprofile: {}\nmodel: {}\nduration_s: {}\n".format(
668
+ output if output.strip() else "(no output)",
669
+ RESULT_DELIMITER,
670
+ os.path.basename(job_dir),
671
+ status,
672
+ job.get("profile", "?"),
673
+ model_label(job_dir, job),
674
+ "?" if duration is None else duration,
675
+ )
676
+
677
+
678
+ def print_result(job_dir, status):
679
+ try:
680
+ sys.stdout.write(format_result(job_dir, status))
681
+ sys.stdout.flush()
682
+ except OSError:
683
+ # The caller's stdout may be gone when a signal stopped the run.
684
+ pass
685
+ return EXIT_OK if status == "done" else EXIT_FAIL
686
+
687
+
688
+ def append_progress(job_dir, message):
689
+ line = "{} {}\n".format(time.strftime("%H:%M:%S"), message)
690
+ try:
691
+ with open(job_path(job_dir, "progress.log"), "a", encoding="utf-8") as fh:
692
+ fh.write(line)
693
+ except OSError:
694
+ pass
695
+
696
+
697
+ def record_end(job_dir, status, output=None, detail=None):
698
+ """Record the end of the job's current run; output None keeps output.txt.
699
+
700
+ exit_code is written last: once it exists, readers trust the status
701
+ file and the rest of the run's state.
702
+ """
703
+ if output is not None:
704
+ write_atomic(job_path(job_dir, "output.txt"), output + ("\n" if output else ""))
705
+ write_atomic(job_path(job_dir, "ended_at"), utcnow() + "\n")
706
+ write_atomic(job_path(job_dir, "status"), status + "\n")
707
+ write_atomic(
708
+ job_path(job_dir, "exit_code"),
709
+ f"{EXIT_OK if status == 'done' else EXIT_FAIL}\n",
710
+ )
711
+ message = "run {} ended: {}".format(read_int(job_path(job_dir, "runs")), status)
712
+ append_progress(job_dir, message + (f" ({detail})" if detail else ""))
713
+
714
+
715
+ # --- runner ------------------------------------------------------------------
716
+
717
+
718
+ def child_env(job_dir, job):
719
+ env = dict(os.environ)
720
+ for key in REMOVED_ENV:
721
+ env.pop(key, None)
722
+ env["MULTI_AGENT_JOB"] = job_dir
723
+ env["MULTI_AGENT_DEPTH"] = str(job["depth"])
724
+ env["PI_MINIMAL_PERMISSION_SYSTEM_NONINTERACTIVE"] = "1"
725
+ env["PWD"] = job["cwd"]
726
+ return env
727
+
728
+
729
+ def pi_argv(job_dir, job):
730
+ argv = [pi_bin(), "--mode", "rpc", "--session-dir", job["session_dir"]]
731
+ if job.get("fork") and job.get("run_kind") == "new":
732
+ argv += ["--fork", job["fork"]]
733
+ argv += ["--session-id", job["name"]]
734
+ argv += job["pi_args"]
735
+ argv += ["--append-system-prompt", job_path(job_dir, "append-prompt.md")]
736
+ return argv
737
+
738
+
739
+ def short_args(args):
740
+ if isinstance(args, dict):
741
+ for key in ("command", "path", "file_path", "pattern", "query", "url"):
742
+ value = args.get(key)
743
+ if isinstance(value, str) and value:
744
+ return excerpt(value)
745
+ return excerpt(json.dumps(args, ensure_ascii=False, separators=(",", ":")))
746
+
747
+
748
+ class Runner:
749
+ """Owns one pi RPC process for one run of a job.
750
+
751
+ lock_fd holds the job's lock, which the starting command handed over and
752
+ the runner keeps until it exits. server is the job's listening control
753
+ socket. caller_fd is the read end of a pipe whose only write end the
754
+ foreground `run` process holds. End of file on it means that process has
755
+ ended, for any reason, SIGKILL included, and the run stops.
756
+ """
757
+
758
+ def __init__(self, job_dir, lock_fd, server, caller_fd=None):
759
+ self.job_dir = job_dir
760
+ self.job = read_json(job_path(job_dir, "job.json"))
761
+ self.prompt = read_text(job_path(job_dir, "prompt.txt"), "")
762
+ self.lock_fd = lock_fd
763
+ self.server = server
764
+ # Accepted connections whose request line is not complete yet.
765
+ self.requests = {}
766
+ self.caller_fd = caller_fd
767
+ self.caller_gone = False
768
+ self.proc = None
769
+ self.out_fd = None
770
+ self.buffer = b""
771
+ self.eof = False
772
+ self.stdin_open = False
773
+ self.selector = selectors.DefaultSelector()
774
+ self.wake_r = None
775
+ self.signum = None
776
+ self.serial = 0
777
+ self.prompt_response = None
778
+ # RPC id of a forwarded steer -> the connection waiting for pi's answer.
779
+ self.pending_acks = {}
780
+ self.kill_requested = False
781
+ self.settled = False
782
+ self.last_stop = None
783
+ self.last_error = None
784
+ # Text of this run's latest assistant message. pi's
785
+ # get_last_assistant_text would also return an earlier run's text,
786
+ # or the caller's after --fork.
787
+ self.last_text = None
788
+ self.model_used = None
789
+ self.usage = read_json(job_path(job_dir, "usage.json"), None) or {}
790
+ self.abort_sent = False
791
+ self.abort_requested = False
792
+ self.timing_out = False
793
+ self.stop_deadline = None
794
+ self.finishing = False
795
+
796
+ # -- process and pipes --
797
+
798
+ def install_signals(self):
799
+ wake_r, wake_w = os.pipe()
800
+ os.set_blocking(wake_r, False)
801
+ os.set_blocking(wake_w, False)
802
+ self.wake_r = wake_r
803
+ signal.set_wakeup_fd(wake_w, warn_on_full_buffer=False)
804
+ self.selector.register(wake_r, selectors.EVENT_READ, "wake")
805
+
806
+ def on_signal(signum, _frame):
807
+ if self.signum is None:
808
+ self.signum = signum
809
+
810
+ for signum in STOP_SIGNALS:
811
+ # An ignored signal stays ignored, so `nohup` keeps working.
812
+ if signal.getsignal(signum) is not signal.SIG_IGN:
813
+ signal.signal(signum, on_signal)
814
+
815
+ def watch_caller(self):
816
+ if self.caller_fd is None:
817
+ return
818
+ os.set_inheritable(self.caller_fd, False)
819
+ os.set_blocking(self.caller_fd, False)
820
+ self.selector.register(self.caller_fd, selectors.EVENT_READ, "caller")
821
+
822
+ def listen(self):
823
+ # Neither fd may reach pi: a pi that closed the lock fd would make
824
+ # the job look gone while it runs.
825
+ os.set_inheritable(self.lock_fd, False)
826
+ self.server.set_inheritable(False)
827
+ self.server.setblocking(False)
828
+ self.selector.register(self.server, selectors.EVENT_READ, "server")
829
+
830
+ def read_caller(self):
831
+ try:
832
+ data = os.read(self.caller_fd, 512)
833
+ except BlockingIOError:
834
+ return
835
+ except OSError:
836
+ data = b""
837
+ if not data:
838
+ self.caller_gone = True
839
+ self.selector.unregister(self.caller_fd)
840
+ os.close(self.caller_fd)
841
+ self.caller_fd = None
842
+
843
+ def start(self):
844
+ with open(job_path(self.job_dir, "stderr.txt"), "wb") as stderr_fh:
845
+ self.proc = subprocess.Popen(
846
+ pi_argv(self.job_dir, self.job),
847
+ cwd=self.job["cwd"],
848
+ env=child_env(self.job_dir, self.job),
849
+ stdin=subprocess.PIPE,
850
+ stdout=subprocess.PIPE,
851
+ stderr=stderr_fh,
852
+ # pi gets its own session, so signals for the runner's process
853
+ # group never reach it; the runner stops it through its stdin.
854
+ start_new_session=True,
855
+ )
856
+ self.stdin_open = True
857
+ self.out_fd = self.proc.stdout.fileno()
858
+ self.selector.register(self.out_fd, selectors.EVENT_READ, "pi")
859
+
860
+ def send(self, obj):
861
+ if not self.stdin_open:
862
+ return False
863
+ try:
864
+ self.proc.stdin.write(json.dumps(obj).encode("utf-8") + b"\n")
865
+ self.proc.stdin.flush()
866
+ return True
867
+ except OSError:
868
+ self.stdin_open = False
869
+ return False
870
+
871
+ def close_stdin(self):
872
+ if self.proc is not None and self.proc.stdin is not None:
873
+ try:
874
+ self.proc.stdin.close()
875
+ except OSError:
876
+ pass
877
+ self.stdin_open = False
878
+
879
+ def new_id(self, kind):
880
+ self.serial += 1
881
+ return f"{kind}-{self.serial}"
882
+
883
+ def pump(self, timeout):
884
+ """Wait up to timeout seconds (None: no limit) and handle what arrives.
885
+
886
+ That is pi's records, control requests, signals and the end of the
887
+ foreground command.
888
+ """
889
+ for key, _events in self.selector.select(
890
+ None if timeout is None else max(0.0, timeout)
891
+ ):
892
+ if key.data == "wake":
893
+ try:
894
+ while os.read(self.wake_r, 512):
895
+ pass
896
+ except OSError:
897
+ pass
898
+ elif key.data == "caller":
899
+ self.read_caller()
900
+ elif key.data == "server":
901
+ self.accept()
902
+ elif key.data == "request":
903
+ self.read_request(key.fileobj)
904
+ else:
905
+ self.read_pi()
906
+
907
+ def read_pi(self):
908
+ try:
909
+ data = os.read(self.out_fd, 65536)
910
+ except OSError:
911
+ data = b""
912
+ if not data:
913
+ self.eof = True
914
+ self.selector.unregister(self.out_fd)
915
+ if self.buffer:
916
+ self.handle_line(self.buffer)
917
+ self.buffer = b""
918
+ return
919
+ self.buffer += data
920
+ # Records end at LF only. U+2028 and U+2029 are valid inside JSON strings.
921
+ while b"\n" in self.buffer:
922
+ line, self.buffer = self.buffer.split(b"\n", 1)
923
+ self.handle_line(line)
924
+
925
+ def handle_line(self, line):
926
+ line = line.rstrip(b"\r")
927
+ if not line.strip():
928
+ return
929
+ try:
930
+ record = json.loads(line.decode("utf-8", "replace"))
931
+ except ValueError:
932
+ return
933
+ if isinstance(record, dict):
934
+ self.handle(record)
935
+
936
+ def handle(self, record):
937
+ kind = record.get("type")
938
+ if kind == "response":
939
+ rid = record.get("id")
940
+ if rid == "prompt":
941
+ self.prompt_response = record
942
+ elif rid in self.pending_acks:
943
+ conn = self.pending_acks.pop(rid)
944
+ if record.get("success"):
945
+ disposition = (record.get("data") or {}).get(
946
+ "disposition"
947
+ ) or "queued"
948
+ self.reply(conn, {"ok": True, "disposition": disposition})
949
+ else:
950
+ self.reply(
951
+ conn,
952
+ {"ok": False, "error": record.get("error") or "steer failed"},
953
+ )
954
+ elif kind == "tool_execution_start":
955
+ self.progress(
956
+ "{} {}".format(
957
+ record.get("toolName", "?"), short_args(record.get("args"))
958
+ )
959
+ )
960
+ elif kind == "message_end":
961
+ message = record.get("message") or {}
962
+ if message.get("role") == "assistant":
963
+ self.on_assistant_message(message)
964
+ elif kind == "auto_retry_start":
965
+ self.progress(
966
+ "retry {}/{}: {}".format(
967
+ record.get("attempt"),
968
+ record.get("maxAttempts"),
969
+ excerpt(record.get("errorMessage") or ""),
970
+ )
971
+ )
972
+ elif kind == "compaction_start":
973
+ self.progress("compaction ({})".format(record.get("reason", "?")))
974
+ elif kind == "extension_ui_request" and record.get("method") in DIALOG_METHODS:
975
+ # No person can answer. Without a reply the extension would wait, often forever.
976
+ self.send(
977
+ {
978
+ "type": "extension_ui_response",
979
+ "id": record.get("id"),
980
+ "cancelled": True,
981
+ }
982
+ )
983
+ self.progress(
984
+ "denied: {}".format(
985
+ excerpt(record.get("title") or record.get("method"))
986
+ )
987
+ )
988
+ elif kind == "agent_settled":
989
+ self.settled = True
990
+
991
+ def on_assistant_message(self, message):
992
+ self.last_stop = message.get("stopReason")
993
+ self.last_error = message.get("errorMessage")
994
+ content = message.get("content") or []
995
+ # As in pi's get_last_assistant_text, an aborted message without
996
+ # content keeps the text before it.
997
+ if content or self.last_stop != "aborted":
998
+ parts = (
999
+ content
1000
+ if isinstance(content, list)
1001
+ else [{"type": "text", "text": str(content)}]
1002
+ )
1003
+ text = "".join(
1004
+ p.get("text") or ""
1005
+ for p in parts
1006
+ if isinstance(p, dict) and p.get("type") == "text"
1007
+ )
1008
+ self.last_text = text.strip() or None
1009
+ if message.get("provider") and message.get("model"):
1010
+ self.model_used = "{}/{}".format(message["provider"], message["model"])
1011
+ elif message.get("model"):
1012
+ self.model_used = message["model"]
1013
+ usage = message.get("usage") or {}
1014
+ for key in ("input", "output", "cacheRead", "cacheWrite", "totalTokens"):
1015
+ if isinstance(usage.get(key), (int, float)):
1016
+ self.usage[key] = self.usage.get(key, 0) + usage[key]
1017
+ cost = (usage.get("cost") or {}).get("total")
1018
+ if isinstance(cost, (int, float)):
1019
+ self.usage["cost"] = round(self.usage.get("cost", 0) + cost, 6)
1020
+
1021
+ def progress(self, message):
1022
+ append_progress(self.job_dir, message)
1023
+
1024
+ # -- control socket --
1025
+
1026
+ def accept(self):
1027
+ while True:
1028
+ try:
1029
+ conn, _addr = self.server.accept()
1030
+ except OSError:
1031
+ # BlockingIOError: no more connections are waiting.
1032
+ return
1033
+ conn.setblocking(False)
1034
+ self.requests[conn] = b""
1035
+ self.selector.register(conn, selectors.EVENT_READ, "request")
1036
+
1037
+ def read_request(self, conn):
1038
+ """Read from a connection until its request line is complete, then handle it."""
1039
+ try:
1040
+ data = conn.recv(65536)
1041
+ except BlockingIOError:
1042
+ return
1043
+ except OSError:
1044
+ data = b""
1045
+ if data:
1046
+ self.requests[conn] += data
1047
+ if b"\n" not in self.requests[conn]:
1048
+ return
1049
+ self.selector.unregister(conn)
1050
+ line = self.requests.pop(conn).split(b"\n", 1)[0]
1051
+ if not data and not line:
1052
+ conn.close()
1053
+ return
1054
+ try:
1055
+ request = json.loads(line.decode("utf-8", "replace"))
1056
+ except ValueError:
1057
+ request = None
1058
+ self.handle_request(conn, request)
1059
+
1060
+ def handle_request(self, conn, request):
1061
+ kind = request.get("type") if isinstance(request, dict) else None
1062
+ if self.finishing or self.settled:
1063
+ self.reply(conn, {"ok": False, "error": NOT_RUNNING})
1064
+ elif kind == "steer":
1065
+ message = str(request.get("message", ""))
1066
+ rid = self.new_id("steer")
1067
+ if self.send({"id": rid, "type": "steer", "message": message}):
1068
+ # Answered when pi answers the steer.
1069
+ self.pending_acks[rid] = conn
1070
+ self.progress(f"steer: {excerpt(message)}")
1071
+ else:
1072
+ self.reply(conn, {"ok": False, "error": NOT_RUNNING})
1073
+ elif kind == "abort":
1074
+ self.abort_requested = True
1075
+ self.progress("abort requested")
1076
+ self.send_abort()
1077
+ # Answered at once: pi answers RPC abort only after the agent has
1078
+ # settled, when the run is already ending.
1079
+ self.reply(conn, {"ok": True, "disposition": "requested"})
1080
+ elif kind == "kill":
1081
+ # The run loop stops the run as it does for a signal.
1082
+ self.kill_requested = True
1083
+ self.reply(conn, {"ok": True, "disposition": "stopping"})
1084
+ else:
1085
+ self.reply(conn, {"ok": False, "error": f"unknown request type: {kind}"})
1086
+
1087
+ def reply(self, conn, obj):
1088
+ try:
1089
+ conn.settimeout(1)
1090
+ conn.sendall(json.dumps(obj).encode("utf-8") + b"\n")
1091
+ except OSError:
1092
+ # The command gave up waiting.
1093
+ pass
1094
+ conn.close()
1095
+
1096
+ def close_control(self):
1097
+ """Answer every request still open with NOT_RUNNING, then remove the socket.
1098
+
1099
+ Called once the run's end is recorded. The runner still holds the
1100
+ job's lock, so no new run has bound the socket path, and removing it
1101
+ is safe.
1102
+ """
1103
+ for fileobj in [self.server] + list(self.requests):
1104
+ try:
1105
+ self.selector.unregister(fileobj)
1106
+ except (KeyError, ValueError):
1107
+ pass
1108
+ unread = list(self.requests.items())
1109
+ while True:
1110
+ try:
1111
+ conn, _addr = self.server.accept()
1112
+ except OSError:
1113
+ break
1114
+ unread.append((conn, b""))
1115
+ for conn, data in unread:
1116
+ # Read the request first: closing a socket with unread data makes
1117
+ # the command's read fail with ECONNRESET instead of the answer.
1118
+ try:
1119
+ conn.settimeout(0.5)
1120
+ while b"\n" not in data:
1121
+ chunk = conn.recv(65536)
1122
+ if not chunk:
1123
+ break
1124
+ data += chunk
1125
+ except OSError:
1126
+ pass
1127
+ for conn in list(self.pending_acks.values()) + [conn for conn, _data in unread]:
1128
+ self.reply(conn, {"ok": False, "error": NOT_RUNNING})
1129
+ self.pending_acks = {}
1130
+ self.requests = {}
1131
+ remove(socket_path(self.job_dir))
1132
+ self.server.close()
1133
+
1134
+ def send_abort(self):
1135
+ if self.abort_sent:
1136
+ return
1137
+ self.abort_sent = True
1138
+ # pi continues with queued messages after an abort unless the queue is cleared.
1139
+ self.send({"id": self.new_id("clear_queue"), "type": "clear_queue"})
1140
+ self.send({"id": self.new_id("abort"), "type": "abort"})
1141
+ self.stop_deadline = time.monotonic() + SETTLE_TIMEOUT_S
1142
+
1143
+ # -- the run --
1144
+
1145
+ def run(self):
1146
+ self.install_signals()
1147
+ self.watch_caller()
1148
+ self.listen()
1149
+ try:
1150
+ self.start()
1151
+ except OSError as exc:
1152
+ return self.conclude("failed", f"could not start {pi_bin()}: {exc}")
1153
+ self.progress(
1154
+ "run {} started ({})".format(
1155
+ read_int(job_path(self.job_dir, "runs")), self.job["mode"]
1156
+ )
1157
+ )
1158
+ self.send({"id": "prompt", "type": "prompt", "message": self.prompt})
1159
+
1160
+ timeout = self.job.get("timeout")
1161
+ deadline = time.monotonic() + timeout if timeout else None
1162
+ while True:
1163
+ if self.signum is not None:
1164
+ return self.finish_killed(
1165
+ f"stopped by {signal.Signals(self.signum).name}"
1166
+ )
1167
+ if self.caller_gone:
1168
+ return self.finish_killed("stopped: the caller's command ended")
1169
+ if self.kill_requested:
1170
+ return self.finish_killed("stopped by kill")
1171
+ now = time.monotonic()
1172
+ if deadline is not None and now >= deadline and not self.timing_out:
1173
+ self.timing_out = True
1174
+ self.progress(f"timeout after {format_seconds(timeout)}s; aborting")
1175
+ self.send_abort()
1176
+ if self.stop_deadline is not None and now >= self.stop_deadline:
1177
+ self.progress("agent did not settle after the abort")
1178
+ return self.finish("timeout" if self.timing_out else "aborted")
1179
+ if self.settled:
1180
+ return self.finish(self.settled_status())
1181
+ response = self.prompt_response
1182
+ if response is not None and not response.get("success"):
1183
+ return self.finish("failed", response.get("error") or "prompt rejected")
1184
+ if (
1185
+ response is not None
1186
+ and (response.get("data") or {}).get("disposition") == "handled"
1187
+ ):
1188
+ return self.finish("done")
1189
+ if self.eof:
1190
+ return self.finish_eof()
1191
+ # Every event wakes the selector, so it waits for the next
1192
+ # deadline only. Once the timeout has fired, only the settle
1193
+ # deadline is ahead.
1194
+ wait = None
1195
+ for limit in (None if self.timing_out else deadline, self.stop_deadline):
1196
+ if limit is not None:
1197
+ wait = limit - now if wait is None else min(wait, limit - now)
1198
+ self.pump(wait)
1199
+
1200
+ def settled_status(self):
1201
+ if self.timing_out:
1202
+ return "timeout"
1203
+ if self.abort_requested or self.last_stop == "aborted":
1204
+ return "aborted"
1205
+ if self.last_stop == "error":
1206
+ return "failed"
1207
+ return "done"
1208
+
1209
+ def finish(self, status, error=None):
1210
+ self.finishing = True
1211
+ self.stop_pi()
1212
+ output = self.last_text or ""
1213
+ if status == "failed":
1214
+ detail = error or self.last_error
1215
+ if detail:
1216
+ output = (output + "\n\n" if output else "") + "error: " + detail
1217
+ return self.conclude(status, output)
1218
+
1219
+ def finish_eof(self):
1220
+ self.finishing = True
1221
+ self.stop_pi()
1222
+ stderr = read_text(job_path(self.job_dir, "stderr.txt"), "") or ""
1223
+ tail = "\n".join(stderr.rstrip("\n").splitlines()[-20:])
1224
+ output = (
1225
+ f"pi exited before the run finished (exit status {self.proc.returncode})."
1226
+ )
1227
+ if tail:
1228
+ output += " Last lines of its stderr:\n\n" + tail
1229
+ return self.conclude("failed", output)
1230
+
1231
+ def finish_killed(self, reason):
1232
+ self.finishing = True
1233
+ self.progress(reason)
1234
+ self.stop_pi()
1235
+ return self.conclude("killed", self.last_text or "")
1236
+
1237
+ def stop_pi(self):
1238
+ """Close pi's stdin, which makes it shut down; escalate to SIGTERM and SIGKILL."""
1239
+ self.close_stdin()
1240
+ if self.proc is None:
1241
+ return
1242
+ for signum, grace in (
1243
+ (None, PI_EXIT_GRACE_S),
1244
+ (signal.SIGTERM, PI_TERM_GRACE_S),
1245
+ (signal.SIGKILL, 0),
1246
+ ):
1247
+ if self.proc.poll() is not None:
1248
+ break
1249
+ if signum is not None:
1250
+ try:
1251
+ self.proc.send_signal(signum)
1252
+ except OSError:
1253
+ pass
1254
+ end = time.monotonic() + grace
1255
+ while self.proc.poll() is None and time.monotonic() < end:
1256
+ # Keep reading: pi stalls if its stdout pipe fills during
1257
+ # shutdown. Requests meanwhile get NOT_RUNNING.
1258
+ self.pump(0.1)
1259
+ self.proc.wait()
1260
+
1261
+ def conclude(self, status, output):
1262
+ self.finishing = True
1263
+ write_atomic(
1264
+ job_path(self.job_dir, "usage.json"),
1265
+ json.dumps(self.usage, indent=2) + "\n",
1266
+ )
1267
+ if self.model_used:
1268
+ write_atomic(job_path(self.job_dir, "model"), self.model_used + "\n")
1269
+ record_end(self.job_dir, status, output)
1270
+ return status
1271
+
1272
+ def emergency_stop(self):
1273
+ """Best-effort cleanup after an unexpected error in the runner."""
1274
+ self.close_stdin()
1275
+ if self.proc is not None and self.proc.poll() is None:
1276
+ try:
1277
+ self.proc.wait(PI_EXIT_GRACE_S)
1278
+ except subprocess.TimeoutExpired:
1279
+ self.proc.kill()
1280
+
1281
+
1282
+ def format_seconds(value):
1283
+ return (f"{value:g}") if value is not None else "none"
1284
+
1285
+
1286
+ def run_runner(job_dir, lock_fd, server, caller_fd=None):
1287
+ """Run one job run to completion and return its final status.
1288
+
1289
+ The job's lock is released when this process exits.
1290
+ """
1291
+ runner = Runner(job_dir, lock_fd, server, caller_fd)
1292
+ try:
1293
+ return runner.run()
1294
+ except Exception: # noqa: BLE001
1295
+ # Record unexpected failures and stop pi before releasing the job's lock.
1296
+ with open(job_path(job_dir, "runner.log"), "a", encoding="utf-8") as fh:
1297
+ fh.write(f"{utcnow()}\n{traceback.format_exc()}\n")
1298
+ runner.emergency_stop()
1299
+ return runner.conclude(
1300
+ "failed",
1301
+ "the runner failed; see {}".format(job_path(job_dir, "runner.log")),
1302
+ )
1303
+ finally:
1304
+ runner.close_control()
1305
+
1306
+
1307
+ # --- starting runs -----------------------------------------------------------
1308
+
1309
+
1310
+ def caller_depth():
1311
+ try:
1312
+ return int(os.environ.get("MULTI_AGENT_DEPTH") or 0)
1313
+ except ValueError:
1314
+ return 0
1315
+
1316
+
1317
+ def check_depth():
1318
+ depth = caller_depth()
1319
+ try:
1320
+ limit = int(os.environ.get("MULTI_AGENT_MAX_DEPTH") or DEFAULT_MAX_DEPTH)
1321
+ except ValueError:
1322
+ limit = DEFAULT_MAX_DEPTH
1323
+ if depth >= limit:
1324
+ raise CommandError(
1325
+ f"cannot start an agent: this caller is at depth {depth} and MULTI_AGENT_MAX_DEPTH is {limit}"
1326
+ )
1327
+ return depth + 1
1328
+
1329
+
1330
+ PREVIOUS_RUN_FILES = ("exit_code", "ended_at", "output.txt", "model")
1331
+
1332
+
1333
+ def start_run(job_dir, lock_fd, server, job, prompt, background, timeout, run_kind):
1334
+ """Record a new run of the job and start its runner.
1335
+
1336
+ lock_fd holds the job's lock and server is its listening control socket.
1337
+ Both pass to the runner, and this process closes its copies, so the
1338
+ runner holds the lock from here on with no gap.
1339
+
1340
+ The runner is always a detached process. In the foreground this process
1341
+ stays as the runner's caller: it waits for the runner, prints the result,
1342
+ and holds the write end of the pipe whose end of file stops the run.
1343
+ """
1344
+ try:
1345
+ proc, caller_w = start_runner(
1346
+ job_dir, lock_fd, server, job, prompt, background, timeout, run_kind
1347
+ )
1348
+ finally:
1349
+ server.close()
1350
+ os.close(lock_fd)
1351
+ if background:
1352
+ sys.stdout.write(
1353
+ "job: {}\nstatus: running\ntimeout_s: {}\n".format(
1354
+ job["name"], format_seconds(timeout)
1355
+ )
1356
+ )
1357
+ return EXIT_OK
1358
+ note("started job {}".format(job["name"]))
1359
+ return follow_runner(job_dir, proc, caller_w)
1360
+
1361
+
1362
+ def start_runner(job_dir, lock_fd, server, job, prompt, background, timeout, run_kind):
1363
+ """Reset the job's state for a new run and start the runner; return it and the caller's pipe end.
1364
+
1365
+ Called with the job's lock held. If the runner cannot start, the run is
1366
+ recorded as failed.
1367
+ """
1368
+ job["mode"] = "background" if background else "foreground"
1369
+ job["timeout"] = timeout
1370
+ job["run_kind"] = run_kind
1371
+ write_atomic(job_path(job_dir, "job.json"), json.dumps(job, indent=2) + "\n")
1372
+ write_atomic(job_path(job_dir, "prompt.txt"), prompt)
1373
+ for name in PREVIOUS_RUN_FILES:
1374
+ remove(job_path(job_dir, name))
1375
+ runs = (read_int(job_path(job_dir, "runs")) or 0) + 1
1376
+ write_atomic(job_path(job_dir, "runs"), f"{runs}\n")
1377
+ write_atomic(job_path(job_dir, "started_at"), utcnow() + "\n")
1378
+ write_atomic(job_path(job_dir, "status"), "running\n")
1379
+
1380
+ # The job directory follows `__runner` and a space, so a person can find
1381
+ # a stuck runner with `pkill -f` (see kill_job).
1382
+ argv = [
1383
+ sys.executable,
1384
+ SCRIPT_PATH,
1385
+ "__runner",
1386
+ job_dir,
1387
+ str(lock_fd),
1388
+ str(server.fileno()),
1389
+ ]
1390
+ pass_fds = [lock_fd, server.fileno()]
1391
+ caller_r = caller_w = None
1392
+ if not background:
1393
+ # os.pipe() descriptors are not inherited, so only the runner gets
1394
+ # the read end and only this process holds the write end.
1395
+ caller_r, caller_w = os.pipe()
1396
+ argv.append(str(caller_r))
1397
+ pass_fds.append(caller_r)
1398
+ try:
1399
+ proc = subprocess.Popen(
1400
+ argv,
1401
+ cwd=job["cwd"],
1402
+ # Its own session, so the caller's process group and terminal never
1403
+ # signal it. The caller's end is the pipe above.
1404
+ start_new_session=True,
1405
+ # No handle on the caller's stdout or stderr, so a command that runs
1406
+ # this script never waits for the runner.
1407
+ stdin=subprocess.DEVNULL,
1408
+ stdout=subprocess.DEVNULL,
1409
+ stderr=subprocess.DEVNULL,
1410
+ pass_fds=pass_fds,
1411
+ )
1412
+ except OSError as exc:
1413
+ for fd in (caller_r, caller_w):
1414
+ if fd is not None:
1415
+ os.close(fd)
1416
+ message = f"could not start the runner: {exc}"
1417
+ record_end(job_dir, "failed", "error: " + message, message)
1418
+ raise CommandError(message)
1419
+ if caller_r is not None:
1420
+ os.close(caller_r)
1421
+ return proc, caller_w
1422
+
1423
+
1424
+ def follow_runner(job_dir, proc, caller_w):
1425
+ """Foreground: wait for the runner to end the run and print the result.
1426
+
1427
+ SIGTERM, SIGHUP or SIGINT closes caller_w, so the runner stops the run and
1428
+ records it, and this process still prints the result. A second signal
1429
+ ends this process at once. Any other end of this process, SIGKILL
1430
+ included, also closes caller_w.
1431
+ """
1432
+ closed = []
1433
+
1434
+ def on_signal(_signum, _frame):
1435
+ if not closed:
1436
+ closed.append(True)
1437
+ os.close(caller_w)
1438
+ for signum in STOP_SIGNALS:
1439
+ if signal.getsignal(signum) is not signal.SIG_IGN:
1440
+ signal.signal(signum, signal.SIG_DFL)
1441
+
1442
+ for signum in STOP_SIGNALS:
1443
+ if signal.getsignal(signum) is not signal.SIG_IGN:
1444
+ signal.signal(signum, on_signal)
1445
+ proc.wait()
1446
+ return print_result(job_dir, job_state(job_dir))
1447
+
1448
+
1449
+ def cmd_runner(argv):
1450
+ """Internal: the detached runner of a run.
1451
+
1452
+ Its arguments are the job directory, the fd that holds the job's lock,
1453
+ the fd of the listening control socket, and, for a foreground run, the
1454
+ read end of the caller's pipe.
1455
+ """
1456
+ job_dir, lock_fd, server_fd = argv[0], int(argv[1]), int(argv[2])
1457
+ caller_fd = int(argv[3]) if len(argv) > 3 else None
1458
+ run_runner(job_dir, lock_fd, socket.socket(fileno=server_fd), caller_fd)
1459
+ return EXIT_OK
1460
+
1461
+
1462
+ # --- argument parsing --------------------------------------------------------
1463
+
1464
+
1465
+ class Args:
1466
+ """A small parser for `--opt value`, `--opt=value`, flags and positionals."""
1467
+
1468
+ def __init__(self, argv, with_value, flags):
1469
+ self.values = {}
1470
+ self.flags = set()
1471
+ self.positional = []
1472
+ i = 0
1473
+ only_positional = False
1474
+ while i < len(argv):
1475
+ arg = argv[i]
1476
+ if only_positional or arg == "-" or not arg.startswith("-"):
1477
+ self.positional.append(arg)
1478
+ elif arg == "--":
1479
+ only_positional = True
1480
+ elif arg in ("-h", "--help"):
1481
+ self.flags.add("help")
1482
+ else:
1483
+ name, eq, value = arg.partition("=")
1484
+ if name in with_value:
1485
+ if not eq:
1486
+ if i + 1 >= len(argv):
1487
+ raise usage_error(f"{name} needs a value")
1488
+ i += 1
1489
+ value = argv[i]
1490
+ self.values[name] = value
1491
+ elif name in flags and not eq:
1492
+ self.flags.add(name)
1493
+ else:
1494
+ raise usage_error(f"unknown option: {arg}")
1495
+ i += 1
1496
+
1497
+
1498
+ def parse_seconds(value, option):
1499
+ try:
1500
+ seconds = float(value)
1501
+ except ValueError:
1502
+ raise usage_error(f"{option} needs a number of seconds, got: {value}")
1503
+ if seconds < 0 or math.isnan(seconds):
1504
+ raise usage_error(f"{option} needs a number of seconds, got: {value}")
1505
+ return seconds
1506
+
1507
+
1508
+ def run_timeout(args, background):
1509
+ if "--timeout" in args.values:
1510
+ seconds = parse_seconds(args.values["--timeout"], "--timeout")
1511
+ return seconds or None
1512
+ return BACKGROUND_TIMEOUT_S if background else None
1513
+
1514
+
1515
+ def read_message(args, what):
1516
+ if len(args.positional) != 1:
1517
+ raise usage_error(f"expected one {what} (a string, or - to read stdin)")
1518
+ text = args.positional[0]
1519
+ if text == "-":
1520
+ text = sys.stdin.read().rstrip("\n")
1521
+ if not text.strip():
1522
+ raise usage_error(f"empty {what}")
1523
+ return text
1524
+
1525
+
1526
+ def one_name(args, command):
1527
+ if len(args.positional) != 1:
1528
+ raise usage_error(f"{command} needs one job name")
1529
+ return args.positional[0]
1530
+
1531
+
1532
+ # --- commands ----------------------------------------------------------------
1533
+
1534
+
1535
+ def cmd_run(argv):
1536
+ args = Args(
1537
+ argv,
1538
+ ("--profile", "--name", "--timeout", "--model", "--thinking", "--cwd"),
1539
+ ("--background", "--fork"),
1540
+ )
1541
+ if "help" in args.flags:
1542
+ sys.stdout.write(USAGE + "\n")
1543
+ return EXIT_OK
1544
+ background = "--background" in args.flags
1545
+ timeout = run_timeout(args, background)
1546
+ project = logical_cwd()
1547
+ cwd = project
1548
+ if "--cwd" in args.values:
1549
+ cwd = os.path.normpath(
1550
+ os.path.join(project, os.path.expanduser(args.values["--cwd"]))
1551
+ )
1552
+ if not os.path.isdir(cwd):
1553
+ raise usage_error(
1554
+ "--cwd is not a directory: {}".format(args.values["--cwd"])
1555
+ )
1556
+ prompt = read_message(args, "prompt")
1557
+ fork = None
1558
+ if "--fork" in args.flags:
1559
+ fork = os.environ.get("PI_SESSION_FILE")
1560
+ if not fork:
1561
+ raise CommandError(
1562
+ "--fork needs PI_SESSION_FILE, which is not set in this environment"
1563
+ )
1564
+ if not os.path.isfile(fork):
1565
+ raise CommandError(f"--fork: PI_SESSION_FILE does not exist: {fork}")
1566
+ depth = check_depth()
1567
+ profile = find_profile(cwd, args.values.get("--profile") or DEFAULT_PROFILE)
1568
+
1569
+ root = jobs_dir(project)
1570
+ # The socket path has the same length for every name in root.
1571
+ check_socket_path(os.path.join(root, "name"))
1572
+ os.makedirs(root, exist_ok=True)
1573
+ name = args.values.get("--name")
1574
+ if name is not None:
1575
+ validate_name(name)
1576
+ try:
1577
+ os.mkdir(os.path.join(root, name))
1578
+ except FileExistsError:
1579
+ raise CommandError(
1580
+ f"job {name} already exists in this project; use `resume {name}` to continue it, pick"
1581
+ f" another --name, or `rm {name}` first"
1582
+ )
1583
+ else:
1584
+ for _attempt in range(100):
1585
+ name = default_name(profile["name"])
1586
+ try:
1587
+ os.mkdir(os.path.join(root, name))
1588
+ break
1589
+ except FileExistsError:
1590
+ continue
1591
+ else:
1592
+ raise CommandError("could not find a free job name; pass --name")
1593
+ job_dir = os.path.join(root, name)
1594
+
1595
+ lock_fd = None
1596
+ try:
1597
+ # Taken at once, so the job directory never exists without a lock.
1598
+ lock_fd = job_lock(job_dir, f"job {name} is busy; try again")
1599
+ model = resolve_model(args.values.get("--model"), profile)
1600
+ thinking = resolve_thinking(args.values.get("--thinking"), profile)
1601
+ pi_args = profile_args(
1602
+ profile, model, thinking, job_path(job_dir, "system-prompt.md")
1603
+ )
1604
+ write_atomic(
1605
+ job_path(job_dir, "append-prompt.md"), append_prompt(profile, name)
1606
+ )
1607
+ session_dir = sessions_dir(project)
1608
+ os.makedirs(session_dir, exist_ok=True)
1609
+ job = {
1610
+ "name": name,
1611
+ "profile": profile["name"],
1612
+ "profile_source": profile["source"],
1613
+ "cwd": cwd,
1614
+ "project": project,
1615
+ "session_dir": session_dir,
1616
+ "parent": os.environ.get("MULTI_AGENT_JOB") or None,
1617
+ "depth": depth,
1618
+ "model": model,
1619
+ "thinking": thinking,
1620
+ "fork": fork,
1621
+ "pi_args": pi_args,
1622
+ "created_at": utcnow(),
1623
+ }
1624
+ server = listen_control(job_dir)
1625
+ except BaseException:
1626
+ shutil.rmtree(job_dir, ignore_errors=True)
1627
+ if lock_fd is not None:
1628
+ os.close(lock_fd)
1629
+ raise
1630
+ return start_run(job_dir, lock_fd, server, job, prompt, background, timeout, "new")
1631
+
1632
+
1633
+ def cmd_resume(argv):
1634
+ args = Args(argv, ("--timeout",), ("--background",))
1635
+ if len(args.positional) < 1:
1636
+ raise usage_error("resume needs a job name and a prompt")
1637
+ name = args.positional.pop(0)
1638
+ background = "--background" in args.flags
1639
+ timeout = run_timeout(args, background)
1640
+ prompt = read_message(args, "prompt")
1641
+ job_dir = find_job(name)
1642
+ check_depth()
1643
+ # A runner holds the lock while the job runs, so taking it proves that
1644
+ # nothing runs, and keeps it so until the new runner holds it.
1645
+ lock_fd = job_lock(
1646
+ job_dir, f"job {name} is running; use steer, or wait for it to finish"
1647
+ )
1648
+ try:
1649
+ job = read_json(job_path(job_dir, "job.json"))
1650
+ if not job:
1651
+ raise CommandError(f"job {name} has no readable job.json")
1652
+ if not os.path.isdir(job["cwd"]):
1653
+ raise CommandError(
1654
+ "the job's working directory no longer exists: {}".format(job["cwd"])
1655
+ )
1656
+ server = listen_control(job_dir)
1657
+ except BaseException:
1658
+ os.close(lock_fd)
1659
+ raise
1660
+ if not find_session_file(job["session_dir"], name):
1661
+ note(f"no session file found for {name}; the agent starts a new conversation")
1662
+ return start_run(
1663
+ job_dir, lock_fd, server, job, prompt, background, timeout, "resume"
1664
+ )
1665
+
1666
+
1667
+ def cmd_wait(argv):
1668
+ args = Args(argv, ("--timeout",), ())
1669
+ job_dir = find_job(one_name(args, "wait"))
1670
+ timeout = (
1671
+ parse_seconds(args.values["--timeout"], "--timeout")
1672
+ if "--timeout" in args.values
1673
+ else None
1674
+ )
1675
+ end = None if timeout is None else time.monotonic() + timeout
1676
+ while True:
1677
+ status = job_state(job_dir)
1678
+ if status != "running":
1679
+ return print_result(job_dir, status)
1680
+ if end is not None and time.monotonic() >= end:
1681
+ lines = (
1682
+ read_text(job_path(job_dir, "progress.log"), "") or ""
1683
+ ).splitlines()
1684
+ sys.stdout.write(
1685
+ f"job: {os.path.basename(job_dir)}\nstatus: running\nwaited_s: {format_seconds(timeout)}\n"
1686
+ )
1687
+ if lines:
1688
+ sys.stdout.write(f"last_progress: {lines[-1]}\n")
1689
+ return EXIT_STILL_RUNNING
1690
+ time.sleep(
1691
+ WAIT_POLL_S
1692
+ if end is None
1693
+ else max(0.0, min(WAIT_POLL_S, end - time.monotonic()))
1694
+ )
1695
+
1696
+
1697
+ def cmd_result(argv):
1698
+ job_dir = find_job(one_name(Args(argv, (), ()), "result"))
1699
+ status = job_state(job_dir)
1700
+ if status == "running":
1701
+ raise CommandError(
1702
+ f"job {os.path.basename(job_dir)} is still running; use wait"
1703
+ )
1704
+ return print_result(job_dir, status)
1705
+
1706
+
1707
+ def connect_runner(job_dir):
1708
+ """Connect to the job's runner; NotRunning when no runner listens.
1709
+
1710
+ A refused or missing socket means no runner listens. While the lock is
1711
+ held anyway, a starter is between taking the lock and binding the
1712
+ socket, or the runner is exiting, so the connection is retried for a
1713
+ moment.
1714
+ """
1715
+ path = socket_path(job_dir)
1716
+ end = time.monotonic() + LOCK_TIMEOUT_S
1717
+ while True:
1718
+ sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
1719
+ try:
1720
+ sock.connect(path)
1721
+ return sock
1722
+ except OSError:
1723
+ sock.close()
1724
+ if job_state(job_dir) != "running":
1725
+ raise NotRunning()
1726
+ if time.monotonic() >= end:
1727
+ raise CommandError(
1728
+ f"the runner of job {os.path.basename(job_dir)} does not accept requests"
1729
+ )
1730
+ time.sleep(0.05)
1731
+
1732
+
1733
+ def send_request(job_dir, request, timeout_message):
1734
+ """Send one request to the job's runner and return its reply.
1735
+
1736
+ Returns None when the runner closed the connection without a reply,
1737
+ which happens only when it exits. Fails with timeout_message when no
1738
+ reply arrives within ACK_TIMEOUT_S.
1739
+ """
1740
+ sock = connect_runner(job_dir)
1741
+ try:
1742
+ sock.settimeout(ACK_TIMEOUT_S)
1743
+ sock.sendall(json.dumps(request).encode("utf-8") + b"\n")
1744
+ line = sock.makefile("rb").readline()
1745
+ except socket.timeout: # noqa: UP041
1746
+ # Python 3.9 uses a separate exception for socket timeouts.
1747
+ raise CommandError(timeout_message)
1748
+ except OSError:
1749
+ line = b""
1750
+ finally:
1751
+ sock.close()
1752
+ try:
1753
+ reply = json.loads(line.decode("utf-8", "replace"))
1754
+ except ValueError:
1755
+ return None
1756
+ return reply if isinstance(reply, dict) else None
1757
+
1758
+
1759
+ def send_control(job_dir, request):
1760
+ name = os.path.basename(job_dir)
1761
+ reply = send_request(
1762
+ job_dir,
1763
+ request,
1764
+ f"the runner of job {name} did not reply within {int(ACK_TIMEOUT_S)} seconds",
1765
+ )
1766
+ if reply is None:
1767
+ if wait_while(lambda: job_state(job_dir) == "running", 2):
1768
+ raise NotRunning()
1769
+ raise CommandError(
1770
+ f"the runner of job {name} closed the connection without a reply"
1771
+ )
1772
+ if not reply.get("ok"):
1773
+ raise CommandError(reply.get("error") or "request failed")
1774
+ return reply
1775
+
1776
+
1777
+ def cmd_steer(argv):
1778
+ args = Args(argv, (), ())
1779
+ if len(args.positional) < 1:
1780
+ raise usage_error("steer needs a job name and a message")
1781
+ name = args.positional.pop(0)
1782
+ message = read_message(args, "message")
1783
+ job_dir = find_job(name)
1784
+ reply = send_control(job_dir, {"type": "steer", "message": message})
1785
+ sys.stdout.write(
1786
+ "job: {}\nsteer: {}\n".format(
1787
+ os.path.basename(job_dir), reply.get("disposition", "queued")
1788
+ )
1789
+ )
1790
+ return EXIT_OK
1791
+
1792
+
1793
+ def cmd_abort(argv):
1794
+ job_dir = find_job(one_name(Args(argv, (), ()), "abort"))
1795
+ reply = send_control(job_dir, {"type": "abort"})
1796
+ sys.stdout.write(
1797
+ "job: {}\nabort: {}\n".format(
1798
+ os.path.basename(job_dir), reply.get("disposition", "requested")
1799
+ )
1800
+ )
1801
+ return EXIT_OK
1802
+
1803
+
1804
+ def kill_job(job_dir):
1805
+ """Ask the job's runner to stop the run, and wait until it has exited.
1806
+
1807
+ The runner stops pi and records `killed` itself; this writes nothing.
1808
+ Whichever run is listening when the request arrives is the run that
1809
+ stops, so a run started after this returns is left alone.
1810
+ """
1811
+ name = os.path.basename(job_dir)
1812
+ # `[_]` keeps the pattern from matching the command line of a shell that
1813
+ # runs this pkill, as `sh -c` and a bash tool do.
1814
+ pattern = "[_]_runner {} ".format(
1815
+ re.sub(r"([.^$*+?()\[\]{}|\\])", r"\\\1", job_dir)
1816
+ )
1817
+ stuck = f"the runner of job {name} is not responding; to stop it, run: pkill -KILL -f {shlex.quote(pattern)}"
1818
+ try:
1819
+ reply = send_request(job_dir, {"type": "kill"}, stuck)
1820
+ except NotRunning:
1821
+ reply = {"ok": False}
1822
+
1823
+ def running():
1824
+ return job_state(job_dir) == "running"
1825
+
1826
+ # A runner that is finishing answers NOT_RUNNING and records its own
1827
+ # status; one that exits before it answers leaves no reply.
1828
+ if not wait_while(running, RUNNER_STOP_GRACE_S):
1829
+ raise CommandError(stuck)
1830
+ if reply is not None and not reply.get("ok"):
1831
+ return f"job: {name}\nstatus: {job_state(job_dir)} (not running)\n"
1832
+ return f"job: {name}\nstatus: {job_state(job_dir)}\n"
1833
+
1834
+
1835
+ def project_jobs(project):
1836
+ root = jobs_dir(project)
1837
+ try:
1838
+ names = sorted(os.listdir(root))
1839
+ except OSError:
1840
+ return []
1841
+ return [
1842
+ os.path.join(root, n) for n in names if os.path.isdir(os.path.join(root, n))
1843
+ ]
1844
+
1845
+
1846
+ def cmd_kill(argv):
1847
+ args = Args(argv, (), ("--all",))
1848
+ if "--all" in args.flags:
1849
+ if args.positional:
1850
+ raise usage_error("kill takes a job name or --all, not both")
1851
+ running = [d for d in project_jobs(logical_cwd()) if job_state(d) == "running"]
1852
+ if not running:
1853
+ sys.stdout.write("no running jobs in this project\n")
1854
+ return EXIT_OK
1855
+ results = {}
1856
+
1857
+ def kill_one(job_dir):
1858
+ try:
1859
+ results[job_dir] = kill_job(job_dir)
1860
+ except CommandError as exc:
1861
+ results[job_dir] = f"job: {os.path.basename(job_dir)}\nerror: {exc}\n"
1862
+
1863
+ threads = [threading.Thread(target=kill_one, args=(d,)) for d in running]
1864
+ for thread in threads:
1865
+ thread.start()
1866
+ for thread in threads:
1867
+ thread.join()
1868
+ sys.stdout.write("".join(results[d] for d in running))
1869
+ return (
1870
+ EXIT_FAIL
1871
+ if any("\nerror: " in out for out in results.values())
1872
+ else EXIT_OK
1873
+ )
1874
+ job_dir = find_job(one_name(args, "kill"))
1875
+ sys.stdout.write(kill_job(job_dir))
1876
+ return EXIT_OK
1877
+
1878
+
1879
+ def cmd_status(argv):
1880
+ job_dir = find_job(one_name(Args(argv, (), ()), "status"))
1881
+ job = read_json(job_path(job_dir, "job.json"), {})
1882
+ status = job_state(job_dir)
1883
+ name = os.path.basename(job_dir)
1884
+ session_dir = job.get("session_dir", "")
1885
+ sessions = find_session_file(session_dir, name)
1886
+
1887
+ def value(file_name):
1888
+ return (read_text(job_path(job_dir, file_name), "") or "").strip() or "-"
1889
+
1890
+ lines = [
1891
+ ("name", name),
1892
+ ("status", status),
1893
+ ("mode", job.get("mode", "-")),
1894
+ ("profile", job.get("profile", "-")),
1895
+ ("model", model_label(job_dir, job)),
1896
+ ("parent", os.path.basename(job["parent"]) if job.get("parent") else "-"),
1897
+ ("cwd", job.get("cwd", "-")),
1898
+ ("started_at", value("started_at")),
1899
+ ("ended_at", value("ended_at")),
1900
+ ("runs", value("runs")),
1901
+ ("exit_code", value("exit_code")),
1902
+ ("output_file", job_path(job_dir, "output.txt")),
1903
+ ("session_file", sessions[-1] if sessions else "-"),
1904
+ (
1905
+ "session_cmd",
1906
+ "cd {} && pi --session-dir {} --session-id {}".format(
1907
+ shlex.quote(job.get("cwd", ".")),
1908
+ shlex.quote(session_dir),
1909
+ shlex.quote(name),
1910
+ ),
1911
+ ),
1912
+ ]
1913
+ if status == "running":
1914
+ lines.append(("note", "open the session in pi only after the job stops"))
1915
+ sys.stdout.write("".join("{}: {}\n".format(*pair) for pair in lines))
1916
+ return EXIT_OK
1917
+
1918
+
1919
+ def cmd_peek(argv):
1920
+ args = Args(argv, ("-n",), ())
1921
+ count = args.values.get("-n", "20")
1922
+ if not count.isdigit():
1923
+ raise usage_error(f"-n needs a whole number of lines, got: {count}")
1924
+ job_dir = find_job(one_name(args, "peek"))
1925
+ lines = (read_text(job_path(job_dir, "progress.log"), "") or "").splitlines()
1926
+ tail = lines[-int(count) :] if int(count) else []
1927
+ sys.stdout.write("".join(line + "\n" for line in tail))
1928
+ return EXIT_OK
1929
+
1930
+
1931
+ def cmd_list(argv):
1932
+ args = Args(argv, (), ("--all",))
1933
+ if args.positional:
1934
+ raise usage_error("list takes no job name")
1935
+ rows = []
1936
+ now = time.time()
1937
+ for job_dir in project_jobs(logical_cwd()):
1938
+ job = read_json(job_path(job_dir, "job.json"), {})
1939
+ status = job_state(job_dir)
1940
+ started = parse_utc(read_text(job_path(job_dir, "started_at"), "") or "")
1941
+ ended = parse_utc(read_text(job_path(job_dir, "ended_at"), "") or "")
1942
+ since = started if status == "running" else (ended or started)
1943
+ if (
1944
+ "--all" not in args.flags
1945
+ and status != "running"
1946
+ and (since is None or now - since > RECENT_S)
1947
+ ):
1948
+ continue
1949
+ rows.append(
1950
+ (
1951
+ started or 0,
1952
+ os.path.basename(job_dir),
1953
+ status,
1954
+ job.get("mode", "-"),
1955
+ os.path.basename(job["parent"]) if job.get("parent") else "-",
1956
+ format_age(now - since) if since is not None else "-",
1957
+ )
1958
+ )
1959
+ if not rows:
1960
+ sys.stdout.write("no jobs in this project\n")
1961
+ return EXIT_OK
1962
+ rows.sort()
1963
+ table = [("NAME", "STATUS", "MODE", "PARENT", "AGE")] + [row[1:] for row in rows]
1964
+ widths = [max(len(row[i]) for row in table) for i in range(4)]
1965
+ for row in table:
1966
+ cells = [row[i].ljust(widths[i]) for i in range(4)] + [row[4]]
1967
+ sys.stdout.write(" ".join(cells) + "\n")
1968
+ return EXIT_OK
1969
+
1970
+
1971
+ def cmd_profiles(argv):
1972
+ args = Args(argv, (), ())
1973
+ if args.positional:
1974
+ raise usage_error("profiles takes no arguments")
1975
+ profiles = sorted(
1976
+ (p for p in load_profiles(logical_cwd()).values() if p["enabled"]),
1977
+ key=lambda p: p["name"].lower(),
1978
+ )
1979
+ for profile in profiles:
1980
+ sys.stdout.write(
1981
+ "{}\n description: {}\n source: {}\n".format(
1982
+ profile["name"],
1983
+ " ".join(profile["description"].split()),
1984
+ profile["source"],
1985
+ )
1986
+ )
1987
+ if not profiles:
1988
+ sys.stdout.write("no profiles found\n")
1989
+ return EXIT_OK
1990
+
1991
+
1992
+ def cmd_rm(argv):
1993
+ job_dir = find_job(one_name(Args(argv, (), ()), "rm"))
1994
+ name = os.path.basename(job_dir)
1995
+ lock_fd = job_lock(job_dir, f"job {name} is running; kill it first")
1996
+ try:
1997
+ job = read_json(job_path(job_dir, "job.json"), {})
1998
+ for path in find_session_file(job.get("session_dir", ""), name):
1999
+ remove(path)
2000
+ remove(socket_path(job_dir))
2001
+ shutil.rmtree(job_dir, ignore_errors=True)
2002
+ finally:
2003
+ os.close(lock_fd)
2004
+ sys.stdout.write(f"removed: {name}\n")
2005
+ return EXIT_OK
2006
+
2007
+
2008
+ COMMANDS = {
2009
+ "run": cmd_run,
2010
+ "resume": cmd_resume,
2011
+ "wait": cmd_wait,
2012
+ "result": cmd_result,
2013
+ "steer": cmd_steer,
2014
+ "abort": cmd_abort,
2015
+ "kill": cmd_kill,
2016
+ "status": cmd_status,
2017
+ "peek": cmd_peek,
2018
+ "list": cmd_list,
2019
+ "profiles": cmd_profiles,
2020
+ "rm": cmd_rm,
2021
+ }
2022
+
2023
+
2024
+ def main(argv):
2025
+ if sys.platform == "win32":
2026
+ note("Windows is not supported")
2027
+ return EXIT_FAIL
2028
+ if not argv:
2029
+ sys.stderr.write(USAGE + "\n")
2030
+ return EXIT_USAGE
2031
+ command, rest = argv[0], argv[1:]
2032
+ if command in ("-h", "--help", "help"):
2033
+ sys.stdout.write(USAGE + "\n")
2034
+ return EXIT_OK
2035
+ if command == "__runner":
2036
+ return cmd_runner(rest)
2037
+ handler = COMMANDS.get(command)
2038
+ if handler is None:
2039
+ note(f"unknown command: {command} (try --help)")
2040
+ return EXIT_USAGE
2041
+ try:
2042
+ return handler(rest)
2043
+ except CommandError as exc:
2044
+ note(f"error: {exc}")
2045
+ return exc.code
2046
+
2047
+
2048
+ if __name__ == "__main__":
2049
+ try:
2050
+ sys.exit(main(sys.argv[1:]))
2051
+ except BrokenPipeError:
2052
+ sys.exit(EXIT_FAIL)
2053
+ except KeyboardInterrupt:
2054
+ sys.exit(130)