fix(cloud-agent): use StrEnum for authorship/source constants and typed RunState model

Address PR review feedback: replace string constants with PrAuthorshipMode
and RunSource StrEnum classes, and introduce a RunState Pydantic model for
typed access to TaskRun.state fields.

Author: tatoalo

products/tasks/backend/api.py

@@ -63,7 +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
+from .temporal.process_task.utils import PrAuthorshipMode, cache_github_user_token, parse_run_state
 
 logger = logging.getLogger(__name__)
 
@@ -261,30 +261,26 @@ def run(self, request, pk=None, **kwargs):
                 return Response({"detail": "Invalid resume_from_run_id"}, status=400)
 
             # Derive snapshot_external_id from the validated previous run
-            snapshot_ext_id = (previous_run.state or {}).get("snapshot_external_id")
+            prev_state = parse_run_state(previous_run.state)
             extra_state = {
                 "resume_from_run_id": str(resume_from_run_id),
             }
             if pending_user_message is not None:
                 extra_state["pending_user_message"] = pending_user_message
-            if snapshot_ext_id:
-                extra_state["snapshot_external_id"] = snapshot_ext_id
+            if prev_state.snapshot_external_id:
+                extra_state["snapshot_external_id"] = prev_state.snapshot_external_id
 
-            prev_sandbox_env_id = (previous_run.state or {}).get("sandbox_environment_id")
-            if prev_sandbox_env_id and sandbox_environment_id is None:
-                sandbox_environment_id = prev_sandbox_env_id
+            if prev_state.sandbox_environment_id and sandbox_environment_id is None:
+                sandbox_environment_id = prev_state.sandbox_environment_id
 
-            previous_state = previous_run.state or {}
             if pr_authorship_mode is None:
-                pr_authorship_mode = previous_state.get("pr_authorship_mode")
+                pr_authorship_mode = prev_state.pr_authorship_mode
             if run_source is None:
-                run_source = previous_state.get("run_source")
+                run_source = prev_state.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
+                signal_report_id = prev_state.signal_report_id
+            if branch is None and prev_state.pr_base_branch is not None:
+                branch = prev_state.pr_base_branch
 
         for key, value in {
             "pr_base_branch": branch,
@@ -297,7 +293,7 @@ def run(self, request, pk=None, **kwargs):
                 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:
+        if pr_authorship_mode == PrAuthorshipMode.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:
@@ -322,7 +318,7 @@ 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:
+        if github_user_token and pr_authorship_mode == PrAuthorshipMode.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}")
@@ -862,16 +858,15 @@ def connection_token(self, request, pk=None, **kwargs):
     )
     def command(self, request, pk=None, **kwargs):
         task_run = cast(TaskRun, self.get_object())
-        state = task_run.state or {}
+        run_state = parse_run_state(task_run.state)
 
-        sandbox_url = state.get("sandbox_url")
-        if not sandbox_url:
+        if not run_state.sandbox_url:
             return Response(
                 ErrorResponseSerializer({"error": "No active sandbox for this task run"}).data,
                 status=status.HTTP_400_BAD_REQUEST,
             )
 
-        if not self._is_valid_sandbox_url(sandbox_url):
+        if not self._is_valid_sandbox_url(run_state.sandbox_url):
             logger.warning(f"Blocked request to disallowed sandbox URL for task run {task_run.id}")
             return Response(
                 ErrorResponseSerializer({"error": "Invalid sandbox URL"}).data,
@@ -884,8 +879,6 @@ def command(self, request, pk=None, **kwargs):
             distinct_id=request.user.distinct_id,
         )
 
-        sandbox_connect_token = state.get("sandbox_connect_token")
-
         command_payload: dict = {
             "jsonrpc": request.validated_data["jsonrpc"],
             "method": request.validated_data["method"],
@@ -897,9 +890,9 @@ def command(self, request, pk=None, **kwargs):
 
         try:
             agent_response = self._proxy_command_to_agent_server(
-                sandbox_url=sandbox_url,
+                sandbox_url=run_state.sandbox_url,
                 connection_token=connection_token,
-                sandbox_connect_token=sandbox_connect_token,
+                sandbox_connect_token=run_state.sandbox_connect_token,
                 payload=command_payload,
             )
 

products/tasks/backend/serializers.py

@@ -10,12 +10,7 @@
 
 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,
-)
+from .temporal.process_task.utils import PrAuthorshipMode, RunSource
 
 PRESIGNED_URL_CACHE_TTL = 55 * 60  # 55 minutes (less than 1 hour URL expiry)
 
@@ -363,8 +358,8 @@ 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]
+    PR_AUTHORSHIP_MODE_CHOICES = [mode.value for mode in PrAuthorshipMode]
+    RUN_SOURCE_CHOICES = [source.value for source in RunSource]
 
     mode = serializers.ChoiceField(
         choices=["interactive", "background"],

products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py

@@ -20,6 +20,7 @@
     get_sandbox_api_url,
     get_sandbox_github_token,
     get_sandbox_name_for_task,
+    parse_run_state,
 )
 
 from .get_task_processing_context import TaskProcessingContext
@@ -158,14 +159,15 @@ def get_sandbox_for_repository(input: GetSandboxForRepositoryInput) -> GetSandbo
 
         environment_variables.update(get_git_identity_env_vars(task, ctx.state))
 
+        run_state = parse_run_state(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", "")
-        if resume_from_run_id:
-            environment_variables["POSTHOG_RESUME_RUN_ID"] = resume_from_run_id
+        if run_state.resume_from_run_id:
+            environment_variables["POSTHOG_RESUME_RUN_ID"] = run_state.resume_from_run_id
 
         # Check for resume snapshot (takes priority over integration-level snapshots)
-        resume_snapshot_ext_id = (ctx.state or {}).get("snapshot_external_id")
+        resume_snapshot_ext_id = run_state.snapshot_external_id
         if resume_snapshot_ext_id:
             used_snapshot = True
 

products/tasks/backend/temporal/process_task/utils.py

@@ -1,23 +1,54 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from enum import StrEnum
 from typing import TYPE_CHECKING, Any, Optional
 from urllib.parse import urlparse
 
 from django.conf import settings
 from django.core.cache import cache
 
+from pydantic import BaseModel
+
 from posthog.models.integration import GitHubIntegration, Integration
 from posthog.temporal.oauth import PosthogMcpScopes, has_write_scopes
 
 if TYPE_CHECKING:
     from products.tasks.backend.models import Task
 
 
-PR_AUTHORSHIP_MODE_USER = "user"
-PR_AUTHORSHIP_MODE_BOT = "bot"
-RUN_SOURCE_MANUAL = "manual"
-RUN_SOURCE_SIGNAL_REPORT = "signal_report"
+class PrAuthorshipMode(StrEnum):
+    USER = "user"
+    BOT = "bot"
+
+
+class RunSource(StrEnum):
+    MANUAL = "manual"
+    SIGNAL_REPORT = "signal_report"
+
+
+class RunState(BaseModel, extra="allow"):
+    pr_authorship_mode: PrAuthorshipMode | None = None
+    pr_base_branch: str | None = None
+    run_source: RunSource | None = None
+    signal_report_id: str | None = None
+    resume_from_run_id: str | None = None
+    snapshot_external_id: str | None = None
+    sandbox_id: str | None = None
+    sandbox_url: str | None = None
+    sandbox_connect_token: str | None = None
+    sandbox_environment_id: str | None = None
+    pending_user_message: str | None = None
+    pending_user_message_ts: str | None = None
+    slack_thread_url: str | None = None
+    interaction_origin: str | None = None
+    slack_sent_relay_ids: list[str] | None = None
+
+
+def parse_run_state(state: dict[str, Any] | None) -> RunState:
+    return RunState.model_validate(state or {})
+
+
 GITHUB_USER_TOKEN_CACHE_TTL_SECONDS = 6 * 60 * 60
 
 
@@ -119,8 +150,8 @@ def get_cached_github_user_token(run_id: str) -> str | None:
 def get_sandbox_github_token(
     github_integration_id: int | None, *, run_id: str, state: dict[str, Any] | None = None
 ) -> str | None:
-    authorship_mode = (state or {}).get("pr_authorship_mode")
-    if authorship_mode == PR_AUTHORSHIP_MODE_USER:
+    run_state = parse_run_state(state)
+    if run_state.pr_authorship_mode == PrAuthorshipMode.USER:
         github_user_token = get_cached_github_user_token(run_id)
         if not github_user_token:
             raise ValueError(
@@ -184,22 +215,20 @@ def build_sandbox_environment_variables(
     return env_vars
 
 
-def get_pr_authorship_mode(task: Task, state: dict[str, Any] | None = None) -> str:
+def get_pr_authorship_mode(task: Task, state: dict[str, Any] | None = None) -> PrAuthorshipMode:
     """Return the effective PR authorship mode for a run.
 
     Newer cloud runs store the mode in ``TaskRun.state``. Older user-created
     runs fall back to user authorship so they still get a human git identity.
     """
     from products.tasks.backend.models import Task as TaskModel
 
-    mode = (state or {}).get("pr_authorship_mode")
-    if mode in {PR_AUTHORSHIP_MODE_USER, PR_AUTHORSHIP_MODE_BOT}:
-        return mode
+    run_state = parse_run_state(state)
+    if run_state.pr_authorship_mode is not None:
+        return run_state.pr_authorship_mode
 
     return (
-        PR_AUTHORSHIP_MODE_USER
-        if task.origin_product == TaskModel.OriginProduct.USER_CREATED
-        else PR_AUTHORSHIP_MODE_BOT
+        PrAuthorshipMode.USER if task.origin_product == TaskModel.OriginProduct.USER_CREATED else PrAuthorshipMode.BOT
     )
 
 
@@ -210,7 +239,7 @@ def get_git_identity_env_vars(task: Task, state: dict[str, Any] | None = None) -
     Bot-authored runs fall back to the Dockerfile defaults ("PostHog Code" /
     code@posthog.com).
     """
-    if get_pr_authorship_mode(task, state) != PR_AUTHORSHIP_MODE_USER:
+    if get_pr_authorship_mode(task, state) != PrAuthorshipMode.USER:
         return {}
 
     user = task.created_by