relay-codex 0.0.0-stage → 0.5.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,2288 @@
1
+ #!/usr/bin/env python3
2
+ """Private, experimental Codex CLI account controller. Never log RPC payloads."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import argparse
7
+ import asyncio
8
+ import base64
9
+ import contextlib
10
+ import contextvars
11
+ import hashlib
12
+ import json
13
+ import os
14
+ from pathlib import Path
15
+ import re
16
+ import signal
17
+ import shutil
18
+ import subprocess
19
+ import sys
20
+ import time
21
+ import uuid
22
+
23
+ from pool import Pool, blocked, reset_time, failure_kind, file_lock
24
+ from launch import translate, RESUME_FIELDS
25
+ import settings
26
+
27
+ VERSION = "0.5.0"
28
+ _support = Path.home() / "Library/Application Support"
29
+ _legacy = _support / "Codex Accounts"
30
+ _default_state = _legacy if (_legacy / "accounts").exists() else _support / "Relay"
31
+ STATE = Path(
32
+ os.environ.get("RELAY_STATE")
33
+ or os.environ.get("CODEX_ACCOUNTS_STATE")
34
+ or _default_state
35
+ )
36
+ CONTINUE = (
37
+ "Continue the original task from its saved state after a service interruption. "
38
+ "Review the existing conversation and completed tool results before acting. "
39
+ "Do not repeat completed commands or external actions. Preserve all existing "
40
+ "user requirements and permissions. Verify any uncertain partial operation first."
41
+ )
42
+
43
+
44
+ class RouterError(Exception):
45
+ """A message safe to display; never include a server error or credential."""
46
+
47
+
48
+ class RecoveryTransition(Exception):
49
+ """Return to the recovery driver without growing the Python call stack."""
50
+
51
+ def __init__(self, kind):
52
+ self.kind = kind
53
+
54
+
55
+ _recovery_active = contextvars.ContextVar("relay_recovery_active", default=None)
56
+
57
+
58
+ class RemoteError(RouterError):
59
+ def __init__(self, payload):
60
+ super().__init__(
61
+ "Codex rejected an RPC operation; no account change was assumed."
62
+ )
63
+ self.payload = payload
64
+
65
+
66
+ def private_dir(path: Path) -> Path:
67
+ path.mkdir(parents=True, exist_ok=True, mode=0o700)
68
+ if path.is_symlink():
69
+ raise RouterError("State directories must not be symlinks.")
70
+ path.chmod(0o700)
71
+ return path
72
+
73
+
74
+ def clean_env(home: Path) -> dict:
75
+ env = dict(os.environ)
76
+ # The pool must not accidentally select an environment-provided identity.
77
+ for key in ("OPENAI_API_KEY", "CODEX_API_KEY", "CODEX_ACCESS_TOKEN"):
78
+ env.pop(key, None)
79
+ env["CODEX_HOME"] = str(home)
80
+ return env
81
+
82
+
83
+ def codex_path() -> str:
84
+ path = shutil.which("codex")
85
+ if not path:
86
+ raise RouterError("Codex CLI is not installed or not on PATH.")
87
+ return path
88
+
89
+
90
+ def account_home(state: Path, label: str) -> Path:
91
+ if not re.fullmatch(r"[a-zA-Z0-9][a-zA-Z0-9_-]{0,31}", label):
92
+ raise RouterError(
93
+ "Use a short account label containing letters, numbers, - or _."
94
+ )
95
+ return state / "accounts" / label
96
+
97
+
98
+ def labels(state: Path) -> list[str]:
99
+ root = state / "accounts"
100
+ return (
101
+ sorted(
102
+ p.name for p in root.iterdir() if p.is_dir() and (p / "enrolled").exists()
103
+ )
104
+ if root.exists()
105
+ else []
106
+ )
107
+
108
+
109
+ def read_keychain(home: Path, *, include_identity=False) -> dict:
110
+ """Native Codex direct-keyring layout in 0.166.0. stdout remains in memory."""
111
+ key = "cli|" + hashlib.sha256(str(home.resolve()).encode()).hexdigest()[:16]
112
+ try:
113
+ result = subprocess.run(
114
+ [
115
+ "/usr/bin/security",
116
+ "find-generic-password",
117
+ "-s",
118
+ "Codex Auth",
119
+ "-a",
120
+ key,
121
+ "-w",
122
+ ],
123
+ capture_output=True,
124
+ timeout=20,
125
+ )
126
+ except subprocess.TimeoutExpired:
127
+ raise RouterError(
128
+ "Keychain access timed out. Check for a macOS Keychain prompt."
129
+ ) from None
130
+ if result.returncode:
131
+ raise RouterError(
132
+ "Cannot read this account's Keychain login. Unlock Keychain or enroll the account again."
133
+ )
134
+ try:
135
+ auth = json.loads(result.stdout)
136
+ tokens = auth["tokens"]
137
+ token = tokens["access_token"]
138
+ part = token.split(".")[1]
139
+ claims = json.loads(base64.urlsafe_b64decode(part + "=" * (-len(part) % 4)))
140
+ metadata = claims.get("https://api.openai.com/auth", {})
141
+ account_id = tokens.get("account_id") or metadata["chatgpt_account_id"]
142
+ if not token or not account_id or auth.get("OPENAI_API_KEY"):
143
+ raise ValueError()
144
+ if claims.get("exp", 0) <= time.time() + 10:
145
+ raise RouterError(
146
+ "Account token expired; enroll this account again if automatic refresh fails."
147
+ )
148
+ result = {
149
+ "accessToken": token,
150
+ "chatgptAccountId": account_id,
151
+ "chatgptPlanType": metadata.get("chatgpt_plan_type"),
152
+ }
153
+ if include_identity:
154
+ result["email"] = claims.get("email") or claims.get(
155
+ "https://api.openai.com/profile", {}
156
+ ).get("email")
157
+ return result
158
+ except (KeyError, ValueError, TypeError, IndexError):
159
+ raise RouterError(
160
+ "This Keychain entry is not a usable ChatGPT login."
161
+ ) from None
162
+
163
+
164
+ def limited(snapshot: dict, model: str | None = None) -> bool:
165
+ """Only explicit server flags establish a block; percentages are advisory."""
166
+ return blocked(snapshot, model)
167
+
168
+
169
+ class RPC:
170
+ """Multiplex JSONL RPC. Only safe, fixed error messages leave this object."""
171
+
172
+ def __init__(self, proc, handler=None):
173
+ self.proc, self.handler = proc, handler
174
+ self.pending = {}
175
+ self.sequence = 0
176
+ self.writer_lock = asyncio.Lock()
177
+ self.handlers = set()
178
+ self.reader = asyncio.create_task(self.read_loop())
179
+ self.startup_guard = None
180
+
181
+ @classmethod
182
+ async def start(
183
+ cls, binary: str, home: Path, *, helper=False, handler=None, extra_args=None
184
+ ):
185
+ home.mkdir(parents=True, exist_ok=True, mode=0o700)
186
+ guard = file_lock(home / ".accounts-native-startup.lock")
187
+ await guard.__aenter__()
188
+ mode = "keyring" if helper else "ephemeral"
189
+ cmd = [
190
+ binary,
191
+ "app-server",
192
+ "--stdio",
193
+ ]
194
+ cmd.extend(extra_args or [])
195
+ cmd += ["-c", f'cli_auth_credentials_store="{mode}"']
196
+ if helper:
197
+ cmd += ["-c", "features.secret_auth_storage=false"]
198
+ else:
199
+ # Let failures reach our finite 5/10/15-minute schedule instead of
200
+ # the native client's optional unbounded reconnect loop.
201
+ cmd += ["-c", "features.unbounded_connection_retries=false"]
202
+ try:
203
+ proc = await asyncio.create_subprocess_exec(
204
+ *cmd,
205
+ env=clean_env(home),
206
+ stdin=asyncio.subprocess.PIPE,
207
+ stdout=asyncio.subprocess.PIPE,
208
+ stderr=asyncio.subprocess.DEVNULL,
209
+ limit=32 * 1024 * 1024,
210
+ )
211
+ rpc = cls(proc, handler)
212
+ rpc.startup_guard = guard
213
+ return rpc
214
+ except BaseException:
215
+ await guard.__aexit__(None, None, None)
216
+ raise
217
+
218
+ async def send(self, obj):
219
+ async with self.writer_lock:
220
+ try:
221
+ self.proc.stdin.write(
222
+ json.dumps(obj, separators=(",", ":")).encode() + b"\n"
223
+ )
224
+ await self.proc.stdin.drain()
225
+ except (ConnectionError, OSError):
226
+ raise RouterError(
227
+ "Codex backend disconnected; resume the saved thread after restarting the wrapper."
228
+ ) from None
229
+
230
+ async def call(self, method, params=None, timeout=45):
231
+ if self.reader.done():
232
+ raise RouterError(
233
+ "Codex backend disconnected; resume the saved thread after restarting the wrapper."
234
+ )
235
+ self.sequence += 1
236
+ request_id = f"accounts-{self.sequence}"
237
+ future = asyncio.get_running_loop().create_future()
238
+ self.pending[request_id] = future
239
+ request = {"id": request_id, "method": method}
240
+ if params is not None:
241
+ request["params"] = params
242
+ try:
243
+ await self.send(request)
244
+ return await asyncio.wait_for(future, timeout)
245
+ finally:
246
+ self.pending.pop(request_id, None)
247
+
248
+ async def initialize(self):
249
+ try:
250
+ result = await self.call(
251
+ "initialize",
252
+ {
253
+ "clientInfo": {
254
+ "name": "local_codex_accounts",
255
+ "title": "Relay",
256
+ "version": VERSION,
257
+ },
258
+ "capabilities": {"experimentalApi": True},
259
+ },
260
+ )
261
+ await self.send({"method": "initialized"})
262
+ except BaseException:
263
+ await self.close()
264
+ raise
265
+ await self.release_startup()
266
+ return result
267
+
268
+ async def release_startup(self):
269
+ if self.startup_guard:
270
+ guard, self.startup_guard = self.startup_guard, None
271
+ await guard.__aexit__(None, None, None)
272
+
273
+ async def read_loop(self):
274
+ try:
275
+ while raw := await self.proc.stdout.readline():
276
+ message = json.loads(raw)
277
+ if "method" not in message and message.get("id") in self.pending:
278
+ future = self.pending[message["id"]]
279
+ if not future.done():
280
+ if "error" in message:
281
+ future.set_exception(RemoteError(message["error"]))
282
+ else:
283
+ future.set_result(message.get("result", {}))
284
+ elif self.handler:
285
+ task = asyncio.create_task(self.handler(message))
286
+ self.handlers.add(task)
287
+ task.add_done_callback(self.handlers.discard)
288
+ except (ValueError, OSError, asyncio.LimitOverrunError):
289
+ pass
290
+ finally:
291
+ for future in list(self.pending.values()):
292
+ if not future.done():
293
+ future.set_exception(RouterError("Codex backend disconnected."))
294
+
295
+ async def close(self):
296
+ if self.proc.returncode is None:
297
+ self.proc.terminate()
298
+ try:
299
+ await asyncio.wait_for(self.proc.wait(), 5)
300
+ except asyncio.TimeoutError:
301
+ self.proc.kill()
302
+ await self.proc.wait()
303
+ await self.release_startup()
304
+ self.reader.cancel()
305
+ for task in list(self.handlers):
306
+ task.cancel()
307
+ await asyncio.gather(self.reader, *self.handlers, return_exceptions=True)
308
+
309
+
310
+ class Account:
311
+ def __init__(self, label, home, binary, pool=None):
312
+ self.label, self.home, self.binary = label, home, binary
313
+ self.pool = pool
314
+ self.rpc = None
315
+ self.lock = asyncio.Lock()
316
+
317
+ async def start(self):
318
+ if self.rpc and self.rpc.proc.returncode is not None:
319
+ await self.rpc.close()
320
+ self.rpc = None
321
+ if not self.rpc:
322
+ self.rpc = await RPC.start(self.binary, self.home, helper=True)
323
+ await self.rpc.initialize()
324
+
325
+ @contextlib.asynccontextmanager
326
+ async def operation(self):
327
+ async with self.lock:
328
+ async with file_lock(self.home / "credential.lock"):
329
+ # Fresh process reads latest Keychain state. Never retain a stale
330
+ # refresh token while another terminal rotates the same login.
331
+ try:
332
+ await self.start()
333
+ yield self.rpc
334
+ finally:
335
+ await self.close()
336
+ self.rpc = None
337
+
338
+ async def credentials(self, refresh=False):
339
+ async with self.operation() as rpc:
340
+ response = await rpc.call("account/read", {"refreshToken": refresh})
341
+ if not response.get("account"):
342
+ raise RouterError("Account requires sign-in; enroll it again.")
343
+ return await asyncio.to_thread(read_keychain, self.home)
344
+
345
+ async def usage(self, force=False):
346
+ async with self.lock:
347
+ async with file_lock(self.home / "credential.lock"):
348
+ cached = (
349
+ self.pool.cached(self.label) if self.pool and not force else None
350
+ )
351
+ if cached is None and force and self.pool:
352
+ # Concurrent failures may all request a fresh read. Reuse
353
+ # a just-completed, explicitly blocked snapshot for 1s.
354
+ recent = self.pool.cached(self.label, max_age=1)
355
+ if recent and limited(recent):
356
+ cached = recent
357
+ if cached is not None:
358
+ return cached
359
+ try:
360
+ await self.start()
361
+ snapshot = await self.rpc.call("account/rateLimits/read")
362
+ if self.pool:
363
+ self.pool.save(self.label, snapshot)
364
+ return snapshot
365
+ finally:
366
+ await self.close()
367
+ self.rpc = None
368
+
369
+ async def close(self):
370
+ if self.rpc:
371
+ await self.rpc.close()
372
+
373
+
374
+ class Controller:
375
+ def __init__(
376
+ self,
377
+ accounts,
378
+ backend,
379
+ report=print,
380
+ interval=300,
381
+ *,
382
+ pool=None,
383
+ retry_interval=300,
384
+ retry_attempts=3,
385
+ clock=time.time,
386
+ sleep=asyncio.sleep,
387
+ ):
388
+ self.accounts, self.backend, self.report = accounts, backend, report
389
+ self.interval = interval
390
+ self.index = 0
391
+ self.client = None
392
+ self.active_threads = set()
393
+ self.user_generation = 0
394
+ self.last_turn = {}
395
+ self.seen_failures = set()
396
+ self.recovery = None
397
+ self.login_lock = asyncio.Lock()
398
+ self.client_tasks = set()
399
+ self.server_requests = set()
400
+ self.pending_server_requests = {}
401
+ self.replayed_requests = set()
402
+ self.pool = pool
403
+ self.retry_interval, self.retry_attempts = retry_interval, retry_attempts
404
+ self.clock, self.sleep = clock, sleep
405
+ self.episodes = {}
406
+ self.local_deadlines = {}
407
+ self.thread_settings = {}
408
+ self.status_changed = asyncio.Event()
409
+ self.launch_overrides = {}
410
+ self.additional_roots = []
411
+ self.preferences = (
412
+ settings.load(pool.state) if pool else settings.DEFAULTS.copy()
413
+ )
414
+ self.paused = False
415
+ self.background = False
416
+ self.wake = asyncio.Event()
417
+ self.force_retry = asyncio.Event()
418
+ self.turn_ids = {}
419
+ self.recovery_jobs = {}
420
+ self.pending_intents = {}
421
+ self.active_items = set()
422
+ self.backend_factory = None
423
+ self.restart_lock = asyncio.Lock()
424
+ self.restart_task = None
425
+ self.paused_uncertain = set()
426
+ self.control_task = None
427
+ self.stopping = asyncio.Event()
428
+ self.extra_accounts = None
429
+ self.starting = False
430
+ self.on_thread = None
431
+ self.on_detach = None
432
+ self.project = os.getcwd()
433
+ self.profile = None
434
+ self.network_state = None
435
+ self.awake_process = None
436
+ self.fallback_tried = {}
437
+
438
+ def attached(self):
439
+ return bool(self.client) or self.background or self.starting
440
+
441
+ def persist(self, thread, **values):
442
+ if self.pool:
443
+ previous = self.pool.intent(thread) or {}
444
+ previous.update(values)
445
+ previous["settings"] = self.thread_settings.get(thread, {})
446
+ previous.setdefault(
447
+ "role", "foreground" if thread in self.last_turn else "child"
448
+ )
449
+ self.pool.intent(thread, previous)
450
+
451
+ def note(self, kind, **fields):
452
+ if self.pool:
453
+ self.pool.record(kind, **fields)
454
+ if kind in {"working", "completed", "switched", "backend_restored"}:
455
+ fields.setdefault("next_wake", None)
456
+ self.pool.touch(state=kind, account=self.account.label, **fields)
457
+ if self.preferences.get("notifications"):
458
+ from notifications import notify
459
+
460
+ notify(kind, self.account.label)
461
+
462
+ def merge_settings(self, thread, params):
463
+ allowed = {
464
+ "model",
465
+ "cwd",
466
+ "approvalPolicy",
467
+ "approvalsReviewer",
468
+ "sandbox",
469
+ "sandboxPolicy",
470
+ "permissions",
471
+ "effort",
472
+ "serviceTier",
473
+ "summary",
474
+ "environments",
475
+ "disabledPluginIds",
476
+ "runtimeWorkspaceRoots",
477
+ "modelProvider",
478
+ "personality",
479
+ "baseInstructions",
480
+ "developerInstructions",
481
+ "outputSchema",
482
+ "cyberAccessProgram",
483
+ }
484
+ current = self.thread_settings.setdefault(thread, {})
485
+ params = dict(params)
486
+ if params.get("reasoningEffort") is not None:
487
+ params["effort"] = params["reasoningEffort"]
488
+ if params.get("sandboxPolicy"):
489
+ policy = params["sandboxPolicy"].get("type")
490
+ translated = {
491
+ "dangerFullAccess": "danger-full-access",
492
+ "readOnly": "read-only",
493
+ "workspaceWrite": "workspace-write",
494
+ "danger-full-access": "danger-full-access",
495
+ "read-only": "read-only",
496
+ "workspace-write": "workspace-write",
497
+ }.get(policy)
498
+ if translated:
499
+ params["sandbox"] = translated
500
+ for key in allowed:
501
+ if params.get(key) is not None:
502
+ current[key] = params[key]
503
+ if "permissions" in params and params.get("permissions"):
504
+ current.pop("sandbox", None)
505
+ current.pop("sandboxPolicy", None)
506
+ if params.get("sandbox") or params.get("sandboxPolicy"):
507
+ current.pop("permissions", None)
508
+
509
+ def hold_for_review(self, thread, reason="uncertain_action"):
510
+ self.paused_uncertain.add(thread)
511
+ self.persist(
512
+ thread, review_required=True, review_reason=reason, state="needs_attention"
513
+ )
514
+
515
+ def clear_review(self, thread):
516
+ self.paused_uncertain.discard(thread)
517
+ self.persist(thread, review_required=False, review_reason=None)
518
+
519
+ def can_continue(self, thread, generation):
520
+ saved = self.pool.intent(thread) if self.pool else {}
521
+ return (
522
+ generation == self.user_generation
523
+ and self.attached()
524
+ and thread not in self.paused_uncertain
525
+ and not (saved or {}).get("review_required", False)
526
+ )
527
+
528
+ async def latest(self, thread):
529
+ result = await self.backend.call(
530
+ "thread/turns/list",
531
+ {
532
+ "threadId": thread,
533
+ "limit": 1,
534
+ "sortDirection": "desc",
535
+ "itemsView": "summary",
536
+ },
537
+ timeout=15,
538
+ )
539
+ return (result.get("data") or [None])[0]
540
+
541
+ async def reconcile(self, thread, baseline=None):
542
+ """Read acceptance/status. Never replay a request merely because it timed out."""
543
+ try:
544
+ latest = await self.latest(thread)
545
+ info = await self.backend.call(
546
+ "thread/read", {"threadId": thread}, timeout=15
547
+ )
548
+ except (RouterError, asyncio.TimeoutError):
549
+ return "unknown", None
550
+ status = (info.get("thread") or {}).get("status", {}).get("type")
551
+ if status == "active":
552
+ self.active_threads.add(thread)
553
+ elif status in {"idle", "systemError", "notLoaded"}:
554
+ self.active_threads.discard(thread)
555
+ if baseline == "__unknown_before_submission__" and not self.pending_intents.get(
556
+ thread, {}
557
+ ).get("accepted"):
558
+ return "unknown", latest
559
+ if latest and (baseline is None or latest.get("id") != baseline):
560
+ if latest.get("status") == "inProgress":
561
+ return "running", latest
562
+ if latest.get("status") == "completed":
563
+ return "completed", latest
564
+ if latest.get("status") in {"failed", "interrupted"}:
565
+ return "failed", latest
566
+ return "unknown", latest
567
+
568
+ async def uncertain(self, thread, generation, baseline):
569
+ self.note("reconciling", thread=thread)
570
+ for attempt in range(self.retry_attempts + 1):
571
+ if generation != self.user_generation:
572
+ return
573
+ outcome, turn = await self.reconcile(thread, baseline)
574
+ if outcome in {"running", "completed"}:
575
+ self.clear_review(thread)
576
+ if outcome == "completed":
577
+ self.note("completed", thread=thread)
578
+ return
579
+ if outcome == "failed":
580
+ self.clear_review(thread)
581
+ code = (turn.get("error") or {}).get("codexErrorInfo")
582
+ kind = failure_kind(code, (turn.get("error") or {}).get("message"))
583
+ if kind:
584
+ return await self.recover(thread, generation, kind)
585
+ self.note(
586
+ "needs_attention", thread=thread, reason="accepted_turn_failed"
587
+ )
588
+ return
589
+ if (
590
+ getattr(getattr(self.backend, "proc", None), "returncode", None)
591
+ is not None
592
+ and self.backend_factory
593
+ ):
594
+ await self.restart_backend()
595
+ return
596
+ if attempt < self.retry_attempts:
597
+ await self.wait_until(self.clock() + self.retry_interval, generation)
598
+ self.hold_for_review(thread)
599
+ self.note("needs_attention", thread=thread, reason="acceptance_unknown")
600
+ self.report(
601
+ "Could not prove the request was accepted. It was not replayed. Use retry after reviewing the saved thread."
602
+ )
603
+
604
+ async def wait_until(self, deadline, generation, *, allow_early=True):
605
+ if self.pool:
606
+ self.pool.touch(next_wake=deadline)
607
+ # Injected sleeps keep deterministic test clocks working.
608
+ if self.sleep is not asyncio.sleep:
609
+ return await self.sleep(max(0, deadline - self.clock()))
610
+ while (
611
+ generation == self.user_generation
612
+ and self.attached()
613
+ and self.clock() < deadline
614
+ ):
615
+ self.wake.clear()
616
+ try:
617
+ await asyncio.wait_for(
618
+ self.wake.wait(), min(30, max(0.01, deadline - self.clock()))
619
+ )
620
+ except asyncio.TimeoutError:
621
+ continue
622
+ if allow_early or self.force_retry.is_set():
623
+ self.force_retry.clear()
624
+ break
625
+
626
+ async def maintain(self):
627
+ revision = self.pool.revision() if self.pool else 0
628
+ last_watch = 0
629
+ last_heartbeat = 0
630
+ account_revision = self.pool.revision("accounts") if self.pool else 0
631
+ last_tick = time.time()
632
+ last_network = 0
633
+ while not self.stopping.is_set():
634
+ now = time.time()
635
+ if now - last_tick > 60 or now < last_tick:
636
+ self.wake.set()
637
+ last_tick = now
638
+ if self.pool:
639
+ if now - last_heartbeat > 10:
640
+ self.pool.renew_reservation(self.preferences["reservation_seconds"])
641
+ self.pool.touch(
642
+ threads=list(self.last_turn),
643
+ model=next((self.model_for(t) for t in self.last_turn), None),
644
+ paused=self.paused,
645
+ background=self.background,
646
+ )
647
+ last_heartbeat = now
648
+ new_revision = self.pool.revision()
649
+ if new_revision != revision:
650
+ self.wake.set()
651
+ revision = new_revision
652
+ try:
653
+ self.preferences = settings.load(
654
+ self.pool.state, self.project, self.profile
655
+ )
656
+ self.retry_interval = self.preferences["retry_interval"]
657
+ self.retry_attempts = self.preferences["retry_attempts"]
658
+ self.interval = self.preferences["unknown_reset_interval"]
659
+ except ValueError:
660
+ self.note("needs_attention", reason="invalid_settings")
661
+ changed = self.pool.revision("accounts")
662
+ if changed != account_revision and self.extra_accounts:
663
+ self.accounts = await self.extra_accounts(self.accounts)
664
+ account_revision = changed
665
+ self.wake.set()
666
+ for command, op, params in self.pool.take_commands():
667
+ try:
668
+ await self.control(op, params)
669
+ self.pool.ack(command, "done")
670
+ except Exception:
671
+ self.pool.ack(command, "failed")
672
+ self.note("control_failed", operation=op)
673
+ if self.preferences.get("keep_awake") and self.active_threads:
674
+ if (
675
+ self.awake_process is None
676
+ or self.awake_process.returncode is not None
677
+ ):
678
+ self.awake_process = await asyncio.create_subprocess_exec(
679
+ "/usr/bin/caffeinate",
680
+ "-i",
681
+ "-w",
682
+ str(os.getpid()),
683
+ stdout=asyncio.subprocess.DEVNULL,
684
+ stderr=asyncio.subprocess.DEVNULL,
685
+ )
686
+ elif self.awake_process and self.awake_process.returncode is None:
687
+ self.awake_process.terminate()
688
+ await self.awake_process.wait()
689
+ self.awake_process = None
690
+ if (
691
+ self.network_state is False
692
+ and now - last_network >= self.preferences["offline_retry"]
693
+ and hasattr(self.account, "network")
694
+ ):
695
+ last_network = now
696
+ try:
697
+ self.network_state = await self.account.network()
698
+ if self.network_state:
699
+ self.note("connection_restored")
700
+ if self.recovery and not self.recovery.done():
701
+ self.force_retry.set()
702
+ self.wake.set()
703
+ except (RouterError, asyncio.TimeoutError):
704
+ pass
705
+ if now - last_watch >= self.preferences["watchdog_interval"]:
706
+ last_watch = now
707
+ if self.backend.proc.returncode is not None and self.backend_factory:
708
+ if not self.restart_task or self.restart_task.done():
709
+ self.restart_task = asyncio.create_task(self.restart_backend())
710
+ elif self.active_threads:
711
+ for thread in list(self.active_threads):
712
+ with contextlib.suppress(RouterError, asyncio.TimeoutError):
713
+ await self.reconcile(thread)
714
+ self.status_changed.set()
715
+ try:
716
+ await asyncio.wait_for(self.stopping.wait(), 2)
717
+ except asyncio.TimeoutError:
718
+ pass
719
+
720
+ async def control(self, op, params):
721
+ thread = params.get("thread") or next(reversed(self.last_turn), None)
722
+ if op == "pause":
723
+ self.paused = True
724
+ self.note("paused")
725
+ elif op == "resume":
726
+ self.paused = False
727
+ self.wake.set()
728
+ self.note("resumed")
729
+ elif op == "stop":
730
+ self.stopping.set()
731
+ elif op == "detach":
732
+ if self.pool:
733
+ from service import folder
734
+ from settings import atomic_json
735
+
736
+ file = folder(self.pool.state) / (self.pool.session + ".job.json")
737
+ if file.exists():
738
+ spec = json.loads(file.read_text())
739
+ spec["takeover"] = True
740
+ atomic_json(file, spec)
741
+ if self.on_detach:
742
+ await self.on_detach()
743
+ if self.client:
744
+ await self.client.close(
745
+ code=1000, reason="Moving terminal attachment; task continues"
746
+ )
747
+ elif op == "permissions":
748
+ if self.active_threads or self.active_items:
749
+ raise RouterError("Session is busy")
750
+ if thread:
751
+ await self.backend.call("thread/resume", {"threadId": thread, **params})
752
+ self.merge_settings(thread, params)
753
+ self.launch_overrides = {**self.launch_overrides, **params}
754
+ if params.get("permissions"):
755
+ self.launch_overrides.pop("sandbox", None)
756
+ elif params.get("sandbox"):
757
+ self.launch_overrides.pop("permissions", None)
758
+ elif op == "switch":
759
+ label = params.get("account")
760
+ index = next(
761
+ (i for i, a in enumerate(self.accounts) if a.label == label), None
762
+ )
763
+ if index is None:
764
+ raise RouterError("Unknown account")
765
+ if self.active_threads or self.active_items or self.server_requests:
766
+ raise RouterError("Session is busy; switch at an idle boundary")
767
+ options = self.pool.options(label) if self.pool else {}
768
+ model = self.model_for(thread) if thread else None
769
+ if (
770
+ not options.get("enabled", True)
771
+ or options.get("maintenance")
772
+ or options.get("models")
773
+ and model
774
+ and model not in options["models"]
775
+ ):
776
+ raise RouterError(
777
+ "Account is disabled, under maintenance, or unavailable for this model"
778
+ )
779
+ if limited(await self.accounts[index].usage(force=True), model):
780
+ raise RouterError("Account allowance is exhausted")
781
+ if self.pool and not self.pool.reserve(
782
+ label,
783
+ limit=options.get("max_sessions")
784
+ or self.preferences["max_sessions_per_account"],
785
+ ):
786
+ raise RouterError("Account capacity is full or maintenance has started")
787
+ try:
788
+ await self.login(index)
789
+ finally:
790
+ if self.pool:
791
+ self.pool.release_reservation()
792
+ self.note("switched", to=label)
793
+ elif op == "retry":
794
+ self.paused = False
795
+ self.wake.set()
796
+ if self.recovery and not self.recovery.done():
797
+ self.force_retry.set()
798
+ else:
799
+ self.force_retry.clear()
800
+ if self.pool:
801
+ self.pool.clear(
802
+ self.account.label, self.model_for(thread) if thread else None
803
+ )
804
+ self.pool.invalidate(self.account.label)
805
+ if thread and (not self.recovery or self.recovery.done()):
806
+ self.clear_review(thread)
807
+ self.recovery = asyncio.create_task(
808
+ self.continue_thread(thread, self.user_generation)
809
+ )
810
+ elif op == "wake":
811
+ self.wake.set()
812
+ else:
813
+ raise RouterError("Unknown control operation")
814
+ if self.pool:
815
+ self.pool.touch(paused=self.paused)
816
+ self.status_changed.set()
817
+
818
+ async def restart_backend(self):
819
+ async with self.restart_lock:
820
+ if self.backend.proc.returncode is None:
821
+ return
822
+ self.note("restarting")
823
+ in_flight = bool(self.active_items or self.server_requests)
824
+ old = self.backend
825
+ candidate = None
826
+ try:
827
+ candidate = await self.backend_factory()
828
+ candidate.handler = self.event
829
+ await candidate.initialize()
830
+ self.backend = candidate
831
+ await self.login(self.index)
832
+ except Exception:
833
+ self.backend = old
834
+ if candidate:
835
+ await candidate.close()
836
+ self.note("backend_restart_failed", reason="startup_or_login_failed")
837
+ return
838
+ self.active_threads.clear()
839
+ self.server_requests.clear()
840
+ self.pending_server_requests.clear()
841
+ self.replayed_requests.clear()
842
+ self.active_items.clear()
843
+ for thread in list(self.last_turn):
844
+ params = {"threadId": thread}
845
+ params.update(
846
+ {
847
+ k: v
848
+ for k, v in self.thread_settings.get(thread, {}).items()
849
+ if k in RESUME_FIELDS
850
+ }
851
+ )
852
+ if self.launch_overrides.get("config"):
853
+ params["config"] = self.launch_overrides["config"]
854
+ try:
855
+ await self.backend.call("thread/resume", params)
856
+ except (RouterError, asyncio.TimeoutError):
857
+ self.hold_for_review(thread)
858
+ self.note(
859
+ "needs_attention", thread=thread, reason="backend_resume_failed"
860
+ )
861
+ continue
862
+ if in_flight:
863
+ self.hold_for_review(thread)
864
+ self.note(
865
+ "needs_attention",
866
+ thread=thread,
867
+ reason="action_interrupted_by_crash",
868
+ )
869
+ else:
870
+ state, turn = await self.reconcile(thread)
871
+ persisted = self.pool.intent(thread) if self.pool else {}
872
+ if (
873
+ state != "completed"
874
+ and not self.active_threads
875
+ and persisted
876
+ and persisted.get("state") not in {"completed", "idle", None}
877
+ ):
878
+ self.recovery = asyncio.create_task(
879
+ self.restore_recovery(thread, persisted)
880
+ )
881
+ await old.close()
882
+ if self.paused_uncertain:
883
+ self.note("needs_attention", reason="backend_restored_review_required")
884
+ else:
885
+ self.note("backend_restored")
886
+ self.status_changed.set()
887
+
888
+ @property
889
+ def account(self):
890
+ return self.accounts[self.index]
891
+
892
+ async def login(self, index):
893
+ # A sent login cannot be undone by cancelling its waiter. Complete the
894
+ # state transition before allowing a new user turn to use this backend.
895
+ task = asyncio.create_task(self._login(index))
896
+ try:
897
+ return await asyncio.shield(task)
898
+ except asyncio.CancelledError:
899
+ await task
900
+ raise
901
+
902
+ async def _login(self, index):
903
+ async with self.login_lock:
904
+ credentials = await self.accounts[index].credentials()
905
+ await self.backend.call(
906
+ "account/login/start", {"type": "chatgptAuthTokens", **credentials}
907
+ )
908
+ self.index = index
909
+ # Acceptance of token syntax alone isn't a verified, usable login.
910
+ result = await self.backend.call("account/read", {"refreshToken": False})
911
+ if not result.get("account"):
912
+ raise RouterError("The selected account was not confirmed by Codex.")
913
+ if self.pool:
914
+ self.pool.assign(self.account.label)
915
+
916
+ async def emit(self, message):
917
+ if self.client:
918
+ try:
919
+ await self.client.send(json.dumps(message))
920
+ except Exception:
921
+ pass
922
+
923
+ async def event(self, message):
924
+ method, params = message.get("method"), message.get("params", {})
925
+ if method == "account/chatgptAuthTokens/refresh":
926
+ try:
927
+ # Never return another identity in response to a refresh request.
928
+ credentials = await self.account.credentials(refresh=True)
929
+ previous = params.get("previousAccountId")
930
+ if previous and previous != credentials["chatgptAccountId"]:
931
+ raise RouterError(
932
+ "Refresh account mismatch; re-enroll the account."
933
+ )
934
+ await self.backend.send({"id": message["id"], "result": credentials})
935
+ except (RouterError, asyncio.TimeoutError):
936
+ await self.backend.send(
937
+ {
938
+ "id": message["id"],
939
+ "error": {
940
+ "code": -32000,
941
+ "message": "Account requires sign-in.",
942
+ },
943
+ }
944
+ )
945
+ return
946
+ if "id" in message and method:
947
+ self.server_requests.add(message["id"])
948
+ self.pending_server_requests[message["id"]] = message
949
+ if not self.client:
950
+ self.note(
951
+ "needs_attention",
952
+ reason="approval_or_tool_response_required",
953
+ thread=params.get("threadId"),
954
+ )
955
+ if method == "serverRequest/resolved":
956
+ self.server_requests.discard(params.get("requestId"))
957
+ self.pending_server_requests.pop(params.get("requestId"), None)
958
+ if method == "thread/started":
959
+ thread = params.get("thread", {})
960
+ if self.pool and thread.get("id"):
961
+ self.pool.claim(thread["id"])
962
+ if method == "turn/started":
963
+ self.active_threads.add(params["threadId"])
964
+ self.turn_ids[params["threadId"]] = params.get("turn", {}).get("id")
965
+ if params["threadId"] in self.pending_intents:
966
+ self.pending_intents[params["threadId"]]["accepted"] = True
967
+ self.note("working", thread=params["threadId"])
968
+ self.persist(
969
+ params["threadId"],
970
+ accepted=True,
971
+ turn=params.get("turn", {}).get("id"),
972
+ state="working",
973
+ )
974
+ elif method == "item/started" and params.get("item", {}).get("type") in {
975
+ "commandExecution",
976
+ "mcpToolCall",
977
+ "fileChange",
978
+ "dynamicToolCall",
979
+ }:
980
+ self.active_items.add(params["item"]["id"])
981
+ if params.get("threadId"):
982
+ self.persist(params["threadId"], active_items=True)
983
+ elif method == "item/completed":
984
+ self.active_items.discard(params.get("item", {}).get("id"))
985
+ if params.get("threadId"):
986
+ self.persist(params["threadId"], active_items=bool(self.active_items))
987
+ elif method == "model/rerouted":
988
+ if params.get("toModel"):
989
+ self.merge_settings(params["threadId"], {"model": params["toModel"]})
990
+ elif method == "thread/closed":
991
+ self.active_threads.discard(params["threadId"])
992
+ elif method == "thread/status/changed":
993
+ if params.get("status", {}).get("type") == "active":
994
+ self.active_threads.add(params["threadId"])
995
+ else:
996
+ self.active_threads.discard(params["threadId"])
997
+ elif method == "turn/completed":
998
+ thread = params["threadId"]
999
+ self.active_threads.discard(thread)
1000
+ turn = params.get("turn", {})
1001
+ self.turn_ids[thread] = turn.get("id")
1002
+ self.persist(
1003
+ thread,
1004
+ state=turn.get("status"),
1005
+ accepted=True,
1006
+ turn=turn.get("id"),
1007
+ active_items=bool(self.active_items),
1008
+ )
1009
+ failure = (turn.get("error") or {}).get("codexErrorInfo")
1010
+ key = (thread, turn.get("id"))
1011
+ kind = failure_kind(failure, (turn.get("error") or {}).get("message"))
1012
+ if failure == "modelNotFound" and self.preferences.get("fallback_models"):
1013
+ kind = "model"
1014
+ if (
1015
+ turn.get("status") == "failed"
1016
+ and failure is not None
1017
+ and not kind
1018
+ and key in self.seen_failures
1019
+ ):
1020
+ if self.recovery:
1021
+ self.recovery.cancel()
1022
+ self.episodes.pop(thread, None)
1023
+ self.note(
1024
+ "needs_attention",
1025
+ thread=thread,
1026
+ reason="terminal_error_requires_review",
1027
+ )
1028
+ if turn.get("status") == "completed":
1029
+ if (
1030
+ key in self.seen_failures
1031
+ and self.recovery
1032
+ and self.recovery is not asyncio.current_task()
1033
+ ):
1034
+ self.recovery.cancel()
1035
+ self.episodes.pop(thread, None)
1036
+ self.persist(thread, kind=None, episode=None, next_wake=None)
1037
+ self.note("completed", thread=thread)
1038
+ if (
1039
+ turn.get("status") == "failed"
1040
+ and kind
1041
+ and key not in self.seen_failures
1042
+ ):
1043
+ self.seen_failures.add(key)
1044
+ if thread in self.last_turn:
1045
+ if self.recovery and not self.recovery.done():
1046
+ self.recovery.cancel()
1047
+ self.recovery = asyncio.create_task(
1048
+ self.recover(thread, self.user_generation, kind)
1049
+ )
1050
+ elif method == "error" and params.get("willRetry") is False:
1051
+ error = params.get("error") or {}
1052
+ thread = params.get("threadId")
1053
+ key = (thread, params.get("turnId"))
1054
+ if (
1055
+ failure_kind(error.get("codexErrorInfo"), error.get("message"))
1056
+ == "capacity"
1057
+ and thread in self.last_turn
1058
+ and key not in self.seen_failures
1059
+ ):
1060
+ self.seen_failures.add(key)
1061
+ if self.recovery and not self.recovery.done():
1062
+ self.recovery.cancel()
1063
+ self.recovery = asyncio.create_task(
1064
+ self.recover(thread, self.user_generation, "capacity")
1065
+ )
1066
+ self.status_changed.set()
1067
+ await self.emit(message)
1068
+
1069
+ def model_for(self, thread):
1070
+ return self.thread_settings.get(thread, {}).get("model") or self.last_turn.get(
1071
+ thread, {}
1072
+ ).get("model")
1073
+
1074
+ def deadline(self, index, model):
1075
+ return (
1076
+ self.pool.deadline(self.accounts[index].label, model)
1077
+ if self.pool
1078
+ else self.local_deadlines.get(index, 0)
1079
+ )
1080
+
1081
+ async def mark(self, index, model, reason="quota"):
1082
+ snapshot = None
1083
+ if (
1084
+ self.pool
1085
+ and reason == "quota"
1086
+ and self.pool.deadline(self.accounts[index].label, model) > self.clock()
1087
+ ):
1088
+ return self.pool.deadline(self.accounts[index].label, model)
1089
+ if reason == "quota":
1090
+ try:
1091
+ snapshot = await self.accounts[index].usage(force=True)
1092
+ except (RouterError, asyncio.TimeoutError):
1093
+ pass
1094
+ if self.pool:
1095
+ return self.pool.mark(
1096
+ self.accounts[index].label,
1097
+ model,
1098
+ snapshot=snapshot,
1099
+ reason=reason,
1100
+ delay=self.interval,
1101
+ )
1102
+ until = (
1103
+ reset_time(snapshot or {}, model, self.clock())
1104
+ or self.clock() + self.interval
1105
+ )
1106
+ self.local_deadlines[index] = until
1107
+ return until
1108
+
1109
+ async def wait_idle(self, generation):
1110
+ began = time.monotonic()
1111
+ while (
1112
+ self.active_threads
1113
+ or self.active_items
1114
+ or self.server_requests
1115
+ or self.paused
1116
+ ):
1117
+ if generation != self.user_generation or not self.attached():
1118
+ return False
1119
+ self.status_changed.clear()
1120
+ try:
1121
+ await asyncio.wait_for(
1122
+ self.status_changed.wait(), self.preferences["watchdog_interval"]
1123
+ )
1124
+ except asyncio.TimeoutError:
1125
+ for thread in list(self.active_threads):
1126
+ with contextlib.suppress(RouterError, asyncio.TimeoutError):
1127
+ await self.reconcile(thread)
1128
+ if time.monotonic() - began > self.preferences["idle_warning_after"]:
1129
+ self.note(
1130
+ "waiting_for_activity", reason="actions_or_approvals_not_idle"
1131
+ )
1132
+ if self.paused:
1133
+ self.wake.clear()
1134
+ try:
1135
+ await asyncio.wait_for(self.wake.wait(), 2)
1136
+ except asyncio.TimeoutError:
1137
+ pass
1138
+ return generation == self.user_generation and self.attached()
1139
+
1140
+ async def continue_thread(self, thread, generation):
1141
+ if not self.can_continue(thread, generation):
1142
+ return
1143
+ if not await self.wait_idle(generation) or not self.can_continue(
1144
+ thread, generation
1145
+ ):
1146
+ return
1147
+ params = {
1148
+ "threadId": thread,
1149
+ "input": [{"type": "text", "text": CONTINUE, "text_elements": []}],
1150
+ }
1151
+ for key in ("outputSchema", "cyberAccessProgram", "effort", "summary"):
1152
+ if self.thread_settings.get(thread, {}).get(key) is not None:
1153
+ params[key] = self.thread_settings[thread][key]
1154
+ params["clientUserMessageId"] = str(uuid.uuid4())
1155
+ baseline = self.turn_ids.get(thread) or "__unknown_before_submission__"
1156
+ self.pending_intents[thread] = {"baseline": baseline, "accepted": False}
1157
+ self.persist(thread, accepted=False, state="submitting", baseline=baseline)
1158
+ try:
1159
+ result = await self.backend.call("turn/start", params)
1160
+ turn = result.get("turn", {})
1161
+ if turn.get("id"):
1162
+ self.turn_ids[thread] = turn["id"]
1163
+ self.pending_intents[thread]["accepted"] = True
1164
+ self.persist(
1165
+ thread,
1166
+ accepted=True,
1167
+ turn=turn["id"],
1168
+ state=turn.get("status", "working"),
1169
+ )
1170
+ except RemoteError as error:
1171
+ if error.payload.get("code") == -32001:
1172
+ return await self.recover(thread, generation, "transient")
1173
+ raise
1174
+ except (RouterError, asyncio.TimeoutError):
1175
+ return await self.uncertain(thread, generation, baseline)
1176
+ self.report(f"Continued the saved thread with {self.account.label}.")
1177
+
1178
+ async def restore_recovery(self, thread, persisted):
1179
+ """Honor an accepted retry's original deadline after a process restart."""
1180
+ generation = self.user_generation
1181
+ if persisted.get("review_required"):
1182
+ self.hold_for_review(
1183
+ thread, persisted.get("review_reason") or "review_required"
1184
+ )
1185
+ return
1186
+ episode = persisted.get("episode")
1187
+ if episode:
1188
+ episode = dict(episode)
1189
+ label = persisted.get("account_label")
1190
+ episode["account"] = next(
1191
+ (i for i, a in enumerate(self.accounts) if a.label == label),
1192
+ episode.get("account", self.index),
1193
+ )
1194
+ self.episodes[thread] = episode
1195
+ deadline = persisted.get("next_wake")
1196
+ if deadline and deadline > self.clock():
1197
+ self.note("retry_scheduled", thread=thread, next_wake=deadline)
1198
+ await self.wait_until(
1199
+ deadline, generation, allow_early=persisted.get("kind") == "quota"
1200
+ )
1201
+ if persisted.get("active_items") or not persisted.get("accepted", False):
1202
+ self.hold_for_review(thread)
1203
+ self.note(
1204
+ "needs_attention",
1205
+ thread=thread,
1206
+ reason="worker_restarted_during_uncertain_action",
1207
+ )
1208
+ return
1209
+ if not await self.wait_idle(generation):
1210
+ return
1211
+ if persisted.get("kind") == "auth":
1212
+ try:
1213
+ await self.account.credentials(refresh=True)
1214
+ await self.login(self.index)
1215
+ except (RouterError, asyncio.TimeoutError):
1216
+ return await self.recover(thread, generation, "auth")
1217
+ await self.continue_thread(thread, generation)
1218
+
1219
+ async def choose(self, model=None, preferred=None):
1220
+ """Return (account index or None, earliest next read timestamp)."""
1221
+ loads = self.pool.loads() if self.pool else {}
1222
+ start = self.index
1223
+
1224
+ def score(i):
1225
+ account = self.accounts[i]
1226
+ options = self.pool.options(account.label) if self.pool else {}
1227
+ snapshot = self.pool.cached(account.label, 3600) if self.pool else None
1228
+ remaining = min(
1229
+ (
1230
+ 100 - (b.get(k) or {}).get("usedPercent", 0)
1231
+ for b in __import__("pool").buckets(snapshot or {}, model)
1232
+ for k in ("primary", "secondary")
1233
+ ),
1234
+ default=100,
1235
+ )
1236
+ mode = self.preferences["routing"]
1237
+ fav = self.preferences.get("preferred_accounts", [])
1238
+ priority = -options.get("priority", 0)
1239
+ load = loads.get(account.label, 0)
1240
+ if mode == "quota":
1241
+ rank = (-remaining, load, priority)
1242
+ elif mode == "priority":
1243
+ rank = (priority, load, -remaining)
1244
+ elif mode == "round-robin":
1245
+ rank = ((i - start - 1) % len(self.accounts),)
1246
+ else:
1247
+ rank = (load, priority, -remaining)
1248
+ return (
1249
+ 0 if i == preferred else 1,
1250
+ fav.index(account.label) if account.label in fav else len(fav),
1251
+ *rank,
1252
+ (i - start - 1) % len(self.accounts),
1253
+ )
1254
+
1255
+ order = sorted(range(len(self.accounts)), key=score)
1256
+ deadlines = []
1257
+ for index in order:
1258
+ deadline = self.deadline(index, model)
1259
+ if deadline > self.clock():
1260
+ deadlines.append(deadline)
1261
+ continue
1262
+ account = self.accounts[index]
1263
+ options = self.pool.options(account.label) if self.pool else {}
1264
+ if (
1265
+ not options.get("enabled", True)
1266
+ or options.get("maintenance")
1267
+ or options.get("models")
1268
+ and model
1269
+ and model not in options["models"]
1270
+ ):
1271
+ continue
1272
+ try:
1273
+ # A reset timestamp authorizes a new check, not usage itself.
1274
+ snapshot = await account.usage(force=bool(deadline))
1275
+ if limited(snapshot, model):
1276
+ if self.pool:
1277
+ until = self.pool.mark(
1278
+ account.label, model, snapshot=snapshot, delay=self.interval
1279
+ )
1280
+ else:
1281
+ until = (
1282
+ reset_time(snapshot, model, self.clock())
1283
+ or self.clock() + self.interval
1284
+ )
1285
+ self.local_deadlines[index] = until
1286
+ deadlines.append(until)
1287
+ continue
1288
+ if self.pool:
1289
+ self.pool.clear(account.label, model)
1290
+ else:
1291
+ self.local_deadlines.pop(index, None)
1292
+ if self.pool and not self.pool.reserve(
1293
+ account.label,
1294
+ expected_load=loads.get(account.label, 0),
1295
+ limit=options.get("max_sessions")
1296
+ or self.preferences["max_sessions_per_account"],
1297
+ ttl=self.preferences["reservation_seconds"],
1298
+ ):
1299
+ deadlines.append(self.clock() + 2)
1300
+ continue
1301
+ return index, 0
1302
+ except (RouterError, asyncio.TimeoutError):
1303
+ deadlines.append(await self.mark(index, model, "unavailable"))
1304
+ return None, min(deadlines, default=self.clock() + self.interval)
1305
+
1306
+ async def recover(self, thread, generation, kind="quota"):
1307
+ if _recovery_active.get() is asyncio.current_task():
1308
+ raise RecoveryTransition(kind)
1309
+ token = _recovery_active.set(asyncio.current_task())
1310
+ try:
1311
+ while self.can_continue(thread, generation):
1312
+ try:
1313
+ return await self._recover_episode(thread, generation, kind)
1314
+ except RecoveryTransition as transition:
1315
+ kind = transition.kind
1316
+ await asyncio.sleep(0)
1317
+ finally:
1318
+ _recovery_active.reset(token)
1319
+
1320
+ async def _recover_episode(self, thread, generation, kind):
1321
+ try:
1322
+ model = self.model_for(thread)
1323
+ if kind == "model":
1324
+ tried = self.fallback_tried.setdefault(thread, set())
1325
+ tried.add(model)
1326
+ fallback = next(
1327
+ (
1328
+ m
1329
+ for m in self.preferences.get("fallback_models", [])
1330
+ if m not in tried
1331
+ ),
1332
+ None,
1333
+ )
1334
+ if not fallback:
1335
+ self.note(
1336
+ "needs_attention",
1337
+ thread=thread,
1338
+ reason="configured_model_fallbacks_exhausted",
1339
+ )
1340
+ return
1341
+ tried.add(fallback)
1342
+ self.merge_settings(thread, {"model": fallback})
1343
+ await self.backend.call(
1344
+ "thread/resume", {"threadId": thread, "model": fallback}
1345
+ )
1346
+ self.note("model_fallback", thread=thread, model=fallback)
1347
+ return await self.continue_thread(thread, generation)
1348
+ if kind == "transient" and hasattr(self.account, "network"):
1349
+ with contextlib.suppress(RouterError, asyncio.TimeoutError):
1350
+ self.network_state = await self.account.network()
1351
+ if self.network_state is False:
1352
+ self.note("connection_probe_failed", thread=thread)
1353
+ if kind != "quota":
1354
+ episode = self.episodes.setdefault(
1355
+ thread, {"account": self.index, "start": self.clock(), "attempt": 0}
1356
+ )
1357
+ if episode["account"] != self.index:
1358
+ episode = {
1359
+ "account": self.index,
1360
+ "start": self.clock(),
1361
+ "attempt": 0,
1362
+ }
1363
+ self.episodes[thread] = episode
1364
+ if kind == "capacity" and episode["attempt"] >= self.retry_attempts:
1365
+ if self.preferences.get("fallback_models"):
1366
+ return await self.recover(thread, generation, "model")
1367
+ if not self.retry_attempts:
1368
+ self.note(
1369
+ "needs_attention",
1370
+ thread=thread,
1371
+ reason="capacity_retries_disabled",
1372
+ )
1373
+ return
1374
+ # Model-wide capacity does not establish an account problem.
1375
+ # Keep waiting on this account/model instead of cycling logins.
1376
+ episode = {
1377
+ "account": self.index,
1378
+ "start": self.clock(),
1379
+ "attempt": 0,
1380
+ }
1381
+ self.episodes[thread] = episode
1382
+ if episode["attempt"] < self.retry_attempts:
1383
+ episode["attempt"] += 1
1384
+ target = episode["start"] + episode["attempt"] * self.retry_interval
1385
+ self.report(
1386
+ f"Service/authentication failure on {self.account.label}. Retry {episode['attempt']}/{self.retry_attempts} at {time.strftime('%H:%M:%S', time.localtime(target))}."
1387
+ )
1388
+ self.note(
1389
+ "retry_scheduled",
1390
+ thread=thread,
1391
+ next_wake=target,
1392
+ attempt=episode["attempt"],
1393
+ )
1394
+ self.persist(
1395
+ thread,
1396
+ generation=generation,
1397
+ kind=kind,
1398
+ episode=episode,
1399
+ next_wake=target,
1400
+ account_label=self.account.label,
1401
+ state="retry_scheduled",
1402
+ )
1403
+ await self.wait_until(target, generation, allow_early=False)
1404
+ if not await self.wait_idle(generation):
1405
+ return
1406
+ if kind == "auth":
1407
+ try:
1408
+ await self.account.credentials(refresh=True)
1409
+ await self.login(self.index)
1410
+ except (RouterError, asyncio.TimeoutError):
1411
+ # Keep this same schedule even if login is revoked.
1412
+ return await self.recover(thread, generation, kind)
1413
+ return await self.continue_thread(thread, generation)
1414
+ self.report(
1415
+ "Three retries over 15 minutes failed; trying another account."
1416
+ )
1417
+ self.episodes.pop(thread, None)
1418
+ self.report(
1419
+ f"Selecting another account after {kind} failure on {self.account.label}."
1420
+ )
1421
+ self.note("selecting_account", thread=thread, reason=kind)
1422
+ await self.mark(self.index, model, kind)
1423
+ if kind != "quota" and self.preferences.get("fallback_models"):
1424
+ configured = self.preferences["fallback_models"]
1425
+ if kind in {"transient", "capacity"} and model not in configured:
1426
+ self.merge_settings(thread, {"model": configured[0]})
1427
+ await self.backend.call(
1428
+ "thread/resume", {"threadId": thread, "model": configured[0]}
1429
+ )
1430
+ model = configured[0]
1431
+ self.note("model_fallback", thread=thread, model=model)
1432
+ while generation == self.user_generation and self.attached():
1433
+ if not await self.wait_idle(generation):
1434
+ return
1435
+ index, deadline = await self.choose(model)
1436
+ if index is not None:
1437
+ if not await self.wait_idle(generation):
1438
+ return
1439
+ try:
1440
+ await self.login(index)
1441
+ return await self.continue_thread(thread, generation)
1442
+ except (RouterError, asyncio.TimeoutError):
1443
+ if self.pool:
1444
+ self.pool.release_reservation()
1445
+ await self.mark(index, model, "unavailable")
1446
+ continue
1447
+ self.report(
1448
+ f"All accounts unavailable. Sleeping until {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(deadline))}; no usage requests before then."
1449
+ )
1450
+ self.note("waiting_for_reset", thread=thread, next_wake=deadline)
1451
+ self.persist(
1452
+ thread,
1453
+ generation=generation,
1454
+ kind=kind,
1455
+ next_wake=deadline,
1456
+ account_label=self.account.label,
1457
+ state="waiting_for_reset",
1458
+ )
1459
+ await self.wait_until(max(self.clock() + 1, deadline), generation)
1460
+ except (asyncio.CancelledError, RecoveryTransition):
1461
+ raise
1462
+ except Exception:
1463
+ self.hold_for_review(thread, "handoff_failed")
1464
+ self.note("needs_attention", thread=thread, reason="handoff_failed")
1465
+ self.report(
1466
+ "Automatic handoff stopped. Your saved thread is retained; continue manually."
1467
+ )
1468
+
1469
+ async def forward(self, message, init_result):
1470
+ method = message.get("method")
1471
+ request_id = message.get("id")
1472
+ if (
1473
+ method in {"thread/start", "thread/resume", "thread/fork"}
1474
+ and self.launch_overrides
1475
+ ):
1476
+ message = dict(message)
1477
+ message["params"] = {**message.get("params", {}), **self.launch_overrides}
1478
+ if (
1479
+ method in {"thread/start", "thread/resume", "thread/fork"}
1480
+ and self.additional_roots
1481
+ ):
1482
+ message = {**message, "params": dict(message.get("params", {}))}
1483
+ params = message["params"]
1484
+ params["runtimeWorkspaceRoots"] = list(
1485
+ dict.fromkeys(
1486
+ [
1487
+ *(
1488
+ params.get("runtimeWorkspaceRoots")
1489
+ or [params.get("cwd") or self.project]
1490
+ ),
1491
+ *self.additional_roots,
1492
+ ]
1493
+ )
1494
+ )
1495
+ if not method:
1496
+ # Approval responses keep the server's original ID, never our call IDs.
1497
+ if request_id in self.server_requests:
1498
+ self.server_requests.discard(request_id)
1499
+ self.pending_server_requests.pop(request_id, None)
1500
+ await self.backend.send(message)
1501
+ return
1502
+ if method == "initialized":
1503
+ return
1504
+ if method == "initialize":
1505
+ await self.emit({"id": request_id, "result": init_result})
1506
+ return
1507
+ if method.startswith("account/login") or method == "account/logout":
1508
+ await self.emit(
1509
+ {
1510
+ "id": request_id,
1511
+ "error": {
1512
+ "code": -32600,
1513
+ "message": "Close this controller before managing pool logins with relay add.",
1514
+ },
1515
+ }
1516
+ )
1517
+ return
1518
+ if method in {"thread/resume", "turn/start", "turn/steer"}:
1519
+ thread = message.get("params", {}).get("threadId")
1520
+ if self.pool and thread and not self.pool.claim(thread):
1521
+ await self.emit(
1522
+ {
1523
+ "id": request_id,
1524
+ "error": {
1525
+ "code": -32600,
1526
+ "message": "This saved conversation is already open in another wrapper terminal. Use a different thread or close that terminal.",
1527
+ },
1528
+ }
1529
+ )
1530
+ return
1531
+ if method in {
1532
+ "turn/start",
1533
+ "turn/steer",
1534
+ "turn/interrupt",
1535
+ "thread/archive",
1536
+ "thread/delete",
1537
+ }:
1538
+ self.user_generation += 1
1539
+ if self.recovery and not self.recovery.done():
1540
+ self.recovery.cancel()
1541
+ self.episodes.clear()
1542
+ if method == "turn/start" and message.get("params", {}).get("model"):
1543
+ self.fallback_tried.pop(message["params"]["threadId"], None)
1544
+ self.status_changed.set()
1545
+ if method == "turn/start":
1546
+ self.clear_review(message["params"]["threadId"])
1547
+ self.last_turn[message["params"]["threadId"]] = message["params"]
1548
+ self.merge_settings(message["params"]["threadId"], message["params"])
1549
+ try:
1550
+ if request_id is None:
1551
+ await self.backend.send(message)
1552
+ else:
1553
+ if method == "turn/start" and self.pool:
1554
+ thread = message["params"]["threadId"]
1555
+ model = self.model_for(thread)
1556
+ allowed = self.pool.options(self.account.label).get("models", [])
1557
+ if allowed and model and model not in allowed:
1558
+ if (
1559
+ self.active_threads
1560
+ or self.active_items
1561
+ or self.server_requests
1562
+ ):
1563
+ raise RemoteError(
1564
+ {
1565
+ "code": -32600,
1566
+ "message": "This account does not allow that model. Switch accounts after the current work finishes.",
1567
+ }
1568
+ )
1569
+ index, _ = await self.choose(model)
1570
+ if index is None:
1571
+ raise RemoteError(
1572
+ {
1573
+ "code": -32600,
1574
+ "message": "No available account allows that model. Choose another model or update your account model settings.",
1575
+ }
1576
+ )
1577
+ await self.login(index)
1578
+ self.note("switched", to=self.account.label)
1579
+ guard = (
1580
+ self.login_lock
1581
+ if method
1582
+ in {
1583
+ "turn/start",
1584
+ "turn/steer",
1585
+ "turn/interrupt",
1586
+ "thread/start",
1587
+ "thread/resume",
1588
+ }
1589
+ else contextlib.nullcontext()
1590
+ )
1591
+ if method == "turn/start":
1592
+ thread = message["params"]["threadId"]
1593
+ if thread not in self.turn_ids:
1594
+ try:
1595
+ latest = await self.latest(thread)
1596
+ self.turn_ids[thread] = latest.get("id") if latest else None
1597
+ except (RouterError, asyncio.TimeoutError):
1598
+ self.turn_ids[thread] = "__unknown_before_submission__"
1599
+ message["params"].setdefault(
1600
+ "clientUserMessageId", str(uuid.uuid4())
1601
+ )
1602
+ self.pending_intents[thread] = {
1603
+ "baseline": self.turn_ids.get(thread),
1604
+ "generation": self.user_generation,
1605
+ }
1606
+ self.persist(
1607
+ thread,
1608
+ accepted=False,
1609
+ baseline=self.turn_ids.get(thread),
1610
+ state="submitting",
1611
+ )
1612
+ async with guard:
1613
+ result = await self.backend.call(
1614
+ method, message.get("params"), timeout=300
1615
+ )
1616
+ if method == "turn/start" and result.get("turn", {}).get("id"):
1617
+ thread = message["params"]["threadId"]
1618
+ self.turn_ids[thread] = result["turn"]["id"]
1619
+ self.pending_intents[thread]["accepted"] = True
1620
+ self.persist(
1621
+ thread,
1622
+ accepted=True,
1623
+ turn=result["turn"]["id"],
1624
+ state=result["turn"].get("status"),
1625
+ )
1626
+ if method in {"thread/start", "thread/resume", "thread/fork"}:
1627
+ thread = result.get("thread", {})
1628
+ if thread.get("id"):
1629
+ if thread.get("canAcceptDirectInput") is False:
1630
+ if self.pool:
1631
+ self.pool.release_thread(thread["id"])
1632
+ self.report(
1633
+ "This conversation is owned by another Codex backend. Keep using its existing terminal, or fork it into a separate conversation. No running task was stopped."
1634
+ )
1635
+ await self.emit(
1636
+ {
1637
+ "id": request_id,
1638
+ "error": {
1639
+ "code": -32600,
1640
+ "message": "Conversation is open in another Codex backend. Close/unload it there before resuming here, or use relay run -- fork "
1641
+ + thread["id"]
1642
+ + " --yolo for an independent copy. No running task was stopped.",
1643
+ },
1644
+ }
1645
+ )
1646
+ return
1647
+ self.merge_settings(thread["id"], result)
1648
+ self.merge_settings(thread["id"], message.get("params", {}))
1649
+ self.last_turn.setdefault(
1650
+ thread["id"], {"threadId": thread["id"]}
1651
+ )
1652
+ self.persist(thread["id"], role="foreground")
1653
+ if self.on_thread:
1654
+ await self.on_thread(
1655
+ thread["id"], self.thread_settings[thread["id"]]
1656
+ )
1657
+ if self.pool:
1658
+ self.pool.claim(thread["id"])
1659
+ self.pool.touch(
1660
+ thread=thread["id"],
1661
+ cwd=result.get("cwd") or thread.get("cwd"),
1662
+ model=self.model_for(thread["id"]),
1663
+ )
1664
+ await self.emit({"id": request_id, "result": result})
1665
+ if method in {"thread/start", "thread/resume", "thread/fork"}:
1666
+ thread_id = result.get("thread", {}).get("id")
1667
+ for ident, pending in list(self.pending_server_requests.items()):
1668
+ if ident not in self.replayed_requests and pending.get(
1669
+ "params", {}
1670
+ ).get("threadId") in {None, thread_id}:
1671
+ await self.emit(pending)
1672
+ self.replayed_requests.add(ident)
1673
+ except RemoteError as error:
1674
+ await self.emit({"id": request_id, "error": error.payload})
1675
+ # Ingress overload rejects a turn before a turn/completed event
1676
+ # exists. Its documented code is explicit, not message matching.
1677
+ if method == "turn/start" and error.payload.get("code") == -32001:
1678
+ if self.recovery and not self.recovery.done():
1679
+ self.recovery.cancel()
1680
+ thread = message["params"]["threadId"]
1681
+ self.recovery = asyncio.create_task(
1682
+ self.recover(thread, self.user_generation, "transient")
1683
+ )
1684
+ except (RouterError, asyncio.TimeoutError):
1685
+ await self.emit(
1686
+ {
1687
+ "id": request_id,
1688
+ "error": {
1689
+ "code": -32000,
1690
+ "message": "Codex request failed. No automatic retry of this action was made.",
1691
+ },
1692
+ }
1693
+ )
1694
+ if method == "turn/start":
1695
+ thread = message["params"]["threadId"]
1696
+ intent = self.pending_intents.get(thread, {})
1697
+ if self.recovery and not self.recovery.done():
1698
+ self.recovery.cancel()
1699
+ self.recovery = asyncio.create_task(
1700
+ self.uncertain(thread, self.user_generation, intent.get("baseline"))
1701
+ )
1702
+
1703
+ async def connect(self, websocket, init_result):
1704
+ if self.client:
1705
+ await websocket.close(code=1008, reason="One terminal per controller.")
1706
+ return
1707
+ self.client = websocket
1708
+ self.replayed_requests.clear()
1709
+ if self.pool:
1710
+ self.pool.touch(attached=True)
1711
+ try:
1712
+ async for raw in websocket:
1713
+ try:
1714
+ message = json.loads(raw)
1715
+ if not isinstance(message, dict):
1716
+ raise ValueError()
1717
+ except ValueError:
1718
+ await websocket.close(code=1008, reason="Invalid RPC message.")
1719
+ break
1720
+ # Keep notifications ordered; independent long RPC calls must not
1721
+ # block approval responses or interruption requests.
1722
+ task = asyncio.create_task(self.forward(message, init_result))
1723
+ self.client_tasks.add(task)
1724
+ task.add_done_callback(self.client_tasks.discard)
1725
+ finally:
1726
+ self.client = None
1727
+ if self.pool:
1728
+ self.pool.touch(attached=False)
1729
+ if not self.background:
1730
+ self.user_generation += 1
1731
+ if self.recovery:
1732
+ self.recovery.cancel()
1733
+ for task in list(self.client_tasks):
1734
+ task.cancel()
1735
+
1736
+
1737
+ async def run_pool(args, state, binary):
1738
+ import service
1739
+
1740
+ from launch import frontend_arguments, working_directory, prepare_worktree
1741
+
1742
+ invocation_cwd = os.getcwd()
1743
+ cwd = working_directory(args.codex_args, invocation_cwd)
1744
+ settings.load(state, cwd, args.profile)
1745
+ arguments = frontend_arguments(args.codex_args, invocation_cwd)
1746
+ arguments, cwd = prepare_worktree(arguments, state, cwd)
1747
+ frontend, overrides = translate(arguments)
1748
+ target = await resolve_resume_target(
1749
+ frontend,
1750
+ Path(os.environ.get("CODEX_HOME", Path.home() / ".codex")),
1751
+ cwd,
1752
+ binary,
1753
+ )
1754
+ pool = Pool(state)
1755
+ try:
1756
+ owner = pool.owner(target) if target else None
1757
+ finally:
1758
+ pool.close()
1759
+ if owner:
1760
+ print(
1761
+ "Reconnecting to existing wrapper session " + owner[:10] + ".", flush=True
1762
+ )
1763
+ return await service.attach(
1764
+ state, owner, binary, arguments, wait_forever=True, cwd=invocation_cwd
1765
+ )
1766
+ ident = service.start(
1767
+ state,
1768
+ name=args.name or Path(cwd).name,
1769
+ cwd=cwd,
1770
+ first=args.first,
1771
+ arguments=arguments,
1772
+ profile=args.profile,
1773
+ interactive=True,
1774
+ binary=binary,
1775
+ )
1776
+ try:
1777
+ return await service.attach(
1778
+ state, ident, binary, arguments, wait_forever=True, cwd=invocation_cwd
1779
+ )
1780
+ finally:
1781
+ # A foreground launch normally closes its worker after the terminal exits.
1782
+ # Background mode and explicit takeover preserve the supervised worker.
1783
+ spec_path = service.folder(state) / (ident + ".job.json")
1784
+ spec = json.loads(spec_path.read_text()) if spec_path.exists() else {}
1785
+ current_preferences = settings.load(state, cwd, args.profile)
1786
+ if not current_preferences.get("background_on_disconnect") and not spec.get(
1787
+ "takeover"
1788
+ ):
1789
+ pool = Pool(state)
1790
+ pool.command(ident, "stop")
1791
+ pool.close()
1792
+ deadline = time.monotonic() + 10
1793
+ while (
1794
+ service.folder(state) / (ident + ".connection.json")
1795
+ ).exists() and time.monotonic() < deadline:
1796
+ await asyncio.sleep(0.1)
1797
+
1798
+
1799
+ def resume_selection(arguments):
1800
+ from launch import tokens
1801
+
1802
+ parsed = list(tokens(arguments))
1803
+ positionals = [
1804
+ value for name, value, raw in parsed if name is None and len(raw) == 1
1805
+ ]
1806
+ if not positionals or positionals[0] not in {"resume", "fork"}:
1807
+ return None, set()
1808
+ flags = {name for name, _, _ in parsed}
1809
+ return positionals[1] if len(positionals) > 1 else None, flags
1810
+
1811
+
1812
+ def resume_target(arguments, home=None, cwd=None):
1813
+ """Explicit IDs need no native database or backend."""
1814
+ selected, _ = resume_selection(arguments)
1815
+ try:
1816
+ return str(uuid.UUID(selected)) if selected else None
1817
+ except (ValueError, AttributeError):
1818
+ return None
1819
+
1820
+
1821
+ async def resolve_resume_target(arguments, home, cwd, binary=None, backend=None):
1822
+ selected, flags = resume_selection(arguments)
1823
+ explicit = resume_target(arguments)
1824
+ if explicit or not selected and "--last" not in flags:
1825
+ return explicit
1826
+ own_backend = backend is None
1827
+ if own_backend:
1828
+ backend = await RPC.start(binary, home)
1829
+ await backend.initialize()
1830
+ try:
1831
+ params = {
1832
+ "archived": False,
1833
+ "sortKey": "updated_at",
1834
+ "sortDirection": "desc",
1835
+ "limit": 100,
1836
+ }
1837
+ if "--all" not in flags:
1838
+ params["cwd"] = cwd
1839
+ if "--include-non-interactive" in flags:
1840
+ params["sourceKinds"] = [
1841
+ "cli",
1842
+ "vscode",
1843
+ "exec",
1844
+ "appServer",
1845
+ "subAgent",
1846
+ "subAgentReview",
1847
+ "subAgentCompact",
1848
+ "subAgentThreadSpawn",
1849
+ "subAgentOther",
1850
+ "unknown",
1851
+ ]
1852
+ if selected:
1853
+ params["searchTerm"] = selected
1854
+ while True:
1855
+ response = await backend.call("thread/list", params)
1856
+ for thread in response.get("data", []):
1857
+ if not selected or selected in {
1858
+ thread.get("name"),
1859
+ thread.get("title"),
1860
+ thread.get("id"),
1861
+ }:
1862
+ return thread["id"]
1863
+ cursor = response.get("nextCursor")
1864
+ if not cursor:
1865
+ return None
1866
+ params["cursor"] = cursor
1867
+ finally:
1868
+ if own_backend:
1869
+ await backend.close()
1870
+
1871
+
1872
+ def add_account(args, state, binary):
1873
+ home = private_dir(account_home(state, args.label))
1874
+ old_identity = (
1875
+ json.loads((home / "enrolled").read_text()).get("identity")
1876
+ if (home / "enrolled").exists()
1877
+ else None
1878
+ )
1879
+ # Only these isolated homes use the direct Keychain layout. Existing Codex
1880
+ # settings, login entries, session files and binaries are never replaced.
1881
+ config = home / "config.toml"
1882
+ config.write_text(
1883
+ 'cli_auth_credentials_store = "keyring"\n[features]\nsecret_auth_storage = false\n[analytics]\nenabled = false\n'
1884
+ )
1885
+ config.chmod(0o600)
1886
+ command = [binary, "login"] + (["--device-auth"] if args.device_auth else [])
1887
+ print(f"Sign into the ChatGPT account you want to label {args.label}.")
1888
+ result = subprocess.run(command, env=clean_env(home))
1889
+ if result.returncode:
1890
+ raise RouterError("Sign-in did not finish. No account was added to the pool.")
1891
+ # Once native login changes this home's credential, an old marker must no
1892
+ # longer make a failed/duplicate enrollment selectable.
1893
+ (home / "enrolled").unlink(missing_ok=True)
1894
+ credentials = read_keychain(home, include_identity=True)
1895
+ identity = hashlib.sha256(credentials["chatgptAccountId"].encode()).hexdigest()
1896
+ for label in labels(state):
1897
+ if label == args.label:
1898
+ continue
1899
+ other = json.loads((account_home(state, label) / "enrolled").read_text())
1900
+ if identity == other.get("identity"):
1901
+ raise RouterError(
1902
+ "That account is already enrolled under another label. Select a different browser account and try again."
1903
+ )
1904
+ marker = home / "enrolled"
1905
+ marker.write_text(json.dumps({"identity": identity}) + "\n")
1906
+ marker.chmod(0o600)
1907
+ if credentials.get("email"):
1908
+ from settings import atomic_json
1909
+
1910
+ atomic_json(
1911
+ home / "identity.json",
1912
+ {"email": credentials["email"], "identity": identity},
1913
+ )
1914
+ pool = Pool(state)
1915
+ if old_identity != identity:
1916
+ pool.reset_identity(args.label)
1917
+ else:
1918
+ pool.invalidate(args.label)
1919
+ with pool.db() as db:
1920
+ db.execute(
1921
+ "DELETE FROM blocks WHERE label=? AND reason<>'quota'", (args.label,)
1922
+ )
1923
+ pool.set_options(args.label, enabled=True)
1924
+ pool.close()
1925
+ print(
1926
+ f"Saved {args.label} in macOS Keychain. No credentials were printed or saved in the project."
1927
+ )
1928
+
1929
+
1930
+ def transport_guard(token):
1931
+ async def authenticate(connection, request):
1932
+ # Browser origins are rejected even if they present a transport token.
1933
+ expected = token() if callable(token) else token
1934
+ if (
1935
+ request.headers.get("Origin")
1936
+ or request.headers.get("Authorization") != f"Bearer {expected}"
1937
+ ):
1938
+ return connection.respond(403, "Forbidden\n")
1939
+
1940
+ return authenticate
1941
+
1942
+
1943
+ def parser():
1944
+ p = argparse.ArgumentParser(
1945
+ prog="relay", description="Keep Codex moving between your own accounts (macOS)."
1946
+ )
1947
+ p.add_argument("--version", action="version", version="Relay " + VERSION)
1948
+ sub = p.add_subparsers(dest="command", required=True)
1949
+ help_command = sub.add_parser(
1950
+ "help", help="Show Relay help or help for one command"
1951
+ )
1952
+ help_command.add_argument("topic", nargs="*")
1953
+ add = sub.add_parser("add", help="Sign into one account using native Codex login")
1954
+ add.add_argument("label")
1955
+ add.add_argument("--device-auth", action="store_true")
1956
+ setup = sub.add_parser(
1957
+ "setup", help="Enroll any number of accounts using native browser sign-in"
1958
+ )
1959
+ setup.add_argument("--device-auth", action="store_true")
1960
+ setup.add_argument("labels", nargs="*")
1961
+ setup.add_argument("--count", type=int)
1962
+ sub.add_parser("list", help="List enrolled labels; never show credentials")
1963
+ run = sub.add_parser("run", help="Open Codex with automatic account recovery")
1964
+ run.add_argument("--first", help="Account label to start with")
1965
+ run.add_argument(
1966
+ "--profile", help="Recovery profile, separate from the native Codex profile"
1967
+ )
1968
+ run.add_argument("--name", help="Friendly session name")
1969
+ run.add_argument(
1970
+ "--require-recovery",
1971
+ action="store_true",
1972
+ help="Fail rather than run a native mode outside Relay recovery",
1973
+ )
1974
+ run.add_argument(
1975
+ "codex_args",
1976
+ nargs=argparse.REMAINDER,
1977
+ help="Arguments after -- are forwarded to Codex",
1978
+ )
1979
+ sub.add_parser(
1980
+ "check",
1981
+ help="Check CLI version, dependency and account count without signing in",
1982
+ )
1983
+ sub.add_parser(
1984
+ "refresh",
1985
+ help="Fetch current allowance and reset times, even for cached exhausted accounts",
1986
+ )
1987
+ import commands
1988
+
1989
+ commands.add_parsers(sub)
1990
+ sub.metavar = "COMMAND"
1991
+ sub._choices_actions[:] = [
1992
+ item for item in sub._choices_actions if not item.dest.startswith("_")
1993
+ ]
1994
+ return p
1995
+
1996
+
1997
+ def parse_arguments(argv):
1998
+ argv = list(argv)
1999
+ if argv[:1] == ["attach"]:
2000
+ boundary = argv.index("--") if "--" in argv else len(argv)
2001
+ if "--takeover" in argv[1:boundary]:
2002
+ argv.remove("--takeover")
2003
+ argv.insert(1, "--takeover")
2004
+ return parser().parse_args(argv)
2005
+
2006
+
2007
+ async def dispatch(args, state, binary):
2008
+ import commands
2009
+
2010
+ if args.command == "run":
2011
+ from launch import arguments_only, native_mode
2012
+
2013
+ if native_mode(args.codex_args, binary):
2014
+ if getattr(args, "require_recovery", False):
2015
+ raise RouterError(
2016
+ "This mode runs in native Codex outside Relay recovery. Use relay explain -- to inspect the launch."
2017
+ )
2018
+ if sys.stderr.isatty():
2019
+ print(
2020
+ "Running native Codex; automatic account recovery is inactive for this mode.",
2021
+ file=sys.stderr,
2022
+ )
2023
+ process = await asyncio.create_subprocess_exec(
2024
+ binary, *arguments_only(args.codex_args)
2025
+ )
2026
+ return await process.wait()
2027
+
2028
+ if args.command == "_broker":
2029
+ import broker
2030
+
2031
+ return await broker.serve(state, args.binary or binary)
2032
+ if args.command == "_supervise":
2033
+ import service
2034
+
2035
+ return await service.supervise(state, args.session)
2036
+ if args.command == "_serve":
2037
+ import service
2038
+
2039
+ return await service.serve(state, args.session, binary)
2040
+ if args.command not in {
2041
+ "run",
2042
+ "start",
2043
+ "attach",
2044
+ "add",
2045
+ "setup",
2046
+ "check",
2047
+ "refresh",
2048
+ }:
2049
+ result = await commands.dispatch(args, state, binary)
2050
+ if result is not None:
2051
+ return result
2052
+ async with file_lock(state / "compatibility.lock"):
2053
+ try:
2054
+ import runtime
2055
+
2056
+ binary, result, runtime_mode = await asyncio.to_thread(
2057
+ runtime.select, binary, state
2058
+ )
2059
+ except (ValueError, OSError, KeyError, subprocess.SubprocessError):
2060
+ raise RouterError(
2061
+ "Installed Codex does not expose a compatible protocol. Ordinary codex still works; the wrapper needs an adapter update."
2062
+ ) from None
2063
+ if args.command == "check":
2064
+ print(
2065
+ f"{result['version']}; {len(labels(state))} accounts enrolled; capability checks passed."
2066
+ )
2067
+ print(
2068
+ "Parallel terminals enabled; 5/10/15-minute retries; reset-aware scheduling."
2069
+ )
2070
+ print("Native runtime: " + runtime_mode)
2071
+ return 0
2072
+ if args.command in {"start", "attach"}:
2073
+ return await commands.dispatch(args, state, binary)
2074
+ if args.command == "refresh":
2075
+ async with file_lock(state / "enrollment.lock", shared=True):
2076
+ pool = Pool(state)
2077
+ from broker import ClientAccount
2078
+
2079
+ accounts = [
2080
+ ClientAccount(label, account_home(state, label), binary, pool)
2081
+ for label in labels(state)
2082
+ ]
2083
+ try:
2084
+
2085
+ async def refresh(account):
2086
+ try:
2087
+ snapshot = await account.usage(force=True)
2088
+ with pool.db() as db:
2089
+ scopes = [
2090
+ r[0]
2091
+ for r in db.execute(
2092
+ "SELECT scope FROM blocks WHERE label=?",
2093
+ (account.label,),
2094
+ )
2095
+ ]
2096
+ for scope in set(scopes + ["*"]):
2097
+ model = None if scope == "*" else scope
2098
+ if limited(snapshot, model):
2099
+ pool.mark(account.label, model, snapshot=snapshot)
2100
+ else:
2101
+ pool.clear(account.label, model)
2102
+ print(
2103
+ f"{account.label}: allowance and reset times refreshed.",
2104
+ flush=True,
2105
+ )
2106
+ except (RouterError, asyncio.TimeoutError):
2107
+ print(
2108
+ f"{account.label}: refresh unavailable; existing reset schedule retained.",
2109
+ flush=True,
2110
+ )
2111
+
2112
+ await asyncio.gather(*(refresh(a) for a in accounts))
2113
+ finally:
2114
+ await asyncio.gather(
2115
+ *(a.close() for a in accounts), return_exceptions=True
2116
+ )
2117
+ pool.close()
2118
+ pool = Pool(state)
2119
+ pool.signal()
2120
+ pool.close()
2121
+ return 0
2122
+ if args.command in {"add", "setup"}:
2123
+ if args.command == "setup":
2124
+ chosen = list(args.labels)
2125
+ if args.count is not None:
2126
+ if args.count < 1:
2127
+ raise RouterError("Account count must be at least 1")
2128
+ chosen.extend(
2129
+ "account" + str(i)
2130
+ for i in range(1, args.count + 1)
2131
+ if "account" + str(i) not in labels(state)
2132
+ )
2133
+ elif not chosen and not labels(state):
2134
+ chosen = ["account1"]
2135
+ if not chosen:
2136
+ print(
2137
+ f"{len(labels(state))} accounts enrolled. Use setup --count N or add LABEL to add more."
2138
+ )
2139
+ return 0
2140
+ for label in dict.fromkeys(chosen):
2141
+ import argparse
2142
+
2143
+ request = argparse.Namespace(
2144
+ action="add", label=label, device_auth=args.device_auth
2145
+ )
2146
+ await commands.manage(request, state, binary)
2147
+ return 0
2148
+ args.action = "add"
2149
+ return await commands.manage(args, state, binary)
2150
+ if args.first:
2151
+ args.first = commands.resolve_label(state, args.first)
2152
+ import execution
2153
+
2154
+ if execution.command(args.codex_args):
2155
+ return await execution.run(args, state, binary)
2156
+ return await run_pool(args, state, binary)
2157
+
2158
+
2159
+ def selected_binary(args, state):
2160
+ if args.command == "_broker" and args.binary:
2161
+ return args.binary
2162
+ if (
2163
+ args.command == "runtime"
2164
+ and args.action in {"cache", "preflight"}
2165
+ and args.value
2166
+ ):
2167
+ return args.value
2168
+ if args.command == "_serve":
2169
+ job = state / "services" / (args.session + ".job.json")
2170
+ if job.exists():
2171
+ binary = json.loads(job.read_text()).get("binary")
2172
+ if binary:
2173
+ return binary
2174
+ if args.command in {
2175
+ "list",
2176
+ "status",
2177
+ "events",
2178
+ "usage",
2179
+ "config",
2180
+ "pause",
2181
+ "unpause",
2182
+ "retry",
2183
+ "wake",
2184
+ "stop",
2185
+ "switch",
2186
+ "restart",
2187
+ "menubar",
2188
+ "_supervise",
2189
+ "diagnostics",
2190
+ "clean",
2191
+ "workspace",
2192
+ "explain",
2193
+ "startup",
2194
+ }:
2195
+ return None
2196
+ if args.command == "group" and args.action != "start":
2197
+ return None
2198
+ if args.command == "runtime" and args.action in {"list", "use"}:
2199
+ return None
2200
+ if (
2201
+ args.command == "account"
2202
+ and args.action not in {"add", "reauth"}
2203
+ and not getattr(args, "revoke", False)
2204
+ ):
2205
+ return None
2206
+ try:
2207
+ return codex_path()
2208
+ except RouterError:
2209
+ import compatibility
2210
+ import runtime
2211
+
2212
+ for row in reversed(runtime.index(state)["runtimes"]):
2213
+ try:
2214
+ if compatibility.digest(row["binary"]) == row["sha256"]:
2215
+ return row["binary"]
2216
+ except OSError:
2217
+ continue
2218
+ raise
2219
+
2220
+
2221
+ def main():
2222
+ os.umask(0o077)
2223
+ argv = sys.argv[1:]
2224
+ if argv[:2] == ["help", "run"]:
2225
+ argv = ["run", "--help"]
2226
+ before_separator = argv[: argv.index("--")] if "--" in argv else argv
2227
+ if argv[:1] == ["run"] and any(
2228
+ flag in before_separator for flag in ("--help", "-h")
2229
+ ):
2230
+ result = subprocess.call([codex_path(), "--help"])
2231
+ print(
2232
+ "\nRelay launch options (before --):\n"
2233
+ " --first ACCOUNT Start with an enrolled account\n"
2234
+ " --profile NAME Select a Relay recovery profile\n"
2235
+ " --name NAME Name this Relay session\n"
2236
+ " --require-recovery Refuse modes outside account recovery\n\n"
2237
+ "Examples:\n relay run -- resume --yolo\n"
2238
+ " relay run -- -m MODEL -C PATH\n"
2239
+ ' relay run -- exec --json "Run the checks"\n\n'
2240
+ "Native Codex flags go after --. For native command help:\n"
2241
+ " relay run -- exec --help\n relay help codex resume"
2242
+ )
2243
+ return result
2244
+ if not argv or argv[:1] == ["help"]:
2245
+ topic = argv[1:] if argv else []
2246
+ if topic[:1] == ["codex"]:
2247
+ return subprocess.call([codex_path(), *topic[1:], "--help"])
2248
+ return parser().parse_args([*topic, "--help"])
2249
+ args = parse_arguments(argv)
2250
+ if args.command in {"run", "attach"}:
2251
+
2252
+ def interrupted(signum, frame):
2253
+ raise KeyboardInterrupt
2254
+
2255
+ signal.signal(signal.SIGTERM, interrupted)
2256
+ signal.signal(signal.SIGHUP, interrupted)
2257
+ if sys.platform != "darwin":
2258
+ raise RouterError(
2259
+ "This build uses macOS Keychain and currently supports macOS only."
2260
+ )
2261
+ state = private_dir(STATE)
2262
+ private_dir(state / "accounts")
2263
+ binary = selected_binary(args, state)
2264
+ if args.command == "list":
2265
+ import commands
2266
+
2267
+ value = commands.snapshot(state)
2268
+ print(
2269
+ "\n".join(
2270
+ f"{a['id']}: {a.get('email') or a['name']}" for a in value["accounts"]
2271
+ )
2272
+ or "No accounts enrolled."
2273
+ )
2274
+ return 0
2275
+ return asyncio.run(dispatch(args, state, binary))
2276
+
2277
+
2278
+ if __name__ == "__main__":
2279
+ try:
2280
+ sys.exit(main())
2281
+ except KeyboardInterrupt:
2282
+ sys.exit(130)
2283
+ except (RouterError, ValueError, asyncio.TimeoutError) as error:
2284
+ print(
2285
+ f"relay: {error if isinstance(error, (RouterError, ValueError)) else 'Codex operation timed out.'}",
2286
+ file=sys.stderr,
2287
+ )
2288
+ sys.exit(1)