diff --git a/products/tasks/backend/constants.py b/products/tasks/backend/constants.py index b329cc792fed..bcd9d0737352 100644 --- a/products/tasks/backend/constants.py +++ b/products/tasks/backend/constants.py @@ -391,6 +391,14 @@ def vm_sandbox_allowed_origins(*, distinct_id: str, organization_id: str) -> set } ) +# Stripped from the agent-server's process environment at launch (env -u). +# Two categories: +# - code-injection vectors a resume snapshot could smuggle in (NODE_*, LD_*, DYLD_*); +# - the GitHub token, so the agent-server holds no frozen copy of the acting user's +# credentials. The token is delivered per command via the live /tmp/agent-env file +# (re-sourced by BASH_ENV, seeded before this unset), so git/gh still authenticate; +# removing the static process-env copy is what lets a mid-session logout or rebind +# actually take effect instead of being resurrected from os.environ. SANDBOX_AGENT_LAUNCH_UNSET_ENV_VARS: tuple[str, ...] = ( "NODE_OPTIONS", "NODE_REPL_EXTERNAL_MODULE", @@ -399,6 +407,8 @@ def vm_sandbox_allowed_origins(*, distinct_id: str, organization_id: str) -> set "LD_AUDIT", "DYLD_INSERT_LIBRARIES", "DYLD_LIBRARY_PATH", + "GITHUB_TOKEN", + "GH_TOKEN", ) diff --git a/products/tasks/backend/logic/services/agentsh.py b/products/tasks/backend/logic/services/agentsh.py index 058ec7cb42a3..91f5bf3a8641 100644 --- a/products/tasks/backend/logic/services/agentsh.py +++ b/products/tasks/backend/logic/services/agentsh.py @@ -1,4 +1,5 @@ import shlex +from pathlib import Path from urllib.parse import urlparse from django.conf import settings @@ -16,6 +17,19 @@ # Sourced via BASH_ENV on every `bash -c` the agent runs, so git/gh pick up a # mid-session GitHub credential refresh from its dedicated credential file. BASH_ENV_SCRIPT = "/tmp/agentsh-bash-env.sh" + +# The gh PATH shim (first on PATH; sources the credential script so gh authenticates as the current +# actor in any shell mode). Baked into new base images, but also installed at runtime so resumes +# from pre-shim filesystem snapshots — and any window where the image lags this backend — still +# deliver the token to gh. Read from its single source of truth (the file the Dockerfiles COPY). +GH_GUARD_INSTALL_PATH = "/opt/posthog/bin/gh" +_GH_GUARD_SOURCE_PATH = Path(__file__).resolve().parents[2] / "sandbox" / "images" / "gh-guard.sh" + + +def read_gh_guard_script() -> bytes: + return _GH_GUARD_SOURCE_PATH.read_bytes() + + AGENTSH_AUDIT_DB = "/var/lib/agentsh/events.db" INFRASTRUCTURE_DOMAINS = [ "*.posthog.com", diff --git a/products/tasks/backend/logic/services/docker_sandbox.py b/products/tasks/backend/logic/services/docker_sandbox.py index afcac19487b2..3fe76eac06e8 100644 --- a/products/tasks/backend/logic/services/docker_sandbox.py +++ b/products/tasks/backend/logic/services/docker_sandbox.py @@ -35,6 +35,7 @@ from .agentsh import ( BASH_ENV_SCRIPT, ENV_WRAPPER_SCRIPT, + GH_GUARD_INSTALL_PATH, SESSION_ID_FILE, build_exec_prefix, build_setup_script, @@ -42,6 +43,7 @@ generate_config_yaml, generate_env_wrapper, generate_policy_yaml, + read_gh_guard_script, ) from .local_skills import ENV_LOCAL_SKILLS_HOST_PATH, LocalSkillsCache from .sandbox import ( @@ -874,6 +876,15 @@ def _launch_and_check(self, command: str) -> bool: return False return self._wait_for_health_check(max_attempts=20) + def _install_gh_guard(self) -> None: + """Install the gh PATH shim at runtime so it's present regardless of image age. + + New base images bake it in, but a resume from a pre-shim filesystem snapshot (or any window + where the image lags this backend) would otherwise lack it, leaving gh with no token once the + frozen launch-env token is unset.""" + self.write_file(GH_GUARD_INSTALL_PATH, read_gh_guard_script()) + self.execute(f"chmod +x {shlex.quote(GH_GUARD_INSTALL_PATH)}", timeout_seconds=30) + def start_agent_server( self, repository: str | None, @@ -920,6 +931,7 @@ def start_agent_server( # mid-session credential refreshes reach git/gh. Needed for both agentsh # and non-agentsh runs. self.write_file(BASH_ENV_SCRIPT, generate_bash_env_script().encode()) + self._install_gh_guard() if allowed_domains is not None: self._setup_agentsh(WORKING_DIR, allowed_domains) diff --git a/products/tasks/backend/logic/services/modal_sandbox.py b/products/tasks/backend/logic/services/modal_sandbox.py index 61f3b55f87c7..9ef32c2b8674 100644 --- a/products/tasks/backend/logic/services/modal_sandbox.py +++ b/products/tasks/backend/logic/services/modal_sandbox.py @@ -56,6 +56,7 @@ AGENTSH_DAEMON_PORT, BASH_ENV_SCRIPT, ENV_WRAPPER_SCRIPT, + GH_GUARD_INSTALL_PATH, SESSION_ID_FILE, _hostname_from_url, build_exec_prefix, @@ -64,6 +65,7 @@ generate_config_yaml, generate_env_wrapper, generate_policy_yaml, + read_gh_guard_script, ) from products.tasks.backend.logic.services.local_packages import ( get_local_package_runtime_dependencies, @@ -185,6 +187,7 @@ def _resource_create_kwargs(config: SandboxConfig) -> dict[str, object]: } LOCAL_MODAL_INSTALL_SKILLS_SCRIPT = Path("products/tasks/backend/sandbox/images/install-skills.sh") LOCAL_MODAL_GIT_GUARD_SCRIPT = Path("products/tasks/backend/sandbox/images/git-guard.sh") +LOCAL_MODAL_GH_GUARD_SCRIPT = Path("products/tasks/backend/sandbox/images/gh-guard.sh") _image_ref_cache: TTLCache = TTLCache(maxsize=3, ttl=300) @@ -445,11 +448,14 @@ def _prepare_local_modal_build_context(template: SandboxTemplate) -> tuple[str, destination_dockerfile_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(source_dockerfile_path, destination_dockerfile_path) - # Both base and notebook Dockerfiles COPY the git guard, so include it in - # every local build context. + # Both base and notebook Dockerfiles COPY the git and gh guards, so include + # them in every local build context. destination_git_guard_path = context_dir / LOCAL_MODAL_GIT_GUARD_SCRIPT destination_git_guard_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(base_dir / LOCAL_MODAL_GIT_GUARD_SCRIPT, destination_git_guard_path) + destination_gh_guard_path = context_dir / LOCAL_MODAL_GH_GUARD_SCRIPT + destination_gh_guard_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(base_dir / LOCAL_MODAL_GH_GUARD_SCRIPT, destination_gh_guard_path) if template == SandboxTemplate.DEFAULT_BASE: source_install_script_path = base_dir / LOCAL_MODAL_INSTALL_SKILLS_SCRIPT @@ -1101,6 +1107,11 @@ def start_agent_server( repo_path = f"/tmp/workspace/repos/{org}/{repo}" self.write_file(BASH_ENV_SCRIPT, generate_bash_env_script().encode()) + # Install the gh shim at runtime too (see agentsh.GH_GUARD_INSTALL_PATH): a resume from a + # pre-shim filesystem snapshot — or any window where the base image lags this backend — + # would otherwise leave gh with no token once the frozen launch-env token is unset. + self.write_file(GH_GUARD_INSTALL_PATH, read_gh_guard_script()) + self.execute(f"chmod +x {shlex.quote(GH_GUARD_INSTALL_PATH)}", timeout_seconds=30) if allowed_domains is not None: self._setup_agentsh(WORKING_DIR, allowed_domains) diff --git a/products/tasks/backend/logic/services/tests/test_modal_sandbox.py b/products/tasks/backend/logic/services/tests/test_modal_sandbox.py index 679265c42090..898ae7340135 100644 --- a/products/tasks/backend/logic/services/tests/test_modal_sandbox.py +++ b/products/tasks/backend/logic/services/tests/test_modal_sandbox.py @@ -641,6 +641,8 @@ def test_start_agent_server_raises_on_health_check_failure(self, mock_sandbox: A mock_sandbox.execute = MagicMock( side_effect=[ ExecutionResult(stdout="", stderr="", exit_code=0, error=None), + ExecutionResult(stdout="", stderr="", exit_code=0, error=None), # gh shim write (mv) + ExecutionResult(stdout="", stderr="", exit_code=0, error=None), # gh shim chmod ExecutionResult(stdout="", stderr="", exit_code=0, error=None), # --posthogExecPermissionRegex probe ExecutionResult(stdout="", stderr="", exit_code=1, error=None), ExecutionResult(stdout="some log output", stderr="", exit_code=0, error=None), @@ -696,6 +698,8 @@ def test_start_agent_server_frees_port_before_relaunch(self, mock_sandbox: Any): mock_sandbox.execute = MagicMock( side_effect=[ ExecutionResult(stdout="", stderr="", exit_code=0, error=None), + ExecutionResult(stdout="", stderr="", exit_code=0, error=None), # gh shim write (mv) + ExecutionResult(stdout="", stderr="", exit_code=0, error=None), # gh shim chmod ExecutionResult(stdout="", stderr="", exit_code=0, error=None), # --posthogExecPermissionRegex probe ExecutionResult(stdout="", stderr="", exit_code=0, error=None), ExecutionResult(stdout="ok:1", stderr="", exit_code=0, error=None), diff --git a/products/tasks/backend/sandbox/images/Dockerfile.sandbox-base b/products/tasks/backend/sandbox/images/Dockerfile.sandbox-base index 2a847c849d6c..736be0530dee 100644 --- a/products/tasks/backend/sandbox/images/Dockerfile.sandbox-base +++ b/products/tasks/backend/sandbox/images/Dockerfile.sandbox-base @@ -158,6 +158,10 @@ RUN git config --global user.email "code@posthog.com" && \ # Block `git commit`/`git push` so unsigned commits cannot leave the sandbox COPY products/tasks/backend/sandbox/images/git-guard.sh /opt/posthog/bin/git RUN chmod +x /opt/posthog/bin/git +# Inject the per-actor GitHub token into every `gh` call (BASH_ENV only reaches +# non-interactive `bash -c`; the agent's interactive shell needs this shim). +COPY products/tasks/backend/sandbox/images/gh-guard.sh /opt/posthog/bin/gh +RUN chmod +x /opt/posthog/bin/gh ENV PATH="/opt/posthog/bin:${PATH}" # This is required for the Claude Code SDK to allow --dangerously-skip-permissions as the root user diff --git a/products/tasks/backend/sandbox/images/Dockerfile.sandbox-notebook b/products/tasks/backend/sandbox/images/Dockerfile.sandbox-notebook index a0faa1ee9c1d..53594108094d 100644 --- a/products/tasks/backend/sandbox/images/Dockerfile.sandbox-notebook +++ b/products/tasks/backend/sandbox/images/Dockerfile.sandbox-notebook @@ -91,6 +91,10 @@ RUN git config --global user.email "code@posthog.com" && \ # Block `git commit`/`git push` so unsigned commits cannot leave the sandbox. COPY products/tasks/backend/sandbox/images/git-guard.sh /opt/posthog/bin/git RUN chmod +x /opt/posthog/bin/git +# Inject the per-actor GitHub token into every `gh` call (BASH_ENV only reaches +# non-interactive `bash -c`; the agent's interactive shell needs this shim). +COPY products/tasks/backend/sandbox/images/gh-guard.sh /opt/posthog/bin/gh +RUN chmod +x /opt/posthog/bin/gh ENV PATH="/opt/posthog/bin:${PATH}" # This is required for the Claude Code SDK to allow --dangerously-skip-permissions as the root user diff --git a/products/tasks/backend/sandbox/images/gh-guard.sh b/products/tasks/backend/sandbox/images/gh-guard.sh new file mode 100644 index 000000000000..042a6c118d94 --- /dev/null +++ b/products/tasks/backend/sandbox/images/gh-guard.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# gh guard. +# +# Installed first on PATH inside the cloud sandbox image as /opt/posthog/bin/gh. +# The backend delivers the per-actor GitHub token via BASH_ENV, which only +# non-interactive `bash -c` honors — the agent runs its tool commands in an +# interactive shell, so `gh` there would otherwise have no token. This shim +# sources the same credential script the shells do, so `gh` authenticates as the +# current actor regardless of shell mode, and honors logout (an emptied file +# exports nothing, leaving gh unauthenticated rather than falling back to a stale +# token). All arguments pass straight through to the real gh. + +native_gh="" +for candidate in /usr/bin/gh /usr/local/bin/gh /bin/gh; do + if [ -x "$candidate" ] && [ "$candidate" != "/opt/posthog/bin/gh" ]; then + native_gh="$candidate" + break + fi +done +if [ -z "$native_gh" ]; then + echo "gh-guard: could not locate the real gh binary" >&2 + exit 127 +fi + +# Re-source the backend-managed credentials fresh on every call (the file is +# rewritten on each refresh / actor transition). The script's sourced branch +# unsets then re-exports GH_TOKEN/GITHUB_TOKEN from the env file. +if [ -f /tmp/agentsh-bash-env.sh ]; then + # shellcheck source=/dev/null + . /tmp/agentsh-bash-env.sh +fi + +exec "$native_gh" "$@" diff --git a/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py b/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py index 3086ef7cbdd0..45fbf91cbeb2 100644 --- a/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py @@ -9,9 +9,12 @@ from temporalio import activity from temporalio.exceptions import ApplicationError +from posthog.models.integration import Integration +from posthog.models.user_integration import ReauthorizationRequired, UserIntegration from posthog.temporal.common.utils import close_db_connections from posthog.temporal.oauth import PosthogMcpScopes +from products.tasks.backend.exceptions import CredentialUnavailableError from products.tasks.backend.logic.services.agent_command import ( FOLLOWUP_TIMEOUT_SECONDS, REFRESH_TIMEOUT_SECONDS, @@ -27,15 +30,25 @@ from products.tasks.backend.models import TaskRun from products.tasks.backend.redis import get_tasks_stream_redis_sync, run_uses_dedicated_stream from products.tasks.backend.temporal.oauth import create_oauth_access_token_for_run +from products.tasks.backend.temporal.process_task.sandbox_credentials import ( + apply_github_credentials_to_sandbox, + clear_github_credentials_from_sandbox, + sandbox_credential_lock, +) from products.tasks.backend.temporal.process_task.utils import ( + PrAuthorshipMode, get_actor_distinct_id, get_imported_mcp_server_configs, + get_pr_authorship_mode, + get_sandbox_github_identity_user, + get_sandbox_github_token, get_sandbox_mcp_session_user, get_sandbox_ph_mcp_configs, get_task_run_credential_user, get_user_mcp_server_configs, is_slack_interaction_state, loop_mcp_installation_allowlist, + mark_sandbox_github_identity, mark_sandbox_mcp_session, record_message_actor, sandbox_identity_scope, @@ -208,6 +221,13 @@ def _deliver_followup(input: SendFollowupToSandboxInput) -> str | None: ): error_msg = "Could not rebind sandbox MCP credentials for the follow-up actor" raise RuntimeError(f"send_followup failed: {error_msg}") + + # Bind the sandbox's GitHub credentials to this actor: rebind if they have + # access, otherwise log out so the previous actor's identity can't be used. + # Fail closed only if we can't even clear the prior credentials. + if not _refresh_sandbox_github(task_run, actor_user, state): + error_msg = "Could not rebind or clear sandbox GitHub credentials for the follow-up actor" + raise RuntimeError(f"send_followup failed: {error_msg}") artifacts = None artifact_ids = input.artifact_ids or [] if artifact_ids: @@ -431,6 +451,144 @@ def _refresh_sandbox_mcp( return False # rebind never confirmed → fail closed (unknown binding may hide a live session) +def _resolve_live_sandbox(state: dict[str, Any] | None) -> Any: + """The running Sandbox handle for a run's state, or None when unavailable. + + GitHub credentials are written into the sandbox directly (git remote + env + file), so the gate needs the handle. Absent/dead sandbox → None; the + periodic credential-refresh loop reconciles identity in that case. + """ + sandbox_id = (state or {}).get("sandbox_id") + if not sandbox_id: + return None + from products.tasks.backend.logic.services.sandbox import ( + Sandbox, # noqa: PLC0415 — keep the sandbox service off the import path + ) + + try: + sandbox = Sandbox.get_by_id(sandbox_id) + return sandbox if sandbox.is_running() else None + except Exception: + # This None drives a fail-closed follow-up rejection, so keep the cause: it + # distinguishes a genuinely dead sandbox from a transient control-plane lookup + # error, which have different remediation. + logger.warning("resolve_live_sandbox_failed", sandbox_id=sandbox_id, exc_info=True) + return None + + +def _refresh_sandbox_github(task_run: TaskRun, actor_user: Any, state: dict[str, Any] | None) -> bool: + """Bind the sandbox's in-place GitHub credentials to this message's actor. + + On an actor transition: re-inject the new actor's token if they have usable + access, otherwise log the sandbox out (strip the token from the git remote + and env) so the previous actor's GitHub identity can never be used by a + follow-up actor who lacks access. Reauthorization for that actor is surfaced + by the existing credential-refresh path, unchanged. + + Only USER-authored runs carry per-actor identity — BOT runs share one + installation token, so every actor is already the same identity. This + enforces the transition boundary; the periodic credential-refresh loop + keeps a continuous actor's token rotated between transitions. + + Returns ``True`` when the sandbox safely reflects this actor (rebound, logged + out, or nothing to do) and ``False`` only when we could neither rebind nor + even clear — the previous actor's credentials may still be live, so the + caller fails the follow-up closed. + """ + if actor_user is None: + return True + + run_id = str(task_run.id) + scope = sandbox_identity_scope(run_id, state) + if get_sandbox_github_identity_user(scope) == actor_user.id: + return True # sandbox already reflects this actor — cheapest check first + + task = task_run.task + if get_pr_authorship_mode(task, state) != PrAuthorshipMode.USER: + return True + + sandbox = _resolve_live_sandbox(state) + if sandbox is None: + # We are past the same-actor fast path, so this is an unconfirmed transition. The + # follow-up can still reach a live agent through the saved sandbox URL, so proceeding + # would run it under the prior actor's retained credentials. A missing handle (dead + # sandbox, or a transient control-plane lookup failure) is not proof the sandbox is + # safe, so fail closed rather than deliver without a confirmed rebind or clear. + logger.info("refresh_github_no_sandbox_handle_fail_closed", run_id=run_id, user_id=actor_user.id) + return False + + repository = task.repository + token: str | None = None + try: + token = get_sandbox_github_token( + task.github_integration_id, + run_id=run_id, + state=state, + task=task, + actor_user=actor_user, + repository=repository, + ) + except ( + ReauthorizationRequired, + CredentialUnavailableError, + Integration.DoesNotExist, + UserIntegration.DoesNotExist, + ) as e: + # The new actor has no usable GitHub credential for this repo: needs reauthorization, + # no repo access, or the integration was disconnected mid-run. Log the sandbox out + # rather than run under the prior actor's creds, matching the scheduled refresh's + # handling. A transient error (network, timeout) is deliberately not caught here so it + # propagates and the activity retries. + logger.info( + "refresh_github_actor_credential_unavailable", + run_id=run_id, + user_id=actor_user.id, + repository=repository, + error_type=type(e).__name__, + reason=str(e), + ) + token = None + + # Hold the per-sandbox lock across the write and the marker update so a concurrent owner-scoped + # refresh or propagation cannot interleave and land the owner's token after this actor's — the + # owner writers acquire the same lock and re-check the marker this block advances. + with sandbox_credential_lock(sandbox.id) as acquired: + if not acquired: + logger.warning("refresh_github_lock_unavailable_fail_closed", run_id=run_id, user_id=actor_user.id) + return False + + if token: + applied = False + try: + applied = apply_github_credentials_to_sandbox(sandbox, repository, token) + except Exception: + logger.warning("refresh_github_apply_failed", run_id=run_id, exc_info=True) + if applied: + # Record the new actor only on a fully-confirmed rebind. A partial write leaves one + # credential location on the prior actor's token, so fall through to logout instead. + mark_sandbox_github_identity(scope, actor_user.id) + logger.info("refresh_github_rebound", run_id=run_id, user_id=actor_user.id) + return True + logger.warning("refresh_github_apply_incomplete", run_id=run_id, user_id=actor_user.id) + + # No usable rebind (no token, or the rebind write could not be confirmed): log the sandbox + # out. Fail closed only if even the clear can't be confirmed — the previous actor's + # credentials might still be live. The sandbox exec can raise (it stopped between the + # is_running() check and here, or timed out), so guard it like the rebind above and fail + # closed on the exception rather than letting it escape uncontrolled. + try: + cleared = clear_github_credentials_from_sandbox(sandbox, repository) + except Exception: + logger.warning("refresh_github_logout_failed", run_id=run_id, user_id=actor_user.id, exc_info=True) + return False + if cleared: + mark_sandbox_github_identity(scope, actor_user.id) + logger.info("refresh_github_logged_out", run_id=run_id, user_id=actor_user.id) + return True + logger.warning("refresh_github_logout_failed", run_id=run_id, user_id=actor_user.id) + return False + + def _get_stop_reason(result_data: dict[str, Any] | None) -> str: if not isinstance(result_data, dict): return STOP_REASON_END_TURN diff --git a/products/tasks/backend/temporal/process_task/sandbox_credentials.py b/products/tasks/backend/temporal/process_task/sandbox_credentials.py index ce7ac6a7fef6..5ca3f1d0f2a4 100644 --- a/products/tasks/backend/temporal/process_task/sandbox_credentials.py +++ b/products/tasks/backend/temporal/process_task/sandbox_credentials.py @@ -2,8 +2,10 @@ import shlex import logging +import contextlib +from collections.abc import Iterator from dataclasses import dataclass -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING, NamedTuple, Protocol from django.db import transaction @@ -22,11 +24,13 @@ get_github_token, get_pr_authorship_mode, get_readonly_github_token, + get_sandbox_github_identity_user, get_sandbox_github_token, get_task_run_credential_user, is_caller_token_run, is_slack_interaction_state, resolve_user_github_integration_for_task, + sandbox_identity_scope, ) if TYPE_CHECKING: @@ -111,12 +115,27 @@ def replace_sandbox_credentials( return github_updated and oauth_updated -def apply_github_credentials_to_sandbox(sandbox: "SandboxBase", repository: str | None, github_token: str) -> None: - """Re-inject a GitHub token into both places a running sandbox reads it from.""" - if repository: - set_git_remote_token(sandbox, repository, github_token) +def apply_github_credentials_to_sandbox(sandbox: "SandboxBase", repository: str | None, github_token: str) -> bool: + """Re-inject a GitHub token into both places a running sandbox reads it from. + + Returns ``True`` only when every applicable write succeeded. A caller enforcing per-actor + identity must treat a partial write as an unconfirmed rebind: leaving one location on the + previous actor's token would let a follow-up actor act as them. + """ + remote_applied = set_git_remote_token(sandbox, repository, github_token) if repository else True github_payload = b"".join(f"{key}={github_token}\x00".encode() for key in GITHUB_ENV_KEYS) - _write_sandbox_credential_file(sandbox, GITHUB_ENV_FILE, github_payload) + env_applied = _write_sandbox_credential_file(sandbox, GITHUB_ENV_FILE, github_payload) + return remote_applied and env_applied + + +def clear_github_credentials_from_sandbox(sandbox: "SandboxBase", repository: str | None) -> bool: + """Log the sandbox out of GitHub: strip the token from the git remote and blank the GitHub + credential file, so a follow-up actor who lacks access can't reuse the previous actor's token. + Returns ``True`` only when both were cleared. + """ + remote_cleared = set_git_remote_token(sandbox, repository, None) if repository else True + env_cleared = _write_sandbox_credential_file(sandbox, GITHUB_ENV_FILE, b"") + return remote_cleared and env_cleared def _loop_owner_credentials_revoked(task: Task, state: dict | None) -> bool: @@ -137,6 +156,87 @@ def _loop_owner_credentials_revoked(task: Task, state: dict | None) -> bool: return not eligible +def _actor_rebound_away_from_owner(run_id: str, state: dict | None, owner_id: int | None) -> int | None: + """The actor a per-message transition rebound (or logged out) this sandbox to, when that differs + from the run owner (else ``None``). Owner-scoped refresh paths (scheduled refresh, sibling + propagation) carry the owner's token, so re-applying it would resurrect the owner's identity + over the current actor's session — callers skip when this returns a value. Distinct from + `_loop_owner_credentials_revoked`, which gates on owner *eligibility* rather than session rebind.""" + bound_actor = get_sandbox_github_identity_user(sandbox_identity_scope(run_id, state)) + return bound_actor if bound_actor is not None and bound_actor != owner_id else None + + +# TTL derivation: the lock must outlive the whole critical section, or the lease can expire mid-write +# and let a concurrent refresh acquire and interleave (a regression we hit before). Each managed +# write is a chain of in-sandbox execs, each bounded by a 30s timeout — set_git_remote_token (1 exec) +# and _write_sandbox_credential_file (a file write + a chmod exec). The per-message gate can run two +# such chains back-to-back in one lock hold (apply the new token, then, if that fails, log out), so +# the worst case is ~4 × 30s ≈ 2 min. A 5 min TTL clears that with margin; recompute if those exec +# timeouts change. The wait stays under the refresh activity's 2 min timeout, so a contender that +# can't acquire skips (fail-safe) rather than blocking the activity. +_CREDENTIAL_LOCK_TTL_SECONDS = 5 * 60 +_CREDENTIAL_LOCK_WAIT_SECONDS = 15 + + +def _sandbox_credential_lock_key(sandbox_id: str) -> str: + return f"tasks:sandbox_github_creds:{sandbox_id}" + + +@contextlib.contextmanager +def _redis_lock(key: str, *, ttl: int, wait: int) -> Iterator[bool]: + """Acquire a redis lock, yielding whether it was obtained; release only if held. + + Swallows a release-time ``LockError`` so an already-expired lock never crashes the caller.""" + lock = get_client().lock(key, timeout=ttl, blocking_timeout=wait) + acquired = lock.acquire() + try: + yield acquired + finally: + if acquired: + try: + lock.release() + except redis.exceptions.LockError: + logger.warning("redis_lock_release_failed", extra={"lock_key": key}) + + +@contextlib.contextmanager +def sandbox_credential_lock(sandbox_id: str) -> Iterator[bool]: + """Serialize every writer of one sandbox's GitHub credentials — scheduled refresh, sibling + propagation, and the per-message actor gate — so their check → write → marker-update runs + atomically. Without it, a slow owner-token write can land *after* a follow-up rebound the sandbox + to a different actor, resurrecting the owner's identity for the current actor. Yields whether the + lock was acquired; a caller that does not get it must skip the write rather than race.""" + with _redis_lock( + _sandbox_credential_lock_key(sandbox_id), ttl=_CREDENTIAL_LOCK_TTL_SECONDS, wait=_CREDENTIAL_LOCK_WAIT_SECONDS + ) as acquired: + yield acquired + + +def _apply_owner_token_locked( + sandbox: "SandboxBase", repository: str | None, token: str, run_id: str, state: dict | None, owner_id: int | None +) -> bool: + """Apply an owner-scoped token only while the sandbox is still bound to the owner. + + Owner-scoped writers resolve the token from startup context, which can take seconds against + GitHub's API. Serializing the re-check and the write under the per-sandbox lock closes the window + where a per-message transition rebinds the sandbox between the caller's earlier check and this + write. Returns ``True`` only when the token was actually applied.""" + with sandbox_credential_lock(sandbox.id) as acquired: + if not acquired: + logger.warning( + "owner_token_apply_skipped_lock_unavailable", extra={"run_id": run_id, "sandbox_id": sandbox.id} + ) + return False + rebound_actor = _actor_rebound_away_from_owner(run_id, state, owner_id) + if rebound_actor is not None: + logger.info( + "owner_token_apply_skipped_actor_transition", + extra={"run_id": run_id, "bound_actor": rebound_actor, "owner": owner_id}, + ) + return False + return apply_github_credentials_to_sandbox(sandbox, repository, token) + + USER_TOKEN_REFRESH_INTERVAL_SECONDS: float = _GITHUB_REFRESH_INTERVAL_BY_PREFIX["ghu_"] # TTL covers a slow mint + propagation; wait stays under the refresh activity's 2 min timeout. _ROTATION_LOCK_TTL_SECONDS = 120 @@ -147,8 +247,19 @@ def _rotation_lock_key(user_integration_id: int) -> str: return f"tasks:gh_user_token_rotate:{user_integration_id}" -def _live_sandboxes_for_user_integration(user_integration_id: int) -> list[tuple[str, str, str | None]]: - rows: list[tuple[str, str, str | None]] = [] +class LiveSandbox(NamedTuple): + """A live sandbox eligible for owner-token propagation, with the fields the per-sandbox + actor-rebind re-check needs (``state`` and ``owner_id``) carried alongside.""" + + run_id: str + sandbox_id: str + repository: str | None + state: dict | None + owner_id: int | None + + +def _live_sandboxes_for_user_integration(user_integration_id: int) -> list[LiveSandbox]: + rows: list[LiveSandbox] = [] runs = TaskRun.objects.filter( status=TaskRun.Status.IN_PROGRESS, task__github_user_integration_id=user_integration_id, @@ -166,7 +277,11 @@ def _live_sandboxes_for_user_integration(user_integration_id: int) -> list[tuple continue if _loop_owner_credentials_revoked(run.task, run.state): continue - rows.append((str(run.id), sandbox_id, run.task.repository)) + # This loop carries the owner's token; skip a sandbox a per-message transition rebound to + # a different actor, or re-applying it would resurrect the owner's identity for that actor. + if _actor_rebound_away_from_owner(str(run.id), run.state, run.task.created_by_id) is not None: + continue + rows.append(LiveSandbox(str(run.id), sandbox_id, run.task.repository, run.state, run.task.created_by_id)) return rows @@ -174,16 +289,19 @@ def _propagate_user_token(user_integration_id: int, token: str) -> int: from products.tasks.backend.logic.services.sandbox import Sandbox # noqa: PLC0415 applied = 0 - for run_id, sandbox_id, repository in _live_sandboxes_for_user_integration(user_integration_id): + for live in _live_sandboxes_for_user_integration(user_integration_id): try: - sandbox = Sandbox.get_by_id(sandbox_id) - if sandbox.is_running(): - apply_github_credentials_to_sandbox(sandbox, repository, token) + sandbox = Sandbox.get_by_id(live.sandbox_id) + # Re-check the actor binding under the per-sandbox lock: the filter above is not atomic + # with this write, so a transition could have rebound the sandbox in between. + if sandbox.is_running() and _apply_owner_token_locked( + sandbox, live.repository, token, live.run_id, live.state, live.owner_id + ): applied += 1 except Exception: logger.warning( "Failed to propagate refreshed GitHub user token to sibling sandbox", - extra={"integration_id": user_integration_id, "run_id": run_id, "sandbox_id": sandbox_id}, + extra={"integration_id": user_integration_id, "run_id": live.run_id, "sandbox_id": live.sandbox_id}, exc_info=True, ) return applied @@ -200,17 +318,14 @@ def resolve_coordinated_user_token(integration: UserGitHubIntegration) -> str | return integration.get_usable_user_access_token() integration_id = integration.integration.id - lock = get_client().lock( - _rotation_lock_key(integration_id), - timeout=_ROTATION_LOCK_TTL_SECONDS, - blocking_timeout=_ROTATION_LOCK_WAIT_SECONDS, - ) - if not lock.acquire(): - # Waited out the budget — read the current token without minting; the holder's propagation self-heals. - integration.integration.refresh_from_db() - return UserGitHubIntegration(integration.integration).user_access_token + with _redis_lock( + _rotation_lock_key(integration_id), ttl=_ROTATION_LOCK_TTL_SECONDS, wait=_ROTATION_LOCK_WAIT_SECONDS + ) as acquired: + if not acquired: + # Waited out the budget — read the current token without minting; the holder's propagation self-heals. + integration.integration.refresh_from_db() + return UserGitHubIntegration(integration.integration).user_access_token - try: integration.integration.refresh_from_db() current = UserGitHubIntegration(integration.integration) was_expired = current.user_access_token_expired() @@ -223,14 +338,6 @@ def resolve_coordinated_user_token(integration: UserGitHubIntegration) -> str | extra={"integration_id": integration_id, "sibling_sandboxes_updated": propagated}, ) return token - finally: - try: - lock.release() - except redis.exceptions.LockError: - logger.warning( - "GitHub user-token rotation lock already expired/released", - extra={"integration_id": integration_id}, - ) @dataclass @@ -281,6 +388,19 @@ def refresh(self, sandbox: "SandboxBase", ctx: "TaskProcessingContext", task: Ta self.kind, refreshed=False, next_refresh_seconds=DEFAULT_REFRESH_INTERVAL_SECONDS ) + # This scheduled refresh resolves the actor from startup context (ctx.state), so it carries + # the owner's token; skip when a per-message transition rebound the sandbox and leave that + # binding intact — the per-message gate keeps the current actor's token fresh. + rebound_actor = _actor_rebound_away_from_owner(ctx.run_id, ctx.state, task.created_by_id) + if rebound_actor is not None: + logger.info( + "github_refresh_skipped_actor_transition", + extra={"run_id": ctx.run_id, "bound_actor": rebound_actor, "owner": task.created_by_id}, + ) + return CredentialRefreshOutcome( + self.kind, refreshed=False, next_refresh_seconds=DEFAULT_REFRESH_INTERVAL_SECONDS + ) + actor_user = get_task_run_credential_user(task, ctx.state) if is_slack_interaction_state(ctx.state) and actor_user is None: raise ReauthorizationRequired("Slack run requires an acting user before refreshing GitHub credentials.") @@ -340,9 +460,9 @@ def refresh(self, sandbox: "SandboxBase", ctx: "TaskProcessingContext", task: Ta self.kind, refreshed=False, next_refresh_seconds=DEFAULT_REFRESH_INTERVAL_SECONDS ) - apply_github_credentials_to_sandbox(sandbox, ctx.repository, token) + applied = _apply_owner_token_locked(sandbox, ctx.repository, token, ctx.run_id, ctx.state, task.created_by_id) return CredentialRefreshOutcome( - self.kind, refreshed=True, next_refresh_seconds=github_refresh_interval_seconds(token) + self.kind, refreshed=applied, next_refresh_seconds=github_refresh_interval_seconds(token) ) def _refresh_shared_user_integration( @@ -358,16 +478,21 @@ def _refresh_shared_user_integration( return CredentialRefreshOutcome( self.kind, refreshed=False, next_refresh_seconds=DEFAULT_REFRESH_INTERVAL_SECONDS ) - apply_github_credentials_to_sandbox(sandbox, ctx.repository, fallback) + applied = _apply_owner_token_locked( + sandbox, ctx.repository, fallback, ctx.run_id, ctx.state, task.created_by_id + ) return CredentialRefreshOutcome( - self.kind, refreshed=True, next_refresh_seconds=github_refresh_interval_seconds(fallback) + self.kind, refreshed=applied, next_refresh_seconds=github_refresh_interval_seconds(fallback) ) if token and _loop_owner_credentials_revoked(task, ctx.state): token = None - if token: - apply_github_credentials_to_sandbox(sandbox, ctx.repository, token) + applied = ( + _apply_owner_token_locked(sandbox, ctx.repository, token, ctx.run_id, ctx.state, task.created_by_id) + if token + else False + ) return CredentialRefreshOutcome( - self.kind, refreshed=bool(token), next_refresh_seconds=USER_TOKEN_REFRESH_INTERVAL_SECONDS + self.kind, refreshed=applied, next_refresh_seconds=USER_TOKEN_REFRESH_INTERVAL_SECONDS ) def _installation_token_fallback(self, ctx: "TaskProcessingContext", task: Task, cause: Exception) -> str | None: diff --git a/products/tasks/backend/temporal/process_task/tests/test_sandbox_credentials.py b/products/tasks/backend/temporal/process_task/tests/test_sandbox_credentials.py index c338f6cbaa9b..9b370dd25cf8 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_sandbox_credentials.py +++ b/products/tasks/backend/temporal/process_task/tests/test_sandbox_credentials.py @@ -204,6 +204,26 @@ def test_caller_token_run_with_deleted_integration_is_not_orphaned(self): assert outcome.refreshed is True + def test_scheduled_refresh_skips_when_sandbox_bound_to_different_actor(self): + # A per-message actor transition rebound this sandbox to another actor. The scheduled + # refresh resolves the actor from the startup context, so it carries the owner's token; + # applying it would resurrect the owner's identity over the current actor's session. + from products.tasks.backend.temporal.process_task.utils import mark_sandbox_github_identity + + sandbox = MagicMock() + task = MagicMock() + task.github_integration_id = 123 + task.created_by_id = 2 # run owner + mark_sandbox_github_identity("run-transition", 99) # transitioned to a different actor + + with patch(f"{MODULE}.get_sandbox_github_token") as resolve: + outcome = GitHubSandboxCredential().refresh(sandbox, _context(run_id="run-transition"), task) + + assert outcome.refreshed is False + resolve.assert_not_called() # never resolved or applied the owner's token + sandbox.execute.assert_not_called() + sandbox.write_file.assert_not_called() + class TestBuildSandboxCredentials: def test_includes_github_when_credentials_present(self): @@ -234,7 +254,7 @@ def test_refresh_applies_coordinated_token(self): self._as_user_integration_run(stack) stack.enter_context(patch(f"{MODULE}.resolve_user_github_integration_for_task", return_value=MagicMock())) resolve = stack.enter_context(patch(f"{MODULE}.resolve_coordinated_user_token", return_value="ghu_fresh")) - apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox")) + apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox", return_value=True)) sandbox = MagicMock() outcome = GitHubSandboxCredential().refresh(sandbox, _context(), MagicMock()) @@ -270,7 +290,7 @@ def test_reauthorization_falls_back_to_installation_token(self): patch(f"{MODULE}.resolve_coordinated_user_token", side_effect=ReauthorizationRequired("expired")) ) installation_token = stack.enter_context(patch(f"{MODULE}.get_github_token", return_value="ghs_team")) - apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox")) + apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox", return_value=True)) task = MagicMock() task.github_integration_id = 456 @@ -311,7 +331,7 @@ def test_caller_token_run_skips_coordinated_path(self): stack.enter_context(patch(f"{MODULE}.is_caller_token_run", return_value=True)) resolve = stack.enter_context(patch(f"{MODULE}.resolve_user_github_integration_for_task")) stack.enter_context(patch(f"{MODULE}.get_sandbox_github_token", return_value="ghu_caller")) - apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox")) + apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox", return_value=True)) sandbox = MagicMock() sandbox.id = "sb-own" @@ -322,6 +342,92 @@ def test_caller_token_run_skips_coordinated_path(self): apply.assert_called_once_with(sandbox, "explore-science/paper-wizard-frontend", "ghu_caller") +class TestApplyOwnerTokenLocked: + def _lock(self, stack, *, acquired): + lock = MagicMock() + lock.acquire.return_value = acquired + get_client = stack.enter_context(patch(f"{MODULE}.get_client")) + get_client.return_value.lock.return_value = lock + return lock + + def test_applies_while_sandbox_still_bound_to_owner(self): + import contextlib + + from products.tasks.backend.temporal.process_task.sandbox_credentials import _apply_owner_token_locked + + with contextlib.ExitStack() as stack: + self._lock(stack, acquired=True) + stack.enter_context(patch(f"{MODULE}.get_sandbox_github_identity_user", return_value=None)) + apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox", return_value=True)) + sandbox = MagicMock() + sandbox.id = "sb-1" + + assert _apply_owner_token_locked(sandbox, "org/repo", "ghu_x", "run-1", {}, 7) is True + apply.assert_called_once_with(sandbox, "org/repo", "ghu_x") + + def test_skips_when_a_transition_rebound_the_sandbox_to_another_actor(self): + import contextlib + + from products.tasks.backend.temporal.process_task.sandbox_credentials import _apply_owner_token_locked + + with contextlib.ExitStack() as stack: + self._lock(stack, acquired=True) + # Marker moved to actor 99 under the lock — the owner (7) token must not overwrite it. + stack.enter_context(patch(f"{MODULE}.get_sandbox_github_identity_user", return_value=99)) + apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox")) + sandbox = MagicMock() + sandbox.id = "sb-1" + + assert _apply_owner_token_locked(sandbox, "org/repo", "ghu_x", "run-1", {}, 7) is False + apply.assert_not_called() + + def test_fails_closed_without_applying_when_the_lock_is_contended(self): + import contextlib + + from products.tasks.backend.temporal.process_task.sandbox_credentials import _apply_owner_token_locked + + with contextlib.ExitStack() as stack: + lock = self._lock(stack, acquired=False) + apply = stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox")) + sandbox = MagicMock() + sandbox.id = "sb-1" + + assert _apply_owner_token_locked(sandbox, "org/repo", "ghu_x", "run-1", {}, 7) is False + apply.assert_not_called() + lock.release.assert_not_called() + + def test_lock_is_leased_long_enough_to_outlive_a_writer_past_the_old_30s_lease(self): + # A credential write does a git-remote rewrite + env-file write + chmod, each an in-sandbox + # exec bounded by a 30s timeout, so it can run well past the old 30s lease. The redis lock + # must be leased for that whole worst case, or the lease expires mid-write and a concurrent + # refresh could acquire and interleave. + import contextlib + + from products.tasks.backend.temporal.process_task.sandbox_credentials import ( + _CREDENTIAL_LOCK_TTL_SECONDS, + _apply_owner_token_locked, + ) + + assert _CREDENTIAL_LOCK_TTL_SECONDS == 5 * 60 + assert _CREDENTIAL_LOCK_TTL_SECONDS > 2 * 30 # clears a 30s git-remote + 30s chmod worst case + + with contextlib.ExitStack() as stack: + get_client = stack.enter_context(patch(f"{MODULE}.get_client")) + lock = MagicMock() + lock.acquire.return_value = True + get_client.return_value.lock.return_value = lock + stack.enter_context(patch(f"{MODULE}.get_sandbox_github_identity_user", return_value=None)) + stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox", return_value=True)) + sandbox = MagicMock() + sandbox.id = "sb-1" + + _apply_owner_token_locked(sandbox, "org/repo", "ghu_x", "run-1", {}, 7) + + # The lock is leased for the full worst-case write, not the old 30s. + get_client.return_value.lock.assert_called_once() + assert get_client.return_value.lock.call_args.kwargs["timeout"] == _CREDENTIAL_LOCK_TTL_SECONDS + + class TestLoopOwnerRefreshGate: def _as_user_integration_run(self, stack): from products.tasks.backend.temporal.process_task.utils import PrAuthorshipMode @@ -541,7 +647,49 @@ def _task(repo): result = _live_sandboxes_for_user_integration(integration.id) - assert set(result) == { + assert {(r.run_id, r.sandbox_id, r.repository) for r in result} == { (str(live_run.id), "sb-live", "org/live"), (str(eligible_loop_run.id), "sb-loop", "org/loop"), } + + @pytest.mark.parametrize("marker,included", [("none", True), ("owner", True), ("other", False)]) + def test_actor_transition_gates_owner_token_propagation(self, marker, included): + from posthog.models import Organization, Team + from posthog.models.user import User + from posthog.models.user_integration import UserIntegration + + from products.tasks.backend.models import Task, TaskRun + from products.tasks.backend.temporal.process_task.sandbox_credentials import ( + _live_sandboxes_for_user_integration, + ) + from products.tasks.backend.temporal.process_task.utils import mark_sandbox_github_identity + + org = Organization.objects.create(name="o") + team = Team.objects.create(organization=org, name="t") + owner = User.objects.create(email="owner@test.com") + other = User.objects.create(email="other@test.com") + integration = UserIntegration.objects.create( + user=owner, kind=UserIntegration.IntegrationKind.GITHUB, integration_id="i1", config={}, sensitive_config={} + ) + task = Task.objects.create( + team=team, created_by=owner, repository="org/repo", github_user_integration=integration + ) + run = TaskRun.objects.create( + task=task, + team=team, + status=TaskRun.Status.IN_PROGRESS, + state={"sandbox_id": "sb-x", "pr_authorship_mode": "user"}, + ) + # An unset marker (no transition yet) and one bound to the owner both propagate; a marker + # bound to a different per-message actor means the sandbox was logged out / rebound, so the + # owner's rotating token must not overwrite it. + if marker == "owner": + mark_sandbox_github_identity("sb-x", owner.id) + elif marker == "other": + mark_sandbox_github_identity("sb-x", other.id) + + result = _live_sandboxes_for_user_integration(integration.id) + + assert ( + [(r.run_id, r.sandbox_id, r.repository) for r in result] == [(str(run.id), "sb-x", "org/repo")] + ) is included diff --git a/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py b/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py index a5d7a67ef691..e5f0e86fee5f 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py +++ b/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py @@ -5,19 +5,25 @@ from temporalio.exceptions import ApplicationError +from posthog.models.user_integration import ReauthorizationRequired + from products.tasks.backend.logic.services.agent_command import CommandResult from products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox import ( REFRESH_RETRY_DELAY_SECONDS, SEND_FOLLOWUP_MAX_ATTEMPTS, STEER_DECLINED_OUTCOME, SendFollowupToSandboxInput, + _refresh_sandbox_github, _refresh_sandbox_mcp, send_followup_to_sandbox, ) from products.tasks.backend.temporal.process_task.utils import ( McpServerConfig, - _sandbox_mcp_session_cache_key, + PrAuthorshipMode, + _sandbox_identity_cache_key, + get_sandbox_github_identity_user, get_sandbox_mcp_session_user, + mark_sandbox_github_identity, mark_sandbox_mcp_session, ) @@ -311,7 +317,7 @@ def test_replacement_sandbox_starts_unmarked( mock_send_refresh.assert_called_once() assert get_sandbox_mcp_session_user("sb-2") == 42 - assert cache.get(_sandbox_mcp_session_cache_key("run-1")) == 42 # untouched + assert cache.get(_sandbox_identity_cache_key("mcp-session", "run-1")) == 42 # untouched def test_transition_with_no_configs_fails_closed( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh @@ -349,6 +355,160 @@ def test_unknown_binding_with_no_configs_runs( assert get_sandbox_mcp_session_user("run-1") == 42 # binding recorded +_GH_MODULE = "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox" + + +@patch(f"{_GH_MODULE}.clear_github_credentials_from_sandbox") +@patch(f"{_GH_MODULE}.apply_github_credentials_to_sandbox") +@patch(f"{_GH_MODULE}.get_sandbox_github_token") +@patch(f"{_GH_MODULE}._resolve_live_sandbox") +@patch(f"{_GH_MODULE}.get_pr_authorship_mode") +class TestSandboxGithubIdentityGate: + """On an actor transition the sandbox's GitHub credentials rebind to the new + actor when they have access, otherwise the sandbox is logged out so the + previous actor's identity can't be used.""" + + def test_same_actor_skips(self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear): + mock_authorship.return_value = PrAuthorshipMode.USER + mark_sandbox_github_identity("run-1", 42) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is True + mock_resolve.assert_not_called() + mock_get_token.assert_not_called() + mock_apply.assert_not_called() + mock_clear.assert_not_called() + + def test_bot_authorship_skips(self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear): + # BOT runs share a single installation token, so every actor is already + # the same GitHub identity — nothing to rebind. + mock_authorship.return_value = PrAuthorshipMode.BOT + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is True + mock_get_token.assert_not_called() + mock_apply.assert_not_called() + mock_clear.assert_not_called() + + def test_transition_with_access_rebinds( + self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear + ): + mock_authorship.return_value = PrAuthorshipMode.USER + mock_resolve.return_value = MagicMock() + mock_get_token.return_value = "ghu_newtoken" + mock_apply.return_value = True + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is True + mock_apply.assert_called_once() + assert mock_apply.call_args.args[2] == "ghu_newtoken" + mock_clear.assert_not_called() + assert get_sandbox_github_identity_user("run-1") == 42 + + def test_transition_without_access_logs_out( + self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear + ): + mock_authorship.return_value = PrAuthorshipMode.USER + mock_resolve.return_value = MagicMock() + mock_get_token.side_effect = ReauthorizationRequired("no repo access") + mock_clear.return_value = True + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is True + mock_apply.assert_not_called() + mock_clear.assert_called_once() + assert get_sandbox_github_identity_user("run-1") == 42 + + def test_apply_failure_falls_back_to_logout( + self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear + ): + mock_authorship.return_value = PrAuthorshipMode.USER + mock_resolve.return_value = MagicMock() + mock_get_token.return_value = "ghu_newtoken" + mock_apply.side_effect = RuntimeError("write failed") + mock_clear.return_value = True + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is True + mock_apply.assert_called_once() + mock_clear.assert_called_once() # fell through to logout so no stale creds remain + + def test_apply_incomplete_falls_back_to_logout( + self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear + ): + # A partial credential write (one location refused, no exception) is not a confirmed + # rebind: the prior actor's token may still be live in the other location, so log out + # rather than record the new actor. + mock_authorship.return_value = PrAuthorshipMode.USER + mock_resolve.return_value = MagicMock() + mock_get_token.return_value = "ghu_newtoken" + mock_apply.return_value = False + mock_clear.return_value = True + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is True + mock_apply.assert_called_once() + mock_clear.assert_called_once() + assert get_sandbox_github_identity_user("run-1") == 42 # logout confirmed, bound to new actor + + def test_no_sandbox_handle_fails_closed( + self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear + ): + # The handle can't be resolved (dead sandbox or transient lookup failure), but a follow-up + # can still reach a live agent via the saved URL. Fail closed rather than run under the + # prior actor's retained creds. + mock_authorship.return_value = PrAuthorshipMode.USER + mock_resolve.return_value = None + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is False + mock_get_token.assert_not_called() + mock_apply.assert_not_called() + mock_clear.assert_not_called() + assert get_sandbox_github_identity_user("run-1") == 99 # binding unchanged + + def test_logout_failure_fails_closed(self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear): + # New actor has no access and the sandbox can't even be cleared — the + # previous actor's creds may still be live, so fail closed. + mock_authorship.return_value = PrAuthorshipMode.USER + mock_resolve.return_value = MagicMock() + mock_get_token.side_effect = ReauthorizationRequired("no repo access") + mock_clear.return_value = False + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is False + assert get_sandbox_github_identity_user("run-1") == 99 # binding unchanged + + def test_logout_exception_fails_closed(self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear): + # The clear itself raising (sandbox stopped/timed out between is_running and here) must + # fail closed, not escape uncontrolled. + mock_authorship.return_value = PrAuthorshipMode.USER + mock_resolve.return_value = MagicMock() + mock_get_token.side_effect = ReauthorizationRequired("no repo access") + mock_clear.side_effect = RuntimeError("sandbox stopped") + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is False + assert get_sandbox_github_identity_user("run-1") == 99 # binding unchanged + + def test_credential_unavailable_logs_out( + self, mock_authorship, mock_resolve, mock_get_token, mock_apply, mock_clear + ): + # A disconnected/deleted integration mid-run yields no usable credential (not just + # ReauthorizationRequired): log out rather than let the exception escape. + from products.tasks.backend.exceptions import CredentialUnavailableError + + mock_authorship.return_value = PrAuthorshipMode.USER + mock_resolve.return_value = MagicMock() + mock_get_token.side_effect = CredentialUnavailableError("integration disconnected", {}) + mock_clear.return_value = True + mark_sandbox_github_identity("run-1", 99) + + assert _refresh_sandbox_github(_make_task_run_mock(), MagicMock(id=42), None) is True + mock_apply.assert_not_called() + mock_clear.assert_called_once() + assert get_sandbox_github_identity_user("run-1") == 42 + + class TestSendFollowupActivityRefreshOrdering: """Refresh call must precede user_message, and the activity must succeed when refresh fails (non-fatal) as long as user_message succeeds.""" diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index 70849e521b32..4a17297762ac 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -402,8 +402,16 @@ def sandbox_identity_scope(run_id: str, state: dict[str, Any] | None) -> str: return (state or {}).get("sandbox_id") or run_id -def _sandbox_mcp_session_cache_key(scope: str) -> str: - return f"tasks:sandbox-mcp-session:{scope}" +def _sandbox_identity_cache_key(kind: str, scope: str) -> str: + return f"tasks:sandbox-{kind}:{scope}" + + +def _mark_sandbox_identity(kind: str, scope: str, user_id: int) -> None: + get_tasks_cache().set(_sandbox_identity_cache_key(kind, scope), user_id, timeout=MCP_TOKEN_REFRESH_INTERVAL_SECONDS) + + +def _get_sandbox_identity_user(kind: str, scope: str) -> int | None: + return get_tasks_cache().get(_sandbox_identity_cache_key(kind, scope)) def mark_sandbox_mcp_session(scope: str, user_id: int) -> None: @@ -412,13 +420,31 @@ def mark_sandbox_mcp_session(scope: str, user_id: int) -> None: Self-expires after MCP_TOKEN_REFRESH_INTERVAL_SECONDS, so an absent entry always reads as "must refresh". """ - get_tasks_cache().set(_sandbox_mcp_session_cache_key(scope), user_id, timeout=MCP_TOKEN_REFRESH_INTERVAL_SECONDS) + _mark_sandbox_identity("mcp-session", scope, user_id) def get_sandbox_mcp_session_user(scope: str) -> int | None: """User id the sandbox's MCP session was last bound to within the freshness window, or None when unknown.""" - return get_tasks_cache().get(_sandbox_mcp_session_cache_key(scope)) + return _get_sandbox_identity_user("mcp-session", scope) + + +def mark_sandbox_github_identity(scope: str, user_id: int) -> None: + """Record which actor the sandbox's in-place GitHub credentials reflect. + + The value is the actor whose token was applied, or who was logged out (no + usable access) — either way the sandbox no longer carries a *different* + actor's identity. Self-expires after MCP_TOKEN_REFRESH_INTERVAL_SECONDS; an + absent entry reads as "must re-establish", which is always safe because + re-establishing re-applies or clears rather than trusting stale creds. + """ + _mark_sandbox_identity("github-identity", scope, user_id) + + +def get_sandbox_github_identity_user(scope: str) -> int | None: + """Actor id the sandbox's GitHub credentials were last bound to (or logged + out for) within the freshness window, or None when unknown.""" + return _get_sandbox_identity_user("github-identity", scope) @dataclass(frozen=True) diff --git a/products/tasks/backend/tests/test_agentsh.py b/products/tasks/backend/tests/test_agentsh.py index d7e4d8524f14..a5bf7e0bbb1e 100644 --- a/products/tasks/backend/tests/test_agentsh.py +++ b/products/tasks/backend/tests/test_agentsh.py @@ -463,6 +463,8 @@ def execute(command: str, timeout_seconds: int | None = None) -> ExecutionResult if "--taskId" in command: launched.append(command) return ExecutionResult(stdout="", stderr="", exit_code=0) + if "chmod" in command: # gh shim install + return ExecutionResult(stdout="", stderr="", exit_code=0) self.assertIn("grep", command) return ExecutionResult(stdout="", stderr="", exit_code=0 if supported else 1)