multi-codex 0.7.0__py3-none-any.whl

multi_codex/launcher.py

@@ -0,0 +1,151 @@
+"""启动命令 `<bin_dir>/codex-<名称>` 的生成与识别。
+
+启动命令是由 config.json 推导出的产物，内容完全由 (账号名, 账号目录, 代理, 环境变量) 决定，
+所以“内容相同即 unchanged”就是幂等判断。第 2 行的受管标记用来区分本工具生成的文件
+和用户自己的同名文件：没有标记的文件一律不覆盖、不删除。
+"""
+
+import os
+import shlex
+import stat
+from typing import Dict, Optional
+
+from .config import PROXY_ENV_VARS, PROXY_INHERIT, PROXY_OFF, is_socks
+from .fsutil import KIND_FILE, entry_kind, read_text
+
+MARKER_PREFIX = "# managed-by: multi-codex account="
+LAUNCHER_PREFIX = "codex-"
+LAUNCHER_MODE = 0o755
+# 设置了账号环境变量的启动命令只允许本人读写执行：值里可能有 API key。
+PRIVATE_LAUNCHER_MODE = 0o700
+# `multi-codex run` 通过这个环境变量把脚本交给 sh，而不是放进命令行参数（其它用户能用 ps 看到）。
+RUN_SCRIPT_ENV = "MULTI_CODEX_RUN_SCRIPT"
+
+_PROXY_VARS = PROXY_ENV_VARS
+_NO_PROXY_HOSTS = "localhost,127.0.0.1,::1"
+
+
+def launcher_path(bin_dir: str, name: str) -> str:
+    return os.path.join(bin_dir, LAUNCHER_PREFIX + name)
+
+
+def launcher_mode(env: Optional[Dict[str, str]]) -> int:
+    return PRIVATE_LAUNCHER_MODE if env else LAUNCHER_MODE
+
+
+def launcher_file_ok(path: str, content: str, env: Optional[Dict[str, str]]) -> bool:
+    """启动命令是否与期望一致；_plan_launcher 与 launcher_status 共用，两处判断必须相同。
+
+    没有环境变量时只要求内容相同且可执行，与 0.3.0 一样不比较权限位：用户手工改成 0775，
+    或 bin_dir 所在文件系统忽略 chmod 时，不应被反复判为需要重写。
+    有环境变量时还要求组和其他用户没有任何权限，否则密钥会被其他用户读到。
+    """
+    if read_text(path) != content or not os.access(path, os.X_OK):
+        return False
+    if env:
+        try:
+            return stat.S_IMODE(os.stat(path).st_mode) & 0o077 == 0
+        except OSError:
+            return False
+    return True
+
+
+def render(name: str, account_dir: str, proxy: str, env: Optional[Dict[str, str]] = None,
+           command_mode: bool = False) -> str:
+    """生成 POSIX sh 启动命令的完整内容。所有值都经 shlex.quote 转义。
+
+    env 为空时，输出必须与 0.3.0 逐字节相同，否则升级后所有启动命令都会被判为 stale。
+    command_mode 供 `multi-codex run` 使用：环境设置与启动命令完全相同，只是最后执行 "$@" 而不是 codex，
+    这样两者不会各算一套环境变量。
+    """
+    q_dir = shlex.quote(account_dir)
+    lines = [
+        "#!/bin/sh",
+        MARKER_PREFIX + name,
+        "# Generated by multi-codex. Do not edit; use `multi-codex proxy` or `multi-codex apply`.",
+        "account_dir=" + q_dir,
+        # codex 遇到不存在的 CODEX_HOME 只会给出难懂的配置错误，这里提前给出明确提示。
+        'if [ ! -d "$account_dir" ]; then',
+        # 用 printf 而不是 echo：dash 的 echo 会解释路径里的反斜杠。
+        '  printf \'multi-codex: account directory does not exist: %s\\n\' "$account_dir" >&2',
+        "  exit 1",
+        "fi",
+        'CODEX_HOME="$account_dir"',
+        "export CODEX_HOME",
+    ]
+    lines.extend(_proxy_lines(proxy))
+    for key in sorted(env or {}):
+        # 键名已按 [A-Za-z_][A-Za-z0-9_]* 校验，可以直接写；值按字面引用，不经 shell 展开。
+        lines.append("{}={}; export {}".format(key, shlex.quote(env[key]), key))
+    if command_mode:
+        # 脚本经环境变量传入；先清掉它，被运行的命令不会继承这份含密钥的脚本。
+        lines.insert(1, "unset " + RUN_SCRIPT_ENV)
+        lines.extend(['exec "$@"', ""])
+        return "\n".join(lines)
+    lines.extend([
+        "if ! command -v codex >/dev/null 2>&1; then",
+        '  echo "multi-codex: codex not found in PATH" >&2',
+        "  exit 127",
+        "fi",
+        'exec codex "$@"',
+        "",
+    ])
+    return "\n".join(lines)
+
+
+def _proxy_lines(proxy: str) -> list:
+    if proxy == PROXY_INHERIT:
+        # 不动任何代理变量，沿用启动时 shell 里的值。
+        return []
+    if proxy == PROXY_OFF:
+        # 清掉全部 8 个变量；只是“不导出”挡不住从父进程继承来的代理。
+        return ["unset " + " ".join(_PROXY_VARS)]
+    q_proxy = shlex.quote(proxy)
+    lines = [
+        "HTTPS_PROXY={0}; HTTP_PROXY={0}; https_proxy={0}; http_proxy={0}".format(q_proxy),
+        "export HTTPS_PROXY HTTP_PROXY https_proxy http_proxy",
+    ]
+    if is_socks(proxy):
+        lines.append("ALL_PROXY={0}; all_proxy={0}; export ALL_PROXY all_proxy".format(q_proxy))
+    else:
+        # http 代理不写 ALL_PROXY：curl 等程序会把它用于所有协议。同时清掉继承来的值，
+        # 否则父进程里的另一个 ALL_PROXY 可能把部分流量带到别的代理上。
+        lines.append("unset ALL_PROXY all_proxy")
+    # 在用户原有的 NO_PROXY 后追加本地地址，不覆盖用户原来的内网例外；原值为空时不留前导逗号。
+    lines.append('NO_PROXY="${{NO_PROXY:+$NO_PROXY,}}{}"'.format(_NO_PROXY_HOSTS))
+    lines.append('no_proxy="${{no_proxy:+$no_proxy,}}{}"'.format(_NO_PROXY_HOSTS))
+    lines.append("export NO_PROXY no_proxy")
+    return lines
+
+
+def managed_account(path: str) -> Optional[str]:
+    """若 path 是本工具生成的启动命令，返回其中记录的账号名；否则返回 None。
+
+    只认普通文件：软链、目录或读不出文本的文件都视为“不归本工具管理”。
+    """
+    if entry_kind(path) != KIND_FILE:
+        return None
+    content = read_text(path)
+    if content is None:
+        return None
+    lines = content.split("\n", 2)
+    if len(lines) < 2 or not lines[1].startswith(MARKER_PREFIX):
+        return None
+    return lines[1][len(MARKER_PREFIX):].strip() or None
+
+
+def scan_managed(bin_dir: str) -> Dict[str, str]:
+    """列出 bin_dir 中所有受管启动命令，返回 {账号名: 路径}。bin_dir 不存在时返回空。"""
+    found: Dict[str, str] = {}
+    try:
+        names = os.listdir(bin_dir)
+    except OSError:
+        return found
+    for entry in sorted(names):
+        if not entry.startswith(LAUNCHER_PREFIX):
+            continue
+        path = os.path.join(bin_dir, entry)
+        account = managed_account(path)
+        if account is not None:
+            found[account] = path
+    return found

multi_codex/lock.py

@@ -0,0 +1,47 @@
+"""写操作互斥锁。
+
+用 fcntl.flock 实现：进程退出（包括被 SIGKILL）时内核自动释放锁，
+不会出现“持锁进程已死但锁还在”的情况，也不受 PID 复用影响。
+锁文件里的 PID 只用于报错时告诉用户是谁持有锁，不参与判断。
+"""
+
+import fcntl
+import os
+from typing import Optional
+
+from . import platform
+
+
+class LockBusyError(Exception):
+    def __init__(self, holder_pid: Optional[str]) -> None:
+        super().__init__(holder_pid)
+        self.holder_pid = holder_pid
+
+
+class WriteLock(object):
+    def __init__(self) -> None:
+        self.path = os.path.join(platform.state_dir(), "lock")
+        self._fd: Optional[int] = None
+
+    def __enter__(self) -> "WriteLock":
+        os.makedirs(os.path.dirname(self.path), exist_ok=True)
+        # 不能带 O_TRUNC 打开：此时还没拿到锁，截断会清掉持锁进程写的 PID。
+        # 锁文件也不能用 atomic_write 写或被删除：换了 inode，两个进程会各自锁住不同的文件。
+        fd = os.open(self.path, os.O_RDWR | os.O_CREAT, 0o600)
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except (BlockingIOError, PermissionError):
+            holder = os.read(fd, 64).decode("ascii", "replace").strip() or None
+            os.close(fd)
+            raise LockBusyError(holder)
+        os.ftruncate(fd, 0)
+        os.lseek(fd, 0, os.SEEK_SET)
+        os.write(fd, str(os.getpid()).encode("ascii"))
+        self._fd = fd
+        return self
+
+    def __exit__(self, *_exc) -> None:
+        if self._fd is not None:
+            fcntl.flock(self._fd, fcntl.LOCK_UN)
+            os.close(self._fd)
+            self._fd = None