feat(review_hog): stamp ai_session_id on every review turn generation

products/review_hog/ARCHITECTURE.md

@@ -220,8 +220,13 @@ pr_metadata.head_branch` is threaded (as explicit kwargs, alongside `team_id` /
     "View them in PostHog" deep link to the exact report (`/project/<team>/code-review?review=<report id>`,
     a **permanent public contract** — the frontend URL sync and `report_deep_link` must keep agreeing on it).
     After the publish stage the workflow captures a **`reviewhog_review_completed`** product-analytics event —
-    one per finalized turn (published or stored), carrying repository / PR / trigger / finding-count / PR-size
-    properties (`track_review_completed_activity`). Best-effort: telemetry can never fail a review.
+    one per finalized turn (published or stored), carrying repository / PR / trigger / finding-count / PR-size /
+    sandbox-usage properties (`track_review_completed_activity`): the turn's `sandbox_task_ids` (recovered by
+    workflow-id-prefix, joining to `$ai_generation.task_id` for per-turn dollar cost) plus summed token totals.
+    Best-effort: telemetry can never fail a review. Every turn's LLM calls — sandbox stages (stamped via the
+    agent-server's gateway headers) and one-shots (direct gateway headers) — carry `ai_session_id` (gateway-promoted to the native `$ai_session_id`) =
+    `{report_id}:r{run_index}`, so a turn's generations group as one LLM-analytics session, joinable from the
+    completion event.
 
 ---
 

products/review_hog/backend/reviewer/sandbox/direct_llm.py

@@ -34,6 +34,7 @@ async def run_oneshot_review(
     system_prompt: str,
     model_to_validate: type[_ModelT],
     step_name: str,
+    ai_session_id: str | None = None,
 ) -> _ModelT:
     """Run one review step as a single LLM-gateway call and return its validated output.
 
@@ -65,7 +66,12 @@ class the sandbox path retried on cannot occur). Bedrock fallback is deliberatel
                 output_config=cast(OutputConfigParam, {"effort": ONESHOT_REASONING_EFFORT}),
                 output_format=model_to_validate,
                 metadata={"user_id": f"user-{user_id}"},
-                extra_headers={"x-posthog-property-ai_stage": step_name},
+                extra_headers={
+                    "x-posthog-property-ai_stage": step_name,
+                    # Same key the sandbox path stamps via the agent-server, so the one-shot
+                    # groups with its turn's generations in LLM analytics.
+                    **({"x-posthog-property-ai_session_id": ai_session_id} if ai_session_id else {}),
+                },
                 timeout=_TIMEOUT_SECONDS,
             )
         except APIError as e:

products/review_hog/backend/reviewer/sandbox/executor.py

@@ -19,6 +19,7 @@ async def _run_prompt(
     branch: str | None = None,
     step_name: str = "",
     workflow_id_prefix: str | None = None,
+    ai_session_id: str | None = None,
 ) -> _ModelT:
     """Spawn a single-turn sandbox agent and return its validated end-of-turn.
 
@@ -44,6 +45,7 @@ async def _run_prompt(
             origin_product=TaskOriginProduct.REVIEW_HOG,
             internal=True,
             ai_stage=step_name or None,
+            ai_session_id=ai_session_id,
         )
     except Exception:
         logger.exception("Sandbox execution failed")
@@ -65,6 +67,7 @@ async def run_sandbox_review(
     model_to_validate: type[_ModelT],
     step_name: str = "",
     workflow_id_prefix: str | None = None,
+    ai_session_id: str | None = None,
     runtime_adapter: str | None = None,
     model: str | None = None,
     reasoning_effort: str | None = None,
@@ -105,6 +108,7 @@ async def run_sandbox_review(
         branch=branch,
         step_name=step_name,
         workflow_id_prefix=workflow_id_prefix,
+        ai_session_id=ai_session_id,
     )
 
 
@@ -119,6 +123,7 @@ async def start_sandbox_session(
     model_to_validate: type[_ModelT],
     step_name: str = "",
     workflow_id_prefix: str | None = None,
+    ai_session_id: str | None = None,
     runtime_adapter: str | None = None,
     model: str | None = None,
     reasoning_effort: str | None = None,
@@ -156,6 +161,7 @@ async def start_sandbox_session(
             origin_product=TaskOriginProduct.REVIEW_HOG,
             internal=True,
             ai_stage=step_name or None,
+            ai_session_id=ai_session_id,
         )
     except Exception:
         logger.exception("Sandbox session start failed")

products/review_hog/backend/reviewer/tools/issue_deduplicator.py

@@ -103,6 +103,7 @@ async def deduplicate_issues(
     branch: str,
     repository: str,
     workflow_id_prefix: str | None = None,
+    ai_session_id: str | None = None,
 ) -> list[Issue]:
     """Deduplicate the in-scope issues and return the survivors (the canonical post-dedup set).
 
@@ -155,6 +156,7 @@ async def deduplicate_issues(
             system_prompt=_SYSTEM_PROMPT,
             model_to_validate=IssueDeduplication,
             step_name="dedup",
+            ai_session_id=ai_session_id,
         )
     else:
         deduplication_result = await run_sandbox_review(
@@ -166,6 +168,7 @@ async def deduplicate_issues(
             system_prompt=_SYSTEM_PROMPT,
             model_to_validate=IssueDeduplication,
             step_name="dedup",
+            ai_session_id=ai_session_id,
             workflow_id_prefix=workflow_id_prefix,
             runtime_adapter=DEDUP_RUNTIME_ADAPTER,
             model=DEDUP_MODEL,

products/review_hog/backend/temporal/activities.py

@@ -137,6 +137,7 @@
 from products.signals.backend.artefact_schemas import CodeReview, CodeReviewCounts
 from products.signals.backend.models import SignalReport, SignalReportArtefact
 from products.signals.backend.report_generation.resolve_reviewers import resolve_org_github_login_to_users
+from products.tasks.backend.facade.api import list_sandbox_run_usage
 
 logger = logging.getLogger(__name__)
 
@@ -427,6 +428,13 @@
     return f"{activity.info().workflow_id}:{step_name}".lower()
 
 
+def _turn_ai_session_id(report_id: str, run_index: int) -> str:
+    """One `$ai_session_id` per review turn: the turn's generations (sandbox + one-shot) group
+    under it in LLM analytics, and dashboards join them to `reviewhog_review_completed` by the
+    report-id prefix."""
+    return f"{report_id}:r{run_index}"
+
+
 async def _refresh_status_comment(team_id: int, report_id: str) -> None:
     """Refresh the PR's status comment after this activity persisted progress (debounced, best-effort)."""
     await database_sync_to_async(maybe_refresh_status_comment, thread_sensitive=False)(team_id, report_id)
@@ -749,6 +757,7 @@
                 system_prompt=CHUNKING_SYSTEM_PROMPT,
                 model_to_validate=ChunksList,
                 step_name="chunking",
+                ai_session_id=_turn_ai_session_id(input.report_id, input.run_index),
             )
         else:
             chunks = await run_sandbox_review(
@@ -761,6 +770,7 @@
                 model_to_validate=ChunksList,
                 step_name="chunking",
                 workflow_id_prefix=_sandbox_workflow_id_prefix("chunking"),
+                ai_session_id=_turn_ai_session_id(input.report_id, input.run_index),
                 runtime_adapter=CHUNKING_RUNTIME_ADAPTER,
                 model=CHUNKING_MODEL,
                 reasoning_effort=CHUNKING_REASONING_EFFORT,
@@ -832,6 +842,7 @@
             system_prompt=SELECTION_SYSTEM_PROMPT,
             model_to_validate=PerspectiveSelection,
             step_name="perspective_selection",
+            ai_session_id=_turn_ai_session_id(input.report_id, input.run_index),
         )
     # Persist the normalized plan (exactly what the fan-out runs), not the model's raw output — the
     # progress estimate and the skipped-perspective UI read this artefact as ground truth.
@@ -943,6 +954,7 @@
             model_to_validate=IssuesReview,
             step_name=step_name,
             workflow_id_prefix=_sandbox_workflow_id_prefix(step_name),
+            ai_session_id=_turn_ai_session_id(input.report_id, input.run_index),
             runtime_adapter=REVIEW_RUNTIME_ADAPTER,
             model=REVIEW_MODEL,
             reasoning_effort=REVIEW_REASONING_EFFORT,
@@ -1009,6 +1021,7 @@
             branch=input.branch,
             repository=input.repository,
             workflow_id_prefix=_sandbox_workflow_id_prefix("dedup"),
+            ai_session_id=_turn_ai_session_id(input.report_id, input.run_index),
         )
     issue_ids = await database_sync_to_async(persist_findings, thread_sensitive=False)(
         team_id=input.team_id, report_id=input.report_id, issues=survivors, run_index=input.run_index
@@ -1106,6 +1119,7 @@
                             model_to_validate=IssueValidation,
                             step_name=f"validation-c{input.chunk_id}",
                             workflow_id_prefix=_sandbox_workflow_id_prefix(f"validation-c{input.chunk_id}"),
+                            ai_session_id=_turn_ai_session_id(input.report_id, input.run_index),
                             runtime_adapter=VALIDATION_RUNTIME_ADAPTER,
                             model=VALIDATION_MODEL,
                             reasoning_effort=VALIDATION_REASONING_EFFORT,
@@ -1249,7 +1263,49 @@
     )
 
 
-def _track_review_completed(input: TrackReviewCompletedInput) -> None:
+_SANDBOX_TOKEN_KEYS = ("input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens", "total_tokens")
+
+
+def _collect_turn_sandbox_usage(input: TrackReviewCompletedInput, workflow_id: str | None) -> dict:
+    """The turn's sandbox task ids and summed token usage.
+
+    Every sandbox run a turn spawns carries a Temporal workflow id branded with the review
+    workflow's id (`_sandbox_workflow_id_prefix`), so prefix + the turn's start time recovers
+    exactly this turn's runs: Temporal allows one running turn per PR, and matching on the `:` / `/`
+    separators keeps PR 12's prefix from also matching PR 123. `sandbox_task_ids` joins the event to
+    `$ai_generation` (which carries `task_id`) for per-turn dollar cost; the token sums serve
+    token-level views without a join. Best-effort: a lookup failure nulls these properties, never
+    the capture.
+    """
+    absent: dict = {"sandbox_task_ids": None, "sandbox_run_count": None} | {
+        f"llm_{key}": None for key in _SANDBOX_TOKEN_KEYS
+    }
+    if not workflow_id:
+        return absent
+    try:
+        base = workflow_id.lower()
+        runs = list_sandbox_run_usage(
+            input.team_id,
+            workflow_id_prefixes=[f"{base}:", f"{base}/"],
+            created_after=datetime.datetime.fromisoformat(input.workflow_started_at),
+        )
+        totals = dict.fromkeys(_SANDBOX_TOKEN_KEYS, 0)
+        task_ids: set[str] = set()
+        for run in runs:
+            task_ids.add(str(run["task_id"]))
+            for key in _SANDBOX_TOKEN_KEYS:
+                value = run["token_usage"].get(key)
+                if isinstance(value, int | float) and not isinstance(value, bool):
+                    totals[key] += value
+        return {"sandbox_task_ids": sorted(task_ids), "sandbox_run_count": len(runs)} | {
+            f"llm_{key}": totals[key] for key in _SANDBOX_TOKEN_KEYS
+        }
+    except Exception:
+        logger.exception("Failed to collect sandbox usage for report %s; capturing without it", input.report_id)
+        return absent
+
+
+def _track_review_completed(input: TrackReviewCompletedInput, workflow_id: str | None) -> None:
     report = ReviewReport.objects.for_team(input.team_id).select_related("acting_user", "team").get(id=input.report_id)
     findings = load_turn_findings(team_id=input.team_id, report_id=input.report_id, run_index=input.run_index)
     snapshot = load_pr_snapshot(team_id=input.team_id, report_id=input.report_id, head_sha=input.head_sha)
@@ -1290,16 +1346,17 @@
             # the honest denominator for per-line cost.
             "pr_reviewable_additions": count_reviewable_additions(snapshot.pr_files) if snapshot is not None else None,
             "duration_seconds": duration_seconds,
+            **_collect_turn_sandbox_usage(input, workflow_id),
         },
         groups=groups(team=report.team),
         send_feature_flags=True,
     )
 
 
-def _track_review_completed_safe(input: TrackReviewCompletedInput) -> None:
+def _track_review_completed_safe(input: TrackReviewCompletedInput, workflow_id: str | None) -> None:
     # Analytics must never fail a review: any load/capture failure is logged, not raised.
     try:
-        _track_review_completed(input)
+        _track_review_completed(input, workflow_id)
     except Exception:
         logger.exception("Failed to capture reviewhog_review_completed for report %s; continuing", input.report_id)
 
@@ -1315,7 +1372,8 @@
     provide (one review fans out into many sandbox tasks). Best-effort: analytics must never fail
     a review, so any failure is logged, not raised.
     """
-    await database_sync_to_async(_track_review_completed_safe, thread_sensitive=False)(input)
+    workflow_id = activity.info().workflow_id
+    await database_sync_to_async(_track_review_completed_safe, thread_sensitive=False)(input, workflow_id)
 
 
 # --- The PR's live status comment -------------------------------------------------------------------

products/review_hog/backend/tests/test_review_completed_tracking.py

@@ -19,8 +19,10 @@
     _track_review_completed,
     _track_review_completed_safe,
 )
+from products.tasks.backend.models import Task, TaskRun
 
 _PR_URL = "https://github.com/o/r/pull/7"
+_WORKFLOW_ID = "review-pr:2:o/r:7"
 
 
 def _pr_metadata() -> PRMetadata:
@@ -98,7 +100,7 @@ def test_captures_the_turn_with_review_scoped_properties(self, published: bool)
         )
 
         with patch("products.review_hog.backend.temporal.activities.posthoganalytics.capture") as capture:
-            _track_review_completed(self._tracking_input(report_id, published=published))
+            _track_review_completed(self._tracking_input(report_id, published=published), _WORKFLOW_ID)
 
         capture.assert_called_once()
         kwargs = capture.call_args.kwargs
@@ -120,14 +122,18 @@ def test_captures_the_turn_with_review_scoped_properties(self, published: bool)
         assert props["pr_commits"] == 3
         assert props["pr_reviewable_additions"] == 80
         assert 90 <= props["duration_seconds"] < 600
+        # No sandbox runs for this turn — the cost-linkage properties record an empty turn, not null.
+        assert props["sandbox_task_ids"] == []
+        assert props["sandbox_run_count"] == 0
+        assert props["llm_total_tokens"] == 0
 
     def test_missing_snapshot_still_captures_without_pr_size(self) -> None:
         # A turn whose pr_snapshot is unavailable must still count as a review — size props go
         # null rather than the capture (and with it the review count) being lost.
         report_id = self._review_report()
 
         with patch("products.review_hog.backend.temporal.activities.posthoganalytics.capture") as capture:
-            _track_review_completed(self._tracking_input(report_id, published=False))
+            _track_review_completed(self._tracking_input(report_id, published=False), _WORKFLOW_ID)
 
         capture.assert_called_once()
         props = capture.call_args.kwargs["properties"]
@@ -142,8 +148,8 @@ def test_event_uuid_is_stable_across_retries(self) -> None:
         tracking_input = self._tracking_input(report_id)
 
         with patch("products.review_hog.backend.temporal.activities.posthoganalytics.capture") as capture:
-            _track_review_completed(tracking_input)
-            _track_review_completed(tracking_input)
+            _track_review_completed(tracking_input, _WORKFLOW_ID)
+            _track_review_completed(tracking_input, _WORKFLOW_ID)
 
         first, second = capture.call_args_list
         assert first.kwargs["uuid"]
@@ -158,4 +164,50 @@ def test_capture_failure_is_swallowed(self) -> None:
             "products.review_hog.backend.temporal.activities.posthoganalytics.capture",
             side_effect=RuntimeError("analytics down"),
         ):
-            _track_review_completed_safe(self._tracking_input(report_id))
+            _track_review_completed_safe(self._tracking_input(report_id), _WORKFLOW_ID)
+
+    def test_collects_turn_sandbox_usage(self) -> None:
+        # Cost-per-PR dashboards join this event to $ai_generation through sandbox_task_ids — a
+        # broken prefix lookup or token summing silently zeroes them, and a loose prefix match
+        # would leak PR 74's runs into PR 7's turn.
+        report_id = self._review_report()
+        task = Task.objects.create(team=self.team, title="t", origin_product=Task.OriginProduct.REVIEW_HOG)
+        TaskRun.objects.create(
+            task=task,
+            team=self.team,
+            state={
+                "workflow_id_prefix": f"{_WORKFLOW_ID}:chunking",
+                "token_usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
+            },
+        )
+        TaskRun.objects.create(
+            task=task,
+            team=self.team,
+            state={
+                "workflow_id_prefix": f"{_WORKFLOW_ID}/review:issues-review-p1-c1",
+                "token_usage": {"input_tokens": 1, "output_tokens": 2, "total_tokens": 3},
+            },
+        )
+        # PR 74 shares PR 7's prefix as a string — the separator-suffixed match must exclude it.
+        neighbour_pr = Task.objects.create(team=self.team, title="t", origin_product=Task.OriginProduct.REVIEW_HOG)
+        TaskRun.objects.create(
+            task=neighbour_pr,
+            team=self.team,
+            state={"workflow_id_prefix": f"{_WORKFLOW_ID}4:chunking", "token_usage": {"input_tokens": 100}},
+        )
+        # A prior turn's run for the same PR sits before this turn's window.
+        stale = TaskRun.objects.create(
+            task=task, team=self.team, state={"workflow_id_prefix": f"{_WORKFLOW_ID}:chunking"}
+        )
+        TaskRun.objects.filter(id=stale.id).update(created_at=datetime.now(UTC) - timedelta(hours=2))
+
+        with patch("products.review_hog.backend.temporal.activities.posthoganalytics.capture") as capture:
+            # Real workflow ids carry the repo's casing; the stored prefixes are lowercased.
+            _track_review_completed(self._tracking_input(report_id), _WORKFLOW_ID.upper())
+
+        props = capture.call_args.kwargs["properties"]
+        assert props["sandbox_task_ids"] == [str(task.id)]
+        assert props["sandbox_run_count"] == 2
+        assert props["llm_input_tokens"] == 11
+        assert props["llm_output_tokens"] == 7
+        assert props["llm_total_tokens"] == 18