Viewing files at commit 2310099b696dCommit detailsView latest

scripts/oakbench/environment.py 115 lines · 3.9 KB · python
1
"""Process environment policy for measured commands.
2
3
All measured subprocesses run through these environments so that author
4
identity, color suppression, and locale-ish noise are controlled in one place.
5
"""
6
7
from __future__ import annotations
8
9
import atexit
10
import os
11
import shutil
12
import tempfile
13
from pathlib import Path, PurePath
14
from typing import Any, Mapping
15
16
ENV_ISOLATION_VERSION = "2"
17
18
GIT_IDENTITY_KEYS = frozenset({"GIT_AUTHOR_NAME", "GIT_AUTHOR_EMAIL", "GIT_COMMITTER_NAME", "GIT_COMMITTER_EMAIL"})
19
20
21
def sanitize_git_environment(values: Mapping[str, str]) -> dict[str, str]:
22
    """Prevent host plumbing/configuration from redirecting measured commands.
23
24
    Explicit per-command scenario overrides may be applied after this boundary.
25
    Agent credential filtering is a separate, stricter policy.
26
    """
27
    env = {key: value for key, value in values.items()
28
           if not key.startswith("GIT_") or key in GIT_IDENTITY_KEYS}
29
    env.update(GIT_CONFIG_GLOBAL=os.devnull, GIT_CONFIG_NOSYSTEM="1", GIT_NO_REPLACE_OBJECTS="1")
30
    return env
31
32
_AUTO_BENCH_HOME: Path | None = None
33
_CLEANUP_REGISTERED = False
34
35
36
def _cleanup_auto_bench_home() -> None:
37
    if _AUTO_BENCH_HOME is not None:
38
        shutil.rmtree(_AUTO_BENCH_HOME, ignore_errors=True)
39
40
41
def _bench_home(home: Path | None = None) -> str:
42
    global _AUTO_BENCH_HOME, _CLEANUP_REGISTERED
43
    if home is None:
44
        path = Path(tempfile.gettempdir()) / f"oak-bench-home-{os.getpid()}"
45
        _AUTO_BENCH_HOME = path
46
        if not _CLEANUP_REGISTERED:
47
            atexit.register(_cleanup_auto_bench_home)
48
            _CLEANUP_REGISTERED = True
49
    else:
50
        path = home
51
    path.mkdir(mode=0o700, parents=True, exist_ok=True)
52
    (path / ".config").mkdir(exist_ok=True)
53
    (path / ".cache").mkdir(exist_ok=True)
54
    return str(path)
55
56
57
def base_env(
58
    author_name: str = "Oak Bench",
59
    author_email: str = "[email protected]",
60
    oak_author: str = "oakbench",
61
    home: Path | None = None,
62
) -> dict[str, str]:
63
    env = sanitize_git_environment(os.environ)
64
    bench_home = _bench_home(home)
65
    env.update(
66
        {
67
            "GIT_AUTHOR_NAME": author_name,
68
            "GIT_AUTHOR_EMAIL": author_email,
69
            "GIT_COMMITTER_NAME": author_name,
70
            "GIT_COMMITTER_EMAIL": author_email,
71
            "OAK_AUTHOR": oak_author,
72
            "NO_COLOR": "1",
73
            "CLICOLOR": "0",
74
            "HOME": bench_home,
75
            "XDG_CONFIG_HOME": str(Path(bench_home) / ".config"),
76
            "XDG_CACHE_HOME": str(Path(bench_home) / ".cache"),
77
            "ZDOTDIR": bench_home,
78
            "BASH_ENV": os.devnull,
79
            "ENV": os.devnull,
80
            "GIT_CONFIG_GLOBAL": os.devnull,
81
            "GIT_CONFIG_NOSYSTEM": "1",
82
            "LC_ALL": "C",
83
            "LANG": "C",
84
            "TZ": "UTC",
85
            "OAK_BENCH_ENV_ISOLATION_VERSION": ENV_ISOLATION_VERSION,
86
        }
87
    )
88
    return env
89
90
91
def probe_env(**kwargs: str) -> dict[str, str]:
92
    # The determinism probe intentionally does NOT suppress color: it measures
93
    # the tool's own non-TTY behavior. ANSI bytes in piped output are token
94
    # waste and break provider prompt caching.
95
    env = base_env(**kwargs)
96
    env.pop("NO_COLOR", None)
97
    env.pop("CLICOLOR", None)
98
    return env
99
100
101
def command_display(command: Any) -> str:
102
    """Human/token-countable form of a command or command sequence."""
103
    if (
104
        isinstance(command, list)
105
        and len(command) >= 3
106
        and command[0] in {"/bin/zsh", "zsh", "/bin/bash", "bash", "sh"}
107
        and command[1] in {"-c", "-lc"}
108
    ):
109
        return command[2]
110
    if isinstance(command, list) and command and all(isinstance(item, str) for item in command):
111
        executable = PurePath(command[0]).name if os.sep in command[0] else command[0]
112
        return " ".join([executable, *command[1:]])
113
    if isinstance(command, list):
114
        return " && ".join(command_display(item) for item in command)
115
    return str(command)