From cb3ffdd7497cdf3b34f795bedcd7e7e03c937841 Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Fri, 10 Jul 2026 15:02:02 +0200 Subject: [PATCH] feat(tasks): rebind sandbox GitHub identity and git author to the actor Extends the per-message identity rebind: when the acting user has a personal GitHub install covering the task's repository, swap the git remote token, GITHUB_TOKEN/GH_TOKEN, and GIT_AUTHOR_*/GIT_COMMITTER_* in the live sandbox via the agentsh env file (re-sourced per command, so no restart), making this turn's commits and PRs theirs. Actors without a covering install keep the previous identity; bot-authored and caller-token runs are never swapped. The github/github_user identity marks make the credential refresh loop resolve the swapped integration instead of the task creator's, keep token-rotation propagation from reverting a swap, and let same-actor messages skip integration resolution entirely. Sandbox liveness is checked before the rotation lock so a dead sandbox never triggers a token mint-and-revoke. --- products/tasks/backend/facade/api.py | 19 ++- .../process_task/sandbox_credentials.py | 146 +++++++++++++++- .../tests/test_sandbox_credentials.py | 159 ++++++++++++++++++ .../backend/temporal/process_task/utils.py | 21 ++- 4 files changed, 330 insertions(+), 15 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 2608fe15a0a4..faa3747f9819 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -4591,17 +4591,22 @@ def rebind_sandbox_identity_for_user( auth_token: str | None = None, ) -> None: """Rebind a run's live sandbox to ``user_id``: mint a PostHog MCP OAuth - token for them and push fresh MCP configs. + token for them and push fresh MCP configs, then swap the GitHub token and + git author identity when they have a personal install covering the task's + repository (otherwise the previous GitHub identity keeps authoring). Used by Slack follow-ups so a teammate taking over the conversation acts - as themselves — insights, dashboards, and other PostHog writes attribute - to the live actor rather than the task creator. Best-effort by contract: - failures are logged, never raised, so a rebind problem can't block the - message that triggered it. + as themselves — insights, dashboards, commits, and PRs attribute to the + live actor rather than the task creator. Best-effort by contract: each + credential kind is rebound independently and failures are logged, never + raised, so a rebind problem can't block the message that triggered it. """ from products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox import ( # noqa: PLC0415 — keep sandbox deps off the api import path refresh_sandbox_mcp_for_user, ) + from products.tasks.backend.temporal.process_task.sandbox_credentials import ( # noqa: PLC0415 — keep sandbox deps off the api import path + refresh_sandbox_github_for_user, + ) try: run = TaskRun.objects.select_related("task").get(id=run_id) @@ -4620,3 +4625,7 @@ def rebind_sandbox_identity_for_user( refresh_sandbox_mcp_for_user(run, user, scopes="full", auth_token=auth_token) except Exception: logger.exception("Sandbox MCP identity rebind failed", extra={"run_id": str(run_id), "user_id": user_id}) + try: + refresh_sandbox_github_for_user(run, user) + except Exception: + logger.exception("Sandbox GitHub identity rebind failed", extra={"run_id": str(run_id), "user_id": user_id}) diff --git a/products/tasks/backend/temporal/process_task/sandbox_credentials.py b/products/tasks/backend/temporal/process_task/sandbox_credentials.py index 099c34756f0b..2fab1fb2de21 100644 --- a/products/tasks/backend/temporal/process_task/sandbox_credentials.py +++ b/products/tasks/backend/temporal/process_task/sandbox_credentials.py @@ -18,13 +18,20 @@ from products.tasks.backend.temporal.process_task.utils import ( PrAuthorshipMode, get_github_token, + get_last_sandbox_identity, get_pr_authorship_mode, get_sandbox_github_token, + get_user_github_integration, + git_identity_env_for_user, is_caller_token_run, + mark_sandbox_identity, resolve_user_github_integration_for_task, + sandbox_identity_scope, ) if TYPE_CHECKING: + from posthog.models.user import User + from products.tasks.backend.logic.services.sandbox import SandboxBase from .activities.get_task_processing_context import TaskProcessingContext @@ -115,11 +122,20 @@ def update_sandbox_env_file(sandbox: "SandboxBase", updates: dict[str, str]) -> return True -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.""" +def apply_github_credentials_to_sandbox( + sandbox: "SandboxBase", + repository: str | None, + github_token: str, + extra_env: dict[str, str] | None = None, +) -> None: + """Re-inject a GitHub token into both places a running sandbox reads it from. + + ``extra_env`` rides along in the same env-file read-modify-write so callers + that also change the git author identity don't pay a second round-trip. + """ if repository: set_git_remote_token(sandbox, repository, github_token) - update_sandbox_env_file(sandbox, dict.fromkeys(GITHUB_ENV_KEYS, github_token)) + update_sandbox_env_file(sandbox, {**dict.fromkeys(GITHUB_ENV_KEYS, github_token), **(extra_env or {})}) USER_TOKEN_REFRESH_INTERVAL_SECONDS: float = _GITHUB_REFRESH_INTERVAL_BY_PREFIX["ghu_"] @@ -153,6 +169,12 @@ def _live_sandboxes_for_user_integration(user_integration_id: int) -> list[tuple continue if is_caller_token_run(str(run.id), run.state): continue + # A Slack actor may have swapped this sandbox's live identity to a + # different integration; propagating the task owner's token over it + # would silently revert the swap. + swapped_identity = get_last_sandbox_identity(sandbox_identity_scope(str(run.id), run.state), "github") + if swapped_identity is not None and str(swapped_identity) != str(user_integration_id): + continue rows.append((str(run.id), sandbox_id, run.task.repository)) return rows @@ -255,7 +277,11 @@ def refresh(self, sandbox: "SandboxBase", ctx: "TaskProcessingContext", task: Ta if get_pr_authorship_mode(task, ctx.state) == PrAuthorshipMode.USER and not is_caller_token_run( ctx.run_id, ctx.state ): - integration = resolve_user_github_integration_for_task(task, repository=ctx.repository, allow_refresh=True) + integration = _swapped_identity_integration(sandbox_identity_scope(ctx.run_id, ctx.state), task) + if integration is None: + integration = resolve_user_github_integration_for_task( + task, repository=ctx.repository, allow_refresh=True + ) if integration is not None: return self._refresh_shared_user_integration(sandbox, ctx, task, integration) @@ -343,6 +369,118 @@ def _installation_token_fallback(self, ctx: "TaskProcessingContext", task: Task, ) +def _current_github_identity(scope: str, task: Task) -> str: + """The UserIntegration id the sandbox's GitHub credentials are bound to. + + ``scope`` comes from ``sandbox_identity_scope``. Falls back to the task's + own integration — the boot-time identity — when no swap was ever recorded + (or the cache entry was evicted).""" + return str(get_last_sandbox_identity(scope, "github") or task.github_user_integration_id) + + +def _swapped_identity_integration(scope: str, task: Task) -> UserGitHubIntegration | None: + """Return the UserGitHubIntegration a Slack actor swapped this sandbox to, if any. + + None when the sandbox was never swapped, still holds the task's own + identity, or the swapped integration has since been deleted (in which + case the caller falls back to the task's own resolution).""" + swapped_id = _current_github_identity(scope, task) + if swapped_id == str(task.github_user_integration_id): + return None + try: + integration = UserIntegration.objects.get(id=swapped_id, kind="github") + except UserIntegration.DoesNotExist: + return None + return UserGitHubIntegration(integration) + + +def refresh_sandbox_github_for_user(task_run: TaskRun, user: "User") -> bool: + """Rebind a live sandbox's GitHub credentials and git author identity to ``user``. + + Rewrites the git remote token, the agentsh GITHUB_TOKEN/GH_TOKEN entries, + and the GIT_AUTHOR_*/GIT_COMMITTER_* entries so the agent's subsequent + commits and ``gh`` calls (including ``gh pr create``) act as ``user``. The + agentsh exec wrapper re-sources the env file per command, so no sandbox + restart is needed. + + No-ops (returning False) when the run isn't user-authored, is pinned to a + caller-supplied token, the sandbox is gone, ``user`` has no personal GitHub + install covering the task's repository, or the sandbox already holds this + identity — the previous identity keeps authoring in those cases. Raises + only on unexpected errors; expected credential problems are logged. + """ + task = task_run.task + run_id = str(task_run.id) + state = task_run.state or {} + + if get_pr_authorship_mode(task, state) != PrAuthorshipMode.USER or is_caller_token_run(run_id, state): + return False + + sandbox_id = state.get("sandbox_id") + if not sandbox_id: + return False + scope = sandbox_identity_scope(run_id, state) + + # Fast path for the dominant case: same actor as the sandbox's current + # GitHub identity (the creator's boot-time one when never swapped). Skips + # the integration resolution — a DB query and, on a stale repo cache, a + # GitHub API sync — per message; the refresh loop owns token freshness. + current_github_user = get_last_sandbox_identity(scope, "github_user") or task.created_by_id + if user.id == current_github_user: + return False + + integration = get_user_github_integration(user, repository=task.repository, allow_refresh=True) + if integration is None: + logger.info( + "GitHub identity swap skipped: actor has no personal install covering the repository", + extra={"run_id": run_id, "user_id": user.id, "repository": task.repository}, + ) + return False + + current_identity = _current_github_identity(scope, task) + integration_id = str(integration.integration.id) + if integration_id == current_identity: + # Already this identity (stale user mark); record the user and move on. + mark_sandbox_identity(scope, "github_user", user.id) + return False + + from products.tasks.backend.logic.services.sandbox import ( + Sandbox, # noqa: PLC0415 — keep sandbox deps off the module import path + ) + + # Liveness before token resolution: an expired actor token would otherwise + # take the rotation lock and mint (revoking the previous token) for a + # sandbox that's already gone. + sandbox = Sandbox.get_by_id(sandbox_id) + if not sandbox.is_running(): + return False + + try: + token = resolve_coordinated_user_token(integration) + except (ReauthorizationRequired, UserIntegration.DoesNotExist): + logger.info( + "GitHub identity swap skipped: actor integration requires reauthorization", + extra={"run_id": run_id, "user_id": user.id}, + ) + return False + if not token: + return False + + apply_github_credentials_to_sandbox(sandbox, task.repository, token, extra_env=git_identity_env_for_user(user)) + mark_sandbox_identity(scope, "github", integration_id) + mark_sandbox_identity(scope, "github_user", user.id) + logger.info( + "Swapped sandbox GitHub identity to live actor", + extra={ + "run_id": run_id, + "user_id": user.id, + "from_integration_id": current_identity, + "to_integration_id": integration_id, + }, + ) + return True + + def build_sandbox_credentials(ctx: "TaskProcessingContext") -> list[SandboxCredential]: credentials: list[SandboxCredential] = [] if ctx.has_github_credentials: 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 014fd645ff43..d77de086ce3c 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 @@ -462,3 +462,162 @@ def _task(repo): result = _live_sandboxes_for_user_integration(integration.id) assert result == [(str(live_run.id), "sb-live", "org/live")] + + +class TestRefreshSandboxGithubForUser: + """Per-message identity swap: a Slack actor with their own GitHub install + takes over a live sandbox's token and git author identity.""" + + _MODULE = "products.tasks.backend.temporal.process_task.sandbox_credentials" + + def _task_run(self, github_user_integration_id="ui-creator", sandbox_id="sb-1", created_by_id=1): + task = MagicMock() + task.repository = "acme/repo" + task.github_user_integration_id = github_user_integration_id + task.created_by_id = created_by_id + task_run = MagicMock() + task_run.id = "run-gh-1" + task_run.task = task + task_run.state = {"sandbox_id": sandbox_id} if sandbox_id else {} + return task_run + + def _actor(self, user_id=99, name="Bob Builder", email="bob@acme.com"): + user = MagicMock() + user.id = user_id + user.get_full_name.return_value = name + user.email = email + return user + + def _integration(self, integration_id="ui-actor"): + wrapper = MagicMock() + wrapper.integration.id = integration_id + return wrapper + + @pytest.fixture(autouse=True) + def _clear_identity_cache(self): + from django.core.cache import cache + + from products.tasks.backend.temporal.process_task.utils import _sandbox_identity_cache_key + + keys = [_sandbox_identity_cache_key("sb-1", kind) for kind in ("github", "github_user")] + for key in keys: + cache.delete(key) + yield + for key in keys: + cache.delete(key) + + def test_swaps_token_and_git_author_and_marks_identity(self): + from products.tasks.backend.temporal.process_task.sandbox_credentials import refresh_sandbox_github_for_user + from products.tasks.backend.temporal.process_task.utils import PrAuthorshipMode, get_last_sandbox_identity + + sandbox = MagicMock() + sandbox.is_running.return_value = True + actor = self._actor() + with ( + patch(f"{self._MODULE}.get_pr_authorship_mode", return_value=PrAuthorshipMode.USER), + patch(f"{self._MODULE}.is_caller_token_run", return_value=False), + patch(f"{self._MODULE}.get_user_github_integration", return_value=self._integration()), + patch(f"{self._MODULE}.resolve_coordinated_user_token", return_value="ghu_actor_token"), + patch(f"{self._MODULE}.apply_github_credentials_to_sandbox") as mock_apply, + patch("products.tasks.backend.logic.services.sandbox.Sandbox.get_by_id", return_value=sandbox), + ): + assert refresh_sandbox_github_for_user(self._task_run(), actor) is True + + # Token and git author identity travel together through the shared + # apply helper (one env-file read-modify-write). + mock_apply.assert_called_once_with( + sandbox, + "acme/repo", + "ghu_actor_token", + extra_env={ + "GIT_AUTHOR_NAME": "Bob Builder", + "GIT_AUTHOR_EMAIL": "bob@acme.com", + "GIT_COMMITTER_NAME": "Bob Builder", + "GIT_COMMITTER_EMAIL": "bob@acme.com", + }, + ) + assert get_last_sandbox_identity("sb-1", "github") == "ui-actor" + assert get_last_sandbox_identity("sb-1", "github_user") == 99 + + def test_skips_actor_without_covering_integration(self): + from products.tasks.backend.temporal.process_task.sandbox_credentials import refresh_sandbox_github_for_user + from products.tasks.backend.temporal.process_task.utils import PrAuthorshipMode + + with ( + patch(f"{self._MODULE}.get_pr_authorship_mode", return_value=PrAuthorshipMode.USER), + patch(f"{self._MODULE}.is_caller_token_run", return_value=False), + patch(f"{self._MODULE}.get_user_github_integration", return_value=None), + patch(f"{self._MODULE}.apply_github_credentials_to_sandbox") as mock_apply, + ): + assert refresh_sandbox_github_for_user(self._task_run(), self._actor()) is False + + mock_apply.assert_not_called() + + def test_skips_bot_authored_runs(self): + from products.tasks.backend.temporal.process_task.sandbox_credentials import refresh_sandbox_github_for_user + from products.tasks.backend.temporal.process_task.utils import PrAuthorshipMode + + with ( + patch(f"{self._MODULE}.get_pr_authorship_mode", return_value=PrAuthorshipMode.BOT), + patch(f"{self._MODULE}.get_user_github_integration") as mock_resolve, + ): + assert refresh_sandbox_github_for_user(self._task_run(), self._actor()) is False + + mock_resolve.assert_not_called() + + def test_skips_when_sandbox_already_holds_this_identity(self): + from products.tasks.backend.temporal.process_task.sandbox_credentials import refresh_sandbox_github_for_user + from products.tasks.backend.temporal.process_task.utils import PrAuthorshipMode + + with ( + patch(f"{self._MODULE}.get_pr_authorship_mode", return_value=PrAuthorshipMode.USER), + patch(f"{self._MODULE}.is_caller_token_run", return_value=False), + patch(f"{self._MODULE}.get_user_github_integration", return_value=self._integration("ui-creator")), + patch(f"{self._MODULE}.resolve_coordinated_user_token") as mock_token, + ): + assert refresh_sandbox_github_for_user(self._task_run("ui-creator"), self._actor()) is False + + mock_token.assert_not_called() + + def test_switching_back_to_creator_reapplies_their_identity(self): + from products.tasks.backend.temporal.process_task.sandbox_credentials import refresh_sandbox_github_for_user + from products.tasks.backend.temporal.process_task.utils import ( + PrAuthorshipMode, + get_last_sandbox_identity, + mark_sandbox_identity, + ) + + mark_sandbox_identity("sb-1", "github", "ui-actor") + mark_sandbox_identity("sb-1", "github_user", 99) + sandbox = MagicMock() + sandbox.is_running.return_value = True + creator = self._actor(user_id=1, name="Alice A", email="alice@acme.com") + with ( + patch(f"{self._MODULE}.get_pr_authorship_mode", return_value=PrAuthorshipMode.USER), + patch(f"{self._MODULE}.is_caller_token_run", return_value=False), + patch(f"{self._MODULE}.get_user_github_integration", return_value=self._integration("ui-creator")), + patch(f"{self._MODULE}.resolve_coordinated_user_token", return_value="ghu_creator_token"), + patch(f"{self._MODULE}.apply_github_credentials_to_sandbox") as mock_apply, + patch("products.tasks.backend.logic.services.sandbox.Sandbox.get_by_id", return_value=sandbox), + ): + assert refresh_sandbox_github_for_user(self._task_run("ui-creator"), creator) is True + + assert mock_apply.call_args.args[2] == "ghu_creator_token" + assert get_last_sandbox_identity("sb-1", "github") == "ui-creator" + assert get_last_sandbox_identity("sb-1", "github_user") == 1 + + def test_creator_message_on_never_swapped_run_skips_resolution(self): + from products.tasks.backend.temporal.process_task.sandbox_credentials import refresh_sandbox_github_for_user + from products.tasks.backend.temporal.process_task.utils import PrAuthorshipMode + + creator = self._actor(user_id=1) + with ( + patch(f"{self._MODULE}.get_pr_authorship_mode", return_value=PrAuthorshipMode.USER), + patch(f"{self._MODULE}.is_caller_token_run", return_value=False), + patch(f"{self._MODULE}.get_user_github_integration") as mock_resolve, + ): + assert refresh_sandbox_github_for_user(self._task_run(created_by_id=1), creator) is False + + # The dominant case (creator messaging their own never-swapped thread) + # must not pay the integration resolution's DB query per message. + mock_resolve.assert_not_called() diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index 5f1e270613ce..ac9a2f9edafe 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -378,7 +378,7 @@ def should_refresh_mcp_token(scope: str, user_id: int) -> bool: # boot-time one (the task creator / the task's own integration). SANDBOX_IDENTITY_TTL_SECONDS = 7 * 24 * 60 * 60 -SandboxIdentityKind = Literal["mcp"] +SandboxIdentityKind = Literal["mcp", "github", "github_user"] def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: @@ -388,11 +388,15 @@ def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: def mark_sandbox_identity(scope: str, kind: SandboxIdentityKind, value: int | str) -> None: """Record which identity the sandbox currently holds for a credential kind. - ``mcp`` stores the user id the OAuth token was minted for. Written on - every successful rebind so identity transitions (a different Slack actor - taking over the thread) are detected and never silently skipped by the - per-credential freshness rate limits. ``scope`` comes from - ``sandbox_identity_scope``, so a replacement sandbox starts unmarked. + ``mcp`` stores the user id the OAuth token was minted for; ``github`` + stores the UserIntegration id the token and git author belong to, and + ``github_user`` the user id behind it (so same-actor messages can skip + integration resolution). Written on every successful rebind so identity + transitions (a different Slack actor taking over the thread) are detected + and never silently skipped by the per-credential freshness rate limits; + read by the credential refresh loop and token-rotation propagation so + neither reverts a swap. ``scope`` comes from ``sandbox_identity_scope``, + so a replacement sandbox starts unmarked. """ get_tasks_cache().set(_sandbox_identity_cache_key(scope, kind), value, timeout=SANDBOX_IDENTITY_TTL_SECONDS) @@ -917,6 +921,11 @@ def get_git_identity_env_vars(task: Task, state: dict[str, Any] | None = None) - if user is None: return {} + return git_identity_env_for_user(user) + + +def git_identity_env_for_user(user: User) -> dict[str, str]: + """Git author/committer env vars attributing commits to ``user``.""" name = user.get_full_name() or user.first_name or "PostHog User" email = user.email