tenant.openagents/omega

1"""Pier-native agent wrapper for Zed's eval-cli binary.
2
3Pier (a Harbor fork) is the harness for DeepSWE, whose tasks run air-gapped
4(`allow_internet = false`). Pier's installed-agent interface differs from
5Harbor's — it requires `install_spec()` and, crucially, a `network_allowlist()`
6the egress proxy honors so the agent can still reach the model API. Harbor's
7`BaseInstalledAgent` has neither, so this is a separate, Pier-specific class
8(the sanctioned cross-framework exception): everything else in the project has a
9single canonical implementation, but two genuinely different harness SDKs need
10two thin agent shells.
11
12Air-gap notes:
13  - The eval-cli binary is a static musl executable copied from the mounted Modal
14    volume (`/data`, local — not network), so installing it needs no egress.
15  - Network installs (Node, LSPs) are intentionally skipped: there's no egress to
16    fetch them, and eval-cli's core tools (read/edit/terminal/...) don't need
17    them. Only the model API host is allowlisted.
18"""
19
20from __future__ import annotations
21
22import json
23import shlex
24from urllib.parse import urlparse
25
26from pier.agents.installed.base import BaseInstalledAgent, with_prompt_template
27from pier.environments.base import BaseEnvironment
28from pier.models.agent.context import AgentContext
29from pier.models.agent.install import AgentInstallSpec, InstallStep
30from pier.models.agent.network import NetworkAllowlist
31from pier.models.trial.paths import EnvironmentPaths
32
33from .agent_common import (
34    add_anthropic_available_models_env,
35    add_openai_compatible_provider_env,
36    add_zed_eval_env,
37    detect_workdir,
38    eval_cli_with_log_command,
39    patch_command,
40    populate_context_from_result,
41    provider_api_env,
42)
43
44# Default model-provider API hosts the agent must reach even on air-gapped tasks.
45# Mirrors benchmarks.AGENT_API_HOSTS; duplicated here so this module has no
46# dependency on the orchestration package when loaded inside a sandbox.
47DEFAULT_PROVIDER_DOMAINS: dict[str, list[str]] = {
48    "anthropic": ["api.anthropic.com"],
49    "openai": ["api.openai.com"],
50    "google": [".googleapis.com"],
51    "gemini": [".googleapis.com"],
52    "deepseek": ["api.deepseek.com"],
53    "baseten": ["inference.baseten.co"],
54}
55# Base-URL env vars Pier-style allowlist resolution should also honor.
56BASE_URL_ENV_VARS = (
57    "ANTHROPIC_BASE_URL",
58    "OPENAI_BASE_URL",
59    "OPENAI_API_BASE",
60    "GEMINI_API_BASE",
61)
62
63
64class ZedPierAgent(BaseInstalledAgent):
65    """Runs Zed's headless eval-cli binary under Pier."""
66
67    SUPPORTS_ATIF: bool = False
68
69    @staticmethod
70    def name() -> str:
71        return "zed"
72
73    def _container_binary_path(self) -> str:
74        path = self._get_env("EVAL_CLI_CONTAINER_PATH")
75        if not path:
76            raise ValueError(
77                "ZedPierAgent requires EVAL_CLI_CONTAINER_PATH (the eval-cli binary "
78                "on the mounted volume, e.g. /data/builds/<id>/eval-cli)"
79            )
80        return path
81
82    def install_spec(self) -> AgentInstallSpec:
83        # The binary lives on the runtime-mounted volume, which isn't available at
84        # image-build time, so the real placement happens in run() via
85        # exec_as_root. install_spec only needs a non-empty, no-egress step.
86        return AgentInstallSpec(
87            agent_name=self.name(),
88            version=self._version,
89            steps=[
90                InstallStep(
91                    user="root",
92                    run="true  # eval-cli is copied from the volume at run time",
93                )
94            ],
95        )
96
97    def network_allowlist(self) -> NetworkAllowlist:
98        """Hosts the egress proxy must allow so eval-cli can reach the model API
99        on air-gapped DeepSWE tasks."""
100        domains: list[str] = []
101
102        provider = None
103        if self.model_name and "/" in self.model_name:
104            provider = self.model_name.split("/", 1)[0]
105        domains.extend(DEFAULT_PROVIDER_DOMAINS.get(provider or "", []))
106
107        # Any explicitly configured base URL (e.g. an OpenAI-compatible gateway
108        # such as Baseten) wins/adds alongside the provider default.
109        for env_var in BASE_URL_ENV_VARS:
110            value = self._get_env(env_var)
111            if value:
112                host = urlparse(value).hostname
113                if host:
114                    domains.append(host)
115
116        # A Baseten (or other OpenAI-compatible) provider is wired via this JSON;
117        # allow its api_url host too.
118        providers_json = self._get_env("ZED_OPENAI_COMPATIBLE_PROVIDERS")
119        if providers_json:
120            try:
121                parsed = json.loads(providers_json)
122            except json.JSONDecodeError:
123                parsed = {}
124            for provider_config in (parsed or {}).values():
125                api_url = (
126                    provider_config.get("api_url")
127                    if isinstance(provider_config, dict)
128                    else None
129                )
130                host = urlparse(api_url).hostname if api_url else None
131                if host:
132                    domains.append(host)
133
134        if not domains:
135            # Never return an empty allowlist: without the model host the agent
136            # can do nothing on an air-gapped task. Default to Anthropic.
137            domains = ["api.anthropic.com"]
138        return NetworkAllowlist(domains=domains)
139
140    def _api_env(self) -> dict[str, str]:
141        env = provider_api_env(self.model_name, self._get_env)
142        add_openai_compatible_provider_env(
143            env, self._get_env("ZED_OPENAI_COMPATIBLE_PROVIDERS")
144        )
145        add_anthropic_available_models_env(
146            env, self._get_env("ZED_ANTHROPIC_AVAILABLE_MODELS")
147        )
148        add_zed_eval_env(
149            env, self._extra_env, exclude={"ZED_EVAL_INSTRUCTION_SUFFIX_FILE"}
150        )
151        return env
152
153    async def _detect_workdir(self, environment: BaseEnvironment) -> str:
154        return await detect_workdir(
155            environment,
156            self.exec_as_agent,
157            self._get_env,
158            "Could not detect a working directory; set EVAL_CLI_WORKDIR via --ae",
159        )
160
161    @with_prompt_template
162    async def run(
163        self, instruction: str, environment: BaseEnvironment, context: AgentContext
164    ) -> None:
165        # Place the static binary from the mounted volume (no egress needed).
166        await self.exec_as_root(
167            environment,
168            command=(
169                f"install -m755 {shlex.quote(self._container_binary_path())} "
170                "/usr/local/bin/eval-cli && eval-cli --help >/dev/null"
171            ),
172        )
173
174        workdir = await self._detect_workdir(environment)
175        agent_dir = str(EnvironmentPaths.agent_dir)
176        env = self.build_process_env(self._api_env())
177
178        parts = [
179            "eval-cli",
180            f"--workdir {shlex.quote(workdir)}",
181            f"--output-dir {shlex.quote(agent_dir)}",
182        ]
183        if self.model_name:
184            parts.append(f"--model {shlex.quote(self.model_name)}")
185        timeout = self._get_env("EVAL_CLI_TIMEOUT")
186        if timeout:
187            parts.append(f"--timeout {shlex.quote(timeout)}")
188        instruction_suffix_file = self._get_env("ZED_EVAL_INSTRUCTION_SUFFIX_FILE")
189        if instruction_suffix_file:
190            parts.append(
191                f"--instruction-suffix-file {shlex.quote(instruction_suffix_file)}"
192            )
193        staff = self._get_env("EVAL_CLI_STAFF")
194        if staff and staff.lower() == "false":
195            parts.append("--no-staff")
196        reasoning_effort = self._get_env("EVAL_CLI_REASONING_EFFORT")
197        if reasoning_effort:
198            parts.append(f"--reasoning-effort {shlex.quote(reasoning_effort)}")
199        enable_thinking = self._get_env("EVAL_CLI_ENABLE_THINKING")
200        if enable_thinking:
201            if enable_thinking.lower() == "true":
202                parts.append("--thinking true")
203            elif enable_thinking.lower() == "false":
204                parts.append("--thinking false")
205        parts.append(f"--instruction {shlex.quote(instruction)}")
206
207        # Tolerate exit 2 (timeout): deliver whatever the agent produced.
208        await self.exec_as_agent(
209            environment,
210            command=eval_cli_with_log_command(
211                parts,
212                agent_dir + "/eval-cli.txt",
213                timeout_message=None,
214            ),
215            env=env,
216        )
217
218        # SWE-bench-style patch for the verifier (DeepSWE tasks are git repos).
219        await self.exec_as_agent(
220            environment,
221            command=patch_command(agent_dir),
222            cwd=workdir,
223        )
224
225    def populate_context_post_run(self, context: AgentContext) -> None:
226        populate_context_from_result(self.logs_dir, context, self.logger)
227