Viewing files at commit 2735d7f5558eCommit detailsView latest

scripts/oakbench/environment.py 127 lines · 4.4 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 = "3"
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
def benchmark_process_environment(values: Mapping[str, str]) -> dict[str, str]:
33
    """Deterministic engine measurement policy, also enforced after agent filtering.
34
35
    Oak's documented updater otherwise performs a network request after the
36
    first successful command in a fresh HOME, biasing randomized subject order.
37
    Natural first-use updater UX requires a separate measurement contract.
38
    """
39
    env = sanitize_git_environment(values)
40
    env["OAK_NO_UPDATE_CHECK"] = "1"
41
    return env
42
43
44
_AUTO_BENCH_HOME: Path | None = None
45
_CLEANUP_REGISTERED = False
46
47
48
def _cleanup_auto_bench_home() -> None:
49
    if _AUTO_BENCH_HOME is not None:
50
        shutil.rmtree(_AUTO_BENCH_HOME, ignore_errors=True)
51
52
53
def _bench_home(home: Path | None = None) -> str:
54
    global _AUTO_BENCH_HOME, _CLEANUP_REGISTERED
55
    if home is None:
56
        path = Path(tempfile.gettempdir()) / f"oak-bench-home-{os.getpid()}"
57
        _AUTO_BENCH_HOME = path
58
        if not _CLEANUP_REGISTERED:
59
            atexit.register(_cleanup_auto_bench_home)
60
            _CLEANUP_REGISTERED = True
61
    else:
62
        path = home
63
    path.mkdir(mode=0o700, parents=True, exist_ok=True)
64
    (path / ".config").mkdir(exist_ok=True)
65
    (path / ".cache").mkdir(exist_ok=True)
66
    return str(path)
67
68
69
def base_env(
70
    author_name: str = "Oak Bench",
71
    author_email: str = "[email protected]",
72
    oak_author: str = "oakbench",
73
    home: Path | None = None,
74
) -> dict[str, str]:
75
    env = benchmark_process_environment(os.environ)
76
    bench_home = _bench_home(home)
77
    env.update(
78
        {
79
            "GIT_AUTHOR_NAME": author_name,
80
            "GIT_AUTHOR_EMAIL": author_email,
81
            "GIT_COMMITTER_NAME": author_name,
82
            "GIT_COMMITTER_EMAIL": author_email,
83
            "OAK_AUTHOR": oak_author,
84
            "NO_COLOR": "1",
85
            "CLICOLOR": "0",
86
            "HOME": bench_home,
87
            "XDG_CONFIG_HOME": str(Path(bench_home) / ".config"),
88
            "XDG_CACHE_HOME": str(Path(bench_home) / ".cache"),
89
            "ZDOTDIR": bench_home,
90
            "BASH_ENV": os.devnull,
91
            "ENV": os.devnull,
92
            "GIT_CONFIG_GLOBAL": os.devnull,
93
            "GIT_CONFIG_NOSYSTEM": "1",
94
            "LC_ALL": "C",
95
            "LANG": "C",
96
            "TZ": "UTC",
97
            "OAK_BENCH_ENV_ISOLATION_VERSION": ENV_ISOLATION_VERSION,
98
        }
99
    )
100
    return env
101
102
103
def probe_env(**kwargs: str) -> dict[str, str]:
104
    # The determinism probe intentionally does NOT suppress color: it measures
105
    # the tool's own non-TTY behavior. ANSI bytes in piped output are token
106
    # waste and break provider prompt caching.
107
    env = base_env(**kwargs)
108
    env.pop("NO_COLOR", None)
109
    env.pop("CLICOLOR", None)
110
    return env
111
112
113
def command_display(command: Any) -> str:
114
    """Human/token-countable form of a command or command sequence."""
115
    if (
116
        isinstance(command, list)
117
        and len(command) >= 3
118
        and command[0] in {"/bin/zsh", "zsh", "/bin/bash", "bash", "sh"}
119
        and command[1] in {"-c", "-lc"}
120
    ):
121
        return command[2]
122
    if isinstance(command, list) and command and all(isinstance(item, str) for item in command):
123
        executable = PurePath(command[0]).name if os.sep in command[0] else command[0]
124
        return " ".join([executable, *command[1:]])
125
    if isinstance(command, list):
126
        return " && ".join(command_display(item) for item in command)
127
    return str(command)