Skip to content

Commit e557a5c

Browse files
tatoalotatoalo
authored
fix(tasks): preserve refreshed GitHub credentials in sandboxes (#72540)
1 parent 0227a27 commit e557a5c

11 files changed

Lines changed: 453 additions & 157 deletions

File tree

‎products/tasks/backend/logic/services/agentsh.py‎

Lines changed: 108 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,16 @@
55

66
import yaml
77

8+
from products.tasks.backend.constants import SANDBOX_AGENT_LAUNCH_UNSET_ENV_VARS
9+
810
AGENTSH_DAEMON_PORT = 18080
911
SESSION_ID_FILE = "/tmp/agentsh-session-id"
1012
ENV_FILE = "/tmp/agent-env"
13+
GITHUB_ENV_FILE = "/tmp/agent-github-env"
14+
OAUTH_ENV_FILE = "/tmp/agent-oauth-env"
1115
ENV_WRAPPER_SCRIPT = "/tmp/agentsh-env-wrapper.sh"
1216
# Sourced via BASH_ENV on every `bash -c` the agent runs, so git/gh pick up a
13-
# mid-session GitHub credential refresh (the backend rewrites ENV_FILE in place).
17+
# mid-session GitHub credential refresh from its dedicated credential file.
1418
BASH_ENV_SCRIPT = "/tmp/agentsh-bash-env.sh"
1519
AGENTSH_AUDIT_DB = "/var/lib/agentsh/events.db"
1620
INFRASTRUCTURE_DOMAINS = [
@@ -93,7 +97,22 @@ def _get_debug_only_ports() -> list[int]:
9397
return ports
9498

9599

96-
def generate_env_wrapper() -> str:
100+
_MANAGED_CREDENTIAL_ENV_KEYS = ("GH_TOKEN", "GITHUB_TOKEN", "POSTHOG_PERSONAL_API_KEY")
101+
_EXCLUDED_AGENT_ENV_KEYS = (
102+
*SANDBOX_AGENT_LAUNCH_UNSET_ENV_VARS,
103+
"BASH_ENV",
104+
"PROMPT_COMMAND",
105+
"PYTHONSTARTUP",
106+
"PERL5OPT",
107+
"RUBYOPT",
108+
)
109+
110+
111+
def generate_env_wrapper(
112+
env_file: str = ENV_FILE,
113+
github_env_file: str = GITHUB_ENV_FILE,
114+
oauth_env_file: str = OAUTH_ENV_FILE,
115+
) -> str:
97116
"""Generate a wrapper that restores the full sandbox environment.
98117
99118
``agentsh exec`` starts child processes with a heavily stripped
@@ -104,25 +123,107 @@ def generate_env_wrapper() -> str:
104123
Network policy enforcement happens at the syscall level (ptrace) —
105124
it does not depend on proxy environment variables.
106125
"""
126+
quoted_env_file = shlex.quote(env_file)
127+
quoted_github_env_file = shlex.quote(github_env_file)
128+
quoted_oauth_env_file = shlex.quote(oauth_env_file)
129+
excluded_names = " ".join((*_MANAGED_CREDENTIAL_ENV_KEYS, *_EXCLUDED_AGENT_ENV_KEYS))
130+
excluded_entries = "|".join(f"{name}=*" for name in (*_MANAGED_CREDENTIAL_ENV_KEYS, *_EXCLUDED_AGENT_ENV_KEYS))
107131
return f"""\
108132
#!/bin/bash
133+
unset {excluded_names}
134+
while IFS= read -r -d $'\\0' line; do
135+
case "$line" in
136+
{excluded_entries}) ;;
137+
*) export "$line" ;;
138+
esac
139+
done < {quoted_env_file} 2>/dev/null
140+
141+
while IFS= read -r -d $'\\0' line; do
142+
case "$line" in
143+
GH_TOKEN=*|GITHUB_TOKEN=*) export "$line" ;;
144+
esac
145+
done < {quoted_github_env_file} 2>/dev/null
146+
109147
while IFS= read -r -d $'\\0' line; do
110-
export "$line"
111-
done < {ENV_FILE}
148+
case "$line" in
149+
POSTHOG_PERSONAL_API_KEY=*) export "$line" ;;
150+
esac
151+
done < {quoted_oauth_env_file} 2>/dev/null
112152
exec "$@"
113153
"""
114154

115155

116-
def generate_bash_env_script() -> str:
156+
def generate_bash_env_script(
157+
env_file: str = ENV_FILE,
158+
github_env_file: str = GITHUB_ENV_FILE,
159+
oauth_env_file: str = OAUTH_ENV_FILE,
160+
) -> str:
117161
"""
118-
Generate the script sourced via ``BASH_ENV``.
162+
Generate the script sourced via ``BASH_ENV`` and used to initialize its env file.
163+
164+
The explicit invocation runs before the background agent-server launch. It
165+
atomically replaces the full environment with the current sandbox process
166+
environment, excluding launch hooks and credentials. Credential files are
167+
initialized only when absent, so a backend refresh that happened before startup
168+
wins. Sourced invocations stay cheap and only export GitHub credentials.
119169
"""
170+
quoted_env_file = shlex.quote(env_file)
171+
quoted_github_env_file = shlex.quote(github_env_file)
172+
quoted_oauth_env_file = shlex.quote(oauth_env_file)
173+
excluded_entries = "|".join(f"{name}=*" for name in (*_MANAGED_CREDENTIAL_ENV_KEYS, *_EXCLUDED_AGENT_ENV_KEYS))
120174
return f"""\
175+
if [[ "${{BASH_SOURCE[0]}}" == "$0" ]]; then
176+
set -euo pipefail
177+
umask 077
178+
env_tmp="$(mktemp {quoted_env_file}.tmp.XXXXXX)"
179+
github_tmp="$(mktemp {quoted_github_env_file}.tmp.XXXXXX)"
180+
oauth_tmp="$(mktemp {quoted_oauth_env_file}.tmp.XXXXXX)"
181+
trap 'rm -f "$env_tmp" "$github_tmp" "$oauth_tmp"' EXIT
182+
183+
while IFS= read -r -d $'\\0' kv 2>/dev/null; do
184+
case "$kv" in
185+
{excluded_entries}) ;;
186+
*) printf '%s\\0' "$kv" >> "$env_tmp" ;;
187+
esac
188+
done < <(env -0)
189+
chmod 600 "$env_tmp"
190+
mv "$env_tmp" {quoted_env_file}
191+
192+
github_token="${{GITHUB_TOKEN:-${{GH_TOKEN:-}}}}"
193+
if [[ -n "$github_token" ]]; then
194+
printf 'GITHUB_TOKEN=%s\\0GH_TOKEN=%s\\0' "$github_token" "$github_token" > "$github_tmp"
195+
fi
196+
chmod 600 "$github_tmp"
197+
if [[ -e {quoted_github_env_file} || -L {quoted_github_env_file} ]]; then
198+
[[ -f {quoted_github_env_file} && ! -L {quoted_github_env_file} ]]
199+
chmod 600 {quoted_github_env_file}
200+
else
201+
if ! ln "$github_tmp" {quoted_github_env_file} 2>/dev/null; then
202+
[[ -f {quoted_github_env_file} && ! -L {quoted_github_env_file} ]]
203+
fi
204+
fi
205+
206+
if [[ -n "${{POSTHOG_PERSONAL_API_KEY:-}}" ]]; then
207+
printf 'POSTHOG_PERSONAL_API_KEY=%s\\0' "$POSTHOG_PERSONAL_API_KEY" > "$oauth_tmp"
208+
fi
209+
chmod 600 "$oauth_tmp"
210+
if [[ -e {quoted_oauth_env_file} || -L {quoted_oauth_env_file} ]]; then
211+
[[ -f {quoted_oauth_env_file} && ! -L {quoted_oauth_env_file} ]]
212+
chmod 600 {quoted_oauth_env_file}
213+
else
214+
if ! ln "$oauth_tmp" {quoted_oauth_env_file} 2>/dev/null; then
215+
[[ -f {quoted_oauth_env_file} && ! -L {quoted_oauth_env_file} ]]
216+
fi
217+
fi
218+
exit 0
219+
fi
220+
221+
unset GH_TOKEN GITHUB_TOKEN
121222
while IFS= read -r -d $'\\0' kv 2>/dev/null; do
122223
case "$kv" in
123224
GH_TOKEN=*|GITHUB_TOKEN=*) export "$kv" ;;
124225
esac
125-
done < {ENV_FILE} 2>/dev/null
226+
done < {quoted_github_env_file} 2>/dev/null
126227
"""
127228

128229

‎products/tasks/backend/logic/services/docker_sandbox.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434

3535
from .agentsh import (
3636
BASH_ENV_SCRIPT,
37-
ENV_FILE,
3837
ENV_WRAPPER_SCRIPT,
3938
SESSION_ID_FILE,
4039
build_exec_prefix,
@@ -852,16 +851,15 @@ def _build_agent_server_command(
852851
'export NO_PROXY="host.docker.internal,${NO_PROXY:-localhost,127.0.0.1}"; export no_proxy="$NO_PROXY"; '
853852
)
854853
inner = f"cd /scripts && {no_proxy_export}{server_cmd} > /tmp/agent-server.log 2>&1"
854+
initialize_env_file = f"bash {shlex.quote(BASH_ENV_SCRIPT)}"
855855

856856
if allowed_domains is not None:
857857
return (
858-
f"cd /scripts && env -0 > {ENV_FILE} && "
859-
f"{build_exec_prefix()} {ENV_WRAPPER_SCRIPT} bash -c {shlex.quote(inner)} &"
858+
f"cd /scripts && {initialize_env_file} && "
859+
f"({build_exec_prefix()} {ENV_WRAPPER_SCRIPT} bash -c {shlex.quote(inner)} &)"
860860
)
861861
else:
862-
# Write the env file even without agentsh so BASH_ENV (and the
863-
# in-process token resolver) can re-read a backend-refreshed token.
864-
return f"cd /scripts && env -0 > {ENV_FILE} && nohup {server_cmd} > /tmp/agent-server.log 2>&1 &"
862+
return f"cd /scripts && {initialize_env_file} && (nohup {server_cmd} > /tmp/agent-server.log 2>&1 &)"
865863

866864
def _launch_and_check(self, command: str) -> bool:
867865
"""Execute the agent-server command and wait for the health check.

‎products/tasks/backend/logic/services/local_packages.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818

1919
from django.conf import settings
2020

21+
import yaml
22+
2123
logger = logging.getLogger(__name__)
2224

2325
BUILD_OUTPUT_SUBDIR = "dist"
@@ -81,6 +83,7 @@ def get_local_posthog_code_packages() -> tuple[LocalPackage, ...] | None:
8183
def get_local_package_runtime_dependencies(packages: tuple[LocalPackage, ...]) -> dict[str, dict[str, str]]:
8284
"""Collect registry dependencies needed by the overlaid local package builds."""
8385
dependencies: dict[str, dict[str, str]] = {}
86+
dependency_catalogs: dict[Path, tuple[dict[str, str], dict[str, dict[str, str]]]] = {}
8487

8588
for package in packages:
8689
manifest_path = package.source_path / "package.json"
@@ -95,6 +98,57 @@ def get_local_package_runtime_dependencies(packages: tuple[LocalPackage, ...]) -
9598
raise ValueError(f"Expected dependency names and versions to be strings in {manifest_path}")
9699
if version.startswith("workspace:"):
97100
continue
101+
if version.startswith("catalog:"):
102+
workspace_manifest_path = next(
103+
(
104+
parent / "pnpm-workspace.yaml"
105+
for parent in package.source_path.parents
106+
if (parent / "pnpm-workspace.yaml").is_file()
107+
),
108+
None,
109+
)
110+
if workspace_manifest_path is None:
111+
raise ValueError(f"Could not resolve catalog dependency {name} from {manifest_path}")
112+
113+
if workspace_manifest_path not in dependency_catalogs:
114+
workspace_manifest = yaml.safe_load(workspace_manifest_path.read_text())
115+
catalog = workspace_manifest.get("catalog", {}) if isinstance(workspace_manifest, dict) else None
116+
catalogs = workspace_manifest.get("catalogs", {}) if isinstance(workspace_manifest, dict) else None
117+
if not isinstance(catalog, dict) or not all(
118+
isinstance(catalog_name, str) and isinstance(catalog_version, str)
119+
for catalog_name, catalog_version in catalog.items()
120+
):
121+
raise ValueError(
122+
f"Expected catalog to contain string dependencies in {workspace_manifest_path}"
123+
)
124+
if not isinstance(catalogs, dict) or not all(
125+
isinstance(catalog_name, str)
126+
and isinstance(named_catalog, dict)
127+
and all(
128+
isinstance(dependency_name, str) and isinstance(dependency_version, str)
129+
for dependency_name, dependency_version in named_catalog.items()
130+
)
131+
for catalog_name, named_catalog in catalogs.items()
132+
):
133+
raise ValueError(
134+
f"Expected catalogs to contain named string dependencies in {workspace_manifest_path}"
135+
)
136+
dependency_catalogs[workspace_manifest_path] = (catalog, catalogs)
137+
138+
catalog, catalogs = dependency_catalogs[workspace_manifest_path]
139+
catalog_reference = version.removeprefix("catalog:")
140+
if catalog_reference in {"", "*"}:
141+
selected_catalog = catalog
142+
else:
143+
named_catalog = catalogs.get(catalog_reference)
144+
if named_catalog is None:
145+
raise ValueError(f"Catalog {catalog_reference} is missing from {workspace_manifest_path}")
146+
selected_catalog = named_catalog
147+
148+
resolved_version = selected_catalog.get(name)
149+
if resolved_version is None:
150+
raise ValueError(f"Catalog dependency {name} is missing from {workspace_manifest_path}")
151+
version = resolved_version
98152
runtime_dependencies[name] = version
99153

100154
if runtime_dependencies:

‎products/tasks/backend/logic/services/modal_sandbox.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,6 @@
5555
from products.tasks.backend.logic.services.agentsh import (
5656
AGENTSH_DAEMON_PORT,
5757
BASH_ENV_SCRIPT,
58-
ENV_FILE,
5958
ENV_WRAPPER_SCRIPT,
6059
SESSION_ID_FILE,
6160
_hostname_from_url,
@@ -1000,14 +999,15 @@ def _build_agent_server_command(
1000999
server_cmd = f"bash -c {shlex.quote(wait_for_repo)}"
10011000

10021001
inner = f"cd /scripts && {server_cmd} > /tmp/agent-server.log 2>&1"
1002+
initialize_env_file = f"bash {shlex.quote(BASH_ENV_SCRIPT)}"
10031003

10041004
if allowed_domains is not None:
10051005
return (
1006-
f"cd /scripts && env -0 > {ENV_FILE} && "
1007-
f"{build_exec_prefix()} {ENV_WRAPPER_SCRIPT} bash -c {shlex.quote(inner)} &"
1006+
f"cd /scripts && {initialize_env_file} && "
1007+
f"({build_exec_prefix()} {ENV_WRAPPER_SCRIPT} bash -c {shlex.quote(inner)} &)"
10081008
)
10091009
else:
1010-
return f"cd /scripts && env -0 > {ENV_FILE} && nohup {server_cmd} > /tmp/agent-server.log 2>&1 &"
1010+
return f"cd /scripts && {initialize_env_file} && (nohup {server_cmd} > /tmp/agent-server.log 2>&1 &)"
10111011

10121012
def _diagnose_startup_failure(self, allowed_domains: list[str] | None) -> dict[str, str]:
10131013
diagnostics: dict[str, str] = {}

‎products/tasks/backend/logic/services/tests/test_local_packages.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
from products.tasks.backend.logic.services.local_packages import (
99
BUILD_OUTPUT_SUBDIR,
1010
PACKAGE_NAMES,
11+
LocalPackage,
12+
get_local_package_runtime_dependencies,
1113
get_local_posthog_code_packages,
1214
)
1315

@@ -67,3 +69,37 @@ def test_returns_packages_when_everything_present(self, fake_monorepo: Path) ->
6769
assert packages[0].sandbox_install_path == "/scripts/node_modules/@posthog/agent"
6870
assert packages[0].sandbox_build_output_path == "/scripts/node_modules/@posthog/agent/dist"
6971
assert packages[0].build_output_path == fake_monorepo / "packages" / "agent" / "dist"
72+
73+
74+
def test_resolves_default_catalog_runtime_dependencies(fake_monorepo: Path) -> None:
75+
(fake_monorepo / "pnpm-workspace.yaml").write_text(
76+
"catalog:\n"
77+
" catalog-runtime: 1.2.3\n"
78+
" star-catalog-runtime: 2.3.4\n"
79+
"catalogs:\n"
80+
" build:\n"
81+
" named-catalog-runtime: 3.4.5\n"
82+
)
83+
agent_source_path = fake_monorepo / "packages" / "agent"
84+
(agent_source_path / "package.json").write_text(
85+
'{"dependencies":{'
86+
'"catalog-runtime":"catalog:",'
87+
'"star-catalog-runtime":"catalog:*",'
88+
'"named-catalog-runtime":"catalog:build",'
89+
'"registry-runtime":"^4.5.6"'
90+
"}}"
91+
)
92+
package = LocalPackage(
93+
name="agent",
94+
source_path=agent_source_path,
95+
sandbox_install_path="/scripts/node_modules/@posthog/agent",
96+
)
97+
98+
assert get_local_package_runtime_dependencies((package,)) == {
99+
"agent": {
100+
"catalog-runtime": "1.2.3",
101+
"named-catalog-runtime": "3.4.5",
102+
"registry-runtime": "^4.5.6",
103+
"star-catalog-runtime": "2.3.4",
104+
}
105+
}

‎products/tasks/backend/logic/services/tests/test_modal_sandbox.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -487,7 +487,7 @@ def test_start_agent_server_wraps_with_agentsh_when_domains_provided(self, mock_
487487
command = _agent_server_launch_command(mock_sandbox.execute)
488488
assert "--createPr true" in command
489489
assert "agentsh exec --client-timeout 2h --timeout 2h" in command
490-
assert "env -0 > /tmp/agent-env" in command
490+
assert "bash /tmp/agentsh-bash-env.sh" in command
491491
assert "/tmp/agentsh-env-wrapper.sh" in command
492492
assert "./node_modules/.bin/agent-server" in command
493493

@@ -509,7 +509,7 @@ def test_start_agent_server_wraps_with_agentsh_when_domains_empty(self, mock_san
509509
command = _agent_server_launch_command(mock_sandbox.execute)
510510
assert "--allowedDomains" not in command
511511
assert "agentsh exec --client-timeout 2h --timeout 2h" in command
512-
assert "env -0 > /tmp/agent-env" in command
512+
assert "bash /tmp/agentsh-bash-env.sh" in command
513513

514514
@pytest.mark.parametrize(
515515
("create_pr", "expected_flag"),

0 commit comments

Comments
 (0)