feat(cloud-agent): user-authored PRs

tatoalo · tatoalo · commit ebfdcc426b5c · 2026-04-02T18:00:47.000+01:00
chore: cleaner
diff --git a/products/tasks/backend/api.py b/products/tasks/backend/api.py
@@ -63,6 +63,7 @@
 from .services.connection_token import create_sandbox_connection_token
 from .stream.redis_stream import TaskRunRedisStream, TaskRunStreamError, get_task_run_stream_key
 from .temporal.client import execute_posthog_code_agent_relay_workflow, execute_task_processing_workflow
+from .temporal.process_task.utils import PR_AUTHORSHIP_MODE_USER, cache_github_user_token
 
 logger = logging.getLogger(__name__)
 
@@ -247,6 +248,10 @@ def run(self, request, pk=None, **kwargs):
         resume_from_run_id = request.validated_data.get("resume_from_run_id")
         pending_user_message = request.validated_data.get("pending_user_message")
         sandbox_environment_id = request.validated_data.get("sandbox_environment_id")
+        pr_authorship_mode = request.validated_data.get("pr_authorship_mode")
+        run_source = request.validated_data.get("run_source")
+        signal_report_id = request.validated_data.get("signal_report_id")
+        github_user_token = request.validated_data.get("github_user_token")
 
         extra_state = None
         if resume_from_run_id:
@@ -269,6 +274,32 @@ def run(self, request, pk=None, **kwargs):
             if prev_sandbox_env_id and sandbox_environment_id is None:
                 sandbox_environment_id = prev_sandbox_env_id
 
+            previous_state = previous_run.state or {}
+            if pr_authorship_mode is None:
+                pr_authorship_mode = previous_state.get("pr_authorship_mode")
+            if run_source is None:
+                run_source = previous_state.get("run_source")
+            if signal_report_id is None:
+                signal_report_id = previous_state.get("signal_report_id")
+            if branch is None:
+                previous_base_branch = previous_state.get("pr_base_branch")
+                if isinstance(previous_base_branch, str):
+                    branch = previous_base_branch
+
+        for key, value in {
+            "pr_base_branch": branch,
+            "pr_authorship_mode": pr_authorship_mode,
+            "run_source": run_source,
+            "signal_report_id": signal_report_id,
+        }.items():
+            if value is not None:
+                extra_state = extra_state or {}
+                extra_state[key] = value
+
+        # Only require a user token when the task has a repo (no-repo cloud runs skip GitHub operations)
+        if pr_authorship_mode == PR_AUTHORSHIP_MODE_USER and task.repository and not github_user_token:
+            return Response({"detail": "github_user_token is required for user-authored cloud runs"}, status=400)
+
         if sandbox_environment_id is not None:
             sandbox_environment = SandboxEnvironment.objects.filter(id=sandbox_environment_id, team=task.team).first()
             if not sandbox_environment:
@@ -291,6 +322,9 @@ def run(self, request, pk=None, **kwargs):
 
         task_run = task.create_run(mode=mode, branch=branch, extra_state=extra_state)
 
+        if github_user_token and pr_authorship_mode == PR_AUTHORSHIP_MODE_USER:
+            cache_github_user_token(str(task_run.id), github_user_token)
+
         logger.info(f"Triggering workflow for task {task.id}, run {task_run.id}")
 
         self._trigger_workflow(task, task_run)
@@ -379,34 +413,50 @@ def update(self, request, *args, **kwargs):
     )
     def partial_update(self, request, *args, **kwargs):
         task_run = cast(TaskRun, self.get_object())
-        old_status = task_run.status
-
-        # Update fields from validated data
-        for key, value in request.validated_data.items():
-            setattr(task_run, key, value)
+        has_output_merge = "output" in request.validated_data and isinstance(request.validated_data["output"], dict)
 
-        new_status = request.validated_data.get("status")
-        terminal_statuses = [
-            TaskRun.Status.COMPLETED,
-            TaskRun.Status.FAILED,
-            TaskRun.Status.CANCELLED,
-        ]
-
-        # Auto-set completed_at if status is completed or failed
-        if new_status in terminal_statuses:
-            if not task_run.completed_at:
-                task_run.completed_at = timezone.now()
-
-            # Signal Temporal workflow if status changed to terminal state
-            if old_status != new_status:
-                self._signal_workflow_completion(
-                    task_run,
-                    new_status,
-                    request.validated_data.get("error_message"),
-                )
-
-        task_run.save()
-        self._post_slack_update_for_pr(task_run)
+        with transaction.atomic():
+            # Re-fetch with row lock when merging output to prevent concurrent
+            # PATCHes (e.g. branch sync + PR URL) from clobbering each other.
+            if has_output_merge:
+                task_run = TaskRun.objects.select_for_update().get(pk=task_run.pk)
+
+            old_status = task_run.status
+            old_pr_url = (task_run.output or {}).get("pr_url") if isinstance(task_run.output, dict) else None
+
+            # Update fields from validated data
+            for key, value in request.validated_data.items():
+                if key == "output" and isinstance(value, dict):
+                    existing_output = task_run.output if isinstance(task_run.output, dict) else {}
+                    setattr(task_run, key, {**existing_output, **value})
+                    continue
+                setattr(task_run, key, value)
+
+            new_status = request.validated_data.get("status")
+            terminal_statuses = [
+                TaskRun.Status.COMPLETED,
+                TaskRun.Status.FAILED,
+                TaskRun.Status.CANCELLED,
+            ]
+
+            # Auto-set completed_at if status is completed or failed
+            if new_status in terminal_statuses:
+                if not task_run.completed_at:
+                    task_run.completed_at = timezone.now()
+
+            task_run.save()
+
+        # Signal Temporal and post Slack updates after commit to avoid
+        # holding the row lock during external calls.
+        if new_status in terminal_statuses and old_status != new_status:
+            self._signal_workflow_completion(
+                task_run,
+                new_status,
+                request.validated_data.get("error_message"),
+            )
+        new_pr_url = (task_run.output or {}).get("pr_url") if isinstance(task_run.output, dict) else None
+        if new_pr_url and new_pr_url != old_pr_url:
+            self._post_slack_update_for_pr(task_run)
 
         return Response(TaskRunDetailSerializer(task_run, context=self.get_serializer_context()).data)
 
diff --git a/products/tasks/backend/sandbox/images/Dockerfile.sandbox-base b/products/tasks/backend/sandbox/images/Dockerfile.sandbox-base
@@ -101,9 +101,11 @@ RUN chmod +x /tmp/install-skills.sh && \
     /tmp/install-skills.sh /tmp/skills && \
     rm -rf /tmp/install-skills.sh /tmp/skills
 
-# Configure git for commits - TODO: Use user's email and name
-RUN git config --global user.email "array@posthog.com" && \
-    git config --global user.name "Array"
+# Default git identity for bot-authored commits.
+# Runs configured for user-authored pull requests receive
+# GIT_AUTHOR_*/GIT_COMMITTER_* env vars that override these defaults.
+RUN git config --global user.email "code@posthog.com" && \
+    git config --global user.name "PostHog Code"
 
 # This is required for the Claude Code SDK to allow --dangerously-skip-permissions as the root user
 ENV IS_SANDBOX=1
diff --git a/products/tasks/backend/sandbox/images/Dockerfile.sandbox-notebook b/products/tasks/backend/sandbox/images/Dockerfile.sandbox-notebook
@@ -82,9 +82,11 @@ RUN cd /scripts && \
     -o runAgent.mjs && \
     chmod +x runAgent.mjs
 
-# Configure git for commits - TODO: Use user's email and name
-RUN git config --global user.email "array@posthog.com" && \
-    git config --global user.name "Array"
+# Default git identity for bot-authored commits.
+# Runs configured for user-authored pull requests receive
+# GIT_AUTHOR_*/GIT_COMMITTER_* env vars that override these defaults.
+RUN git config --global user.email "code@posthog.com" && \
+    git config --global user.name "PostHog Code"
 
 # This is required for the Claude Code SDK to allow --dangerously-skip-permissions as the root user
 ENV IS_SANDBOX=1
diff --git a/products/tasks/backend/serializers.py b/products/tasks/backend/serializers.py
@@ -10,6 +10,12 @@
 
 from .models import SandboxEnvironment, Task, TaskRun
 from .services.title_generator import generate_task_title
+from .temporal.process_task.utils import (
+    PR_AUTHORSHIP_MODE_BOT,
+    PR_AUTHORSHIP_MODE_USER,
+    RUN_SOURCE_MANUAL,
+    RUN_SOURCE_SIGNAL_REPORT,
+)
 
 PRESIGNED_URL_CACHE_TTL = 55 * 60  # 55 minutes (less than 1 hour URL expiry)
 
@@ -351,6 +357,9 @@ class ConnectionTokenResponseSerializer(serializers.Serializer):
 class TaskRunCreateRequestSerializer(serializers.Serializer):
     """Request body for creating a new task run"""
 
+    PR_AUTHORSHIP_MODE_CHOICES = [PR_AUTHORSHIP_MODE_USER, PR_AUTHORSHIP_MODE_BOT]
+    RUN_SOURCE_CHOICES = [RUN_SOURCE_MANUAL, RUN_SOURCE_SIGNAL_REPORT]
+
     mode = serializers.ChoiceField(
         choices=["interactive", "background"],
         required=False,
@@ -380,6 +389,31 @@ class TaskRunCreateRequestSerializer(serializers.Serializer):
         default=None,
         help_text="Optional sandbox environment to apply for this cloud run.",
     )
+    pr_authorship_mode = serializers.ChoiceField(
+        choices=PR_AUTHORSHIP_MODE_CHOICES,
+        required=False,
+        default=None,
+        help_text="Whether pull requests for this run should be authored by the user or the bot.",
+    )
+    run_source = serializers.ChoiceField(
+        choices=RUN_SOURCE_CHOICES,
+        required=False,
+        default=None,
+        help_text="High-level source that triggered this run, used to distinguish manual and signal-based cloud runs.",
+    )
+    signal_report_id = serializers.CharField(
+        required=False,
+        default=None,
+        allow_blank=False,
+        help_text="Optional signal report identifier when this run was started from Inbox.",
+    )
+    github_user_token = serializers.CharField(
+        required=False,
+        default=None,
+        allow_blank=False,
+        write_only=True,
+        help_text="Ephemeral GitHub user token from PostHog Code for user-authored cloud pull requests.",
+    )
 
 
 class TaskRunCommandRequestSerializer(serializers.Serializer):
diff --git a/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py b/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py
@@ -17,7 +17,8 @@
 from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution
 from products.tasks.backend.temporal.process_task.utils import (
     build_sandbox_environment_variables,
-    get_github_token,
+    get_git_identity_env_vars,
+    get_sandbox_github_token,
     get_sandbox_name_for_task,
 )
 
@@ -70,7 +71,14 @@ def create_sandbox_from_snapshot(input: CreateSandboxFromSnapshotInput) -> Creat
         github_token = ""
         if ctx.github_integration_id is not None:
             try:
-                github_token = get_github_token(ctx.github_integration_id) or ""
+                github_token = (
+                    get_sandbox_github_token(
+                        ctx.github_integration_id,
+                        run_id=ctx.run_id,
+                        state=ctx.state,
+                    )
+                    or ""
+                )
             except Exception as e:
                 raise GitHubAuthenticationError(
                     f"Failed to get GitHub token for integration {ctx.github_integration_id}",
@@ -99,6 +107,7 @@ def create_sandbox_from_snapshot(input: CreateSandboxFromSnapshotInput) -> Creat
             team_id=ctx.team_id,
             sandbox_environment=sandbox_env,
         )
+        environment_variables.update(get_git_identity_env_vars(task, ctx.state))
 
         config = SandboxConfig(
             name=get_sandbox_name_for_task(ctx.task_id),
diff --git a/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py b/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py
@@ -16,8 +16,9 @@
 from products.tasks.backend.temporal.oauth import create_oauth_access_token
 from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution
 from products.tasks.backend.temporal.process_task.utils import (
-    get_github_token,
+    get_git_identity_env_vars,
     get_sandbox_api_url,
+    get_sandbox_github_token,
     get_sandbox_name_for_task,
 )
 
@@ -31,6 +32,7 @@
     "POSTHOG_PROJECT_ID",
     "JWT_PUBLIC_KEY",
     "GITHUB_TOKEN",
+    "GH_TOKEN",
     "LLM_GATEWAY_URL",
     "POSTHOG_RESUME_RUN_ID",
 }
@@ -85,7 +87,14 @@ def get_sandbox_for_repository(input: GetSandboxForRepositoryInput) -> GetSandbo
         if has_repo:
             assert github_integration_id is not None
             try:
-                github_token = get_github_token(github_integration_id) or ""
+                github_token = (
+                    get_sandbox_github_token(
+                        github_integration_id,
+                        run_id=ctx.run_id,
+                        state=ctx.state,
+                    )
+                    or ""
+                )
             except Exception as e:
                 raise GitHubAuthenticationError(
                     f"Failed to get GitHub token for integration {github_integration_id}",
@@ -138,10 +147,13 @@ def get_sandbox_for_repository(input: GetSandboxForRepositoryInput) -> GetSandbo
 
         if github_token:
             environment_variables["GITHUB_TOKEN"] = github_token
+            environment_variables["GH_TOKEN"] = github_token
 
         if settings.SANDBOX_LLM_GATEWAY_URL:
             environment_variables["LLM_GATEWAY_URL"] = settings.SANDBOX_LLM_GATEWAY_URL
 
+        environment_variables.update(get_git_identity_env_vars(task, ctx.state))
+
         # Set resume run ID independently of snapshot so conversation history
         # can be rebuilt from logs even when the filesystem snapshot has expired.
         resume_from_run_id = (ctx.state or {}).get("resume_from_run_id", "")
diff --git a/products/tasks/backend/temporal/process_task/tests/test_utils.py b/products/tasks/backend/temporal/process_task/tests/test_utils.py
diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py
diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py
