From 1c7182ccfc69315e0f5437962fae09015dfedb78 Mon Sep 17 00:00:00 2001 From: Alex Lebedev Date: Thu, 30 Jul 2026 16:50:28 +0200 Subject: [PATCH 1/3] feat(review_hog): link review turns to their LLM cost on the event MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Recover the turn's sandbox runs by the Temporal workflow-id branding every run already carries (_sandbox_workflow_id_prefix), through a new read-only tasks facade helper, and stamp sandbox_task_ids + summed token totals + run count onto reviewhog_review_completed. sandbox_task_ids joins to $ai_generation.task_id for exact per-turn dollar cost — the gateway's cost numbers stay the single source of truth instead of duplicating pricing server-side. Prefix matching is separator-suffixed so PR 12 never matches PR 123, bounded to the turn's window (Temporal allows one running turn per PR), and best-effort: a lookup failure nulls the linkage properties, never the capture. Generated-By: PostHog Code Task-Id: 77db6432-21d6-400a-ad79-7ce81def58cf --- products/review_hog/ARCHITECTURE.md | 6 +- .../review_hog/backend/temporal/activities.py | 53 ++++++++++++++-- .../tests/test_review_completed_tracking.py | 62 +++++++++++++++++-- products/tasks/backend/facade/api.py | 25 ++++++++ 4 files changed, 135 insertions(+), 11 deletions(-) diff --git a/products/review_hog/ARCHITECTURE.md b/products/review_hog/ARCHITECTURE.md index 1109c1da1b54..028e85c2e0ce 100644 --- a/products/review_hog/ARCHITECTURE.md +++ b/products/review_hog/ARCHITECTURE.md @@ -220,8 +220,10 @@ pr_metadata.head_branch` is threaded (as explicit kwargs, alongside `team_id` / "View them in PostHog" deep link to the exact report (`/project//code-review?review=`, 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. --- diff --git a/products/review_hog/backend/temporal/activities.py b/products/review_hog/backend/temporal/activities.py index 366010a81d82..f33b2ba4fb7e 100644 --- a/products/review_hog/backend/temporal/activities.py +++ b/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__) @@ -1249,7 +1250,49 @@ async def publish_review_activity(input: PublishInput) -> PublishResult: ) -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 +1333,17 @@ def _track_review_completed(input: TrackReviewCompletedInput) -> None: # 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 +1359,8 @@ async def track_review_completed_activity(input: TrackReviewCompletedInput) -> N 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 ------------------------------------------------------------------- diff --git a/products/review_hog/backend/tests/test_review_completed_tracking.py b/products/review_hog/backend/tests/test_review_completed_tracking.py index e513ea64c53f..d339c449f6ae 100644 --- a/products/review_hog/backend/tests/test_review_completed_tracking.py +++ b/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,6 +122,10 @@ 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 @@ -127,7 +133,7 @@ def test_missing_snapshot_still_captures_without_pr_size(self) -> None: 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 diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 7d20056f36cb..a6123f8788d3 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -520,6 +520,31 @@ def get_task_id_for_run(run_id: str | UUID, team_id: int) -> UUID | None: return TaskRun.objects.filter(id=run_id, team_id=team_id).values_list("task_id", flat=True).first() +def list_sandbox_run_usage( + team_id: int, *, workflow_id_prefixes: Sequence[str], created_after: datetime +) -> list[dict[str, Any]]: + """Task ids + token usage for the team's runs whose ``workflow_id_prefix`` starts with any prefix. + + Products that brand their sandbox runs' Temporal workflow ids (via ``workflow_id_prefix`` at + creation) can recover everything one logical operation spawned — e.g. a ReviewHog review turn + summing its LLM usage — without holding run ids themselves. ``created_after`` bounds the scan + to the operation's window. Read-only; returns ``{"task_id": UUID, "token_usage": dict}`` rows. + """ + if not workflow_id_prefixes: + return [] + prefix_q = Q() + for prefix in workflow_id_prefixes: + prefix_q |= Q(state__workflow_id_prefix__startswith=prefix) + rows = TaskRun.objects.filter(prefix_q, team_id=team_id, created_at__gte=created_after).values_list( + "task_id", "state" + ) + usage_rows: list[dict[str, Any]] = [] + for task_id, state in rows: + token_usage = state.get("token_usage") if isinstance(state, dict) else None + usage_rows.append({"task_id": task_id, "token_usage": token_usage if isinstance(token_usage, dict) else {}}) + return usage_rows + + def task_exists(task_id: str | UUID, team_id: int) -> bool: """Whether a (non-deleted) task exists for the team.""" return Task.objects.filter(id=task_id, team_id=team_id).exists() From 80efef3aa091648c302824acf2a57c3ca58874ce Mon Sep 17 00:00:00 2001 From: Alex Lebedev Date: Thu, 30 Jul 2026 17:07:09 +0200 Subject: [PATCH 2/3] feat(review_hog): stamp $ai_session_id on every review turn generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Thread an ai_session_id through the tasks facade into run state (mirroring ai_stage exactly), and set it from ReviewHog to {report_id}:r{run_index} on every LLM call a turn makes — sandbox stages via Task.create_and_run, one-shots via a direct gateway header. $ai_session_id is first-class LLM-analytics taxonomy (materialized column + bloom-filter index), so a turn's generations become directly filterable by review turn and group as one session in AI observability. The sandbox path lights up once the agent-server forwards the new state key (companion PostHog/code change); the one-shot path works immediately. Generated-By: PostHog Code Task-Id: 77db6432-21d6-400a-ad79-7ce81def58cf --- products/review_hog/ARCHITECTURE.md | 5 ++++- .../backend/reviewer/sandbox/direct_llm.py | 8 +++++++- .../review_hog/backend/reviewer/sandbox/executor.py | 6 ++++++ .../backend/reviewer/tools/issue_deduplicator.py | 3 +++ products/review_hog/backend/temporal/activities.py | 13 +++++++++++++ .../logic/services/custom_prompt_internals.py | 2 ++ .../services/custom_prompt_multi_turn_runner.py | 4 ++++ products/tasks/backend/models.py | 8 ++++++++ products/tasks/backend/tests/test_models.py | 7 +++++-- 9 files changed, 52 insertions(+), 4 deletions(-) diff --git a/products/review_hog/ARCHITECTURE.md b/products/review_hog/ARCHITECTURE.md index 028e85c2e0ce..5f72a570756a 100644 --- a/products/review_hog/ARCHITECTURE.md +++ b/products/review_hog/ARCHITECTURE.md @@ -223,7 +223,10 @@ pr_metadata.head_branch` is threaded (as explicit kwargs, alongside `team_id` / 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. + 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` = + `{report_id}:r{run_index}`, so a turn's generations group as one LLM-analytics session, joinable from the + completion event. --- diff --git a/products/review_hog/backend/reviewer/sandbox/direct_llm.py b/products/review_hog/backend/reviewer/sandbox/direct_llm.py index b46e3643cb2d..cc9573440602 100644 --- a/products/review_hog/backend/reviewer/sandbox/direct_llm.py +++ b/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: diff --git a/products/review_hog/backend/reviewer/sandbox/executor.py b/products/review_hog/backend/reviewer/sandbox/executor.py index bafab6ff7df2..2f91f7d10ee8 100644 --- a/products/review_hog/backend/reviewer/sandbox/executor.py +++ b/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") diff --git a/products/review_hog/backend/reviewer/tools/issue_deduplicator.py b/products/review_hog/backend/reviewer/tools/issue_deduplicator.py index 77b99d6813c3..2259e473decf 100644 --- a/products/review_hog/backend/reviewer/tools/issue_deduplicator.py +++ b/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, diff --git a/products/review_hog/backend/temporal/activities.py b/products/review_hog/backend/temporal/activities.py index f33b2ba4fb7e..fe1c7d157a8c 100644 --- a/products/review_hog/backend/temporal/activities.py +++ b/products/review_hog/backend/temporal/activities.py @@ -428,6 +428,13 @@ def _sandbox_workflow_id_prefix(step_name: str) -> str: 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) @@ -750,6 +757,7 @@ async def split_chunks_activity(input: SandboxStageInput) -> list[int]: 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( @@ -762,6 +770,7 @@ async def split_chunks_activity(input: SandboxStageInput) -> list[int]: 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, @@ -833,6 +842,7 @@ async def select_perspectives_activity(input: SelectPerspectivesInput) -> Perspe 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. @@ -944,6 +954,7 @@ async def review_chunk_activity(input: ReviewChunkInput) -> bool: 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, @@ -1010,6 +1021,7 @@ async def dedup_activity(input: SandboxStageInput) -> DedupResult: 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 @@ -1107,6 +1119,7 @@ async def validate_chunk_activity(input: ValidateChunkInput) -> ValidateChunkRes 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, diff --git a/products/tasks/backend/logic/services/custom_prompt_internals.py b/products/tasks/backend/logic/services/custom_prompt_internals.py index 11d9fd5f2198..44eef6933086 100644 --- a/products/tasks/backend/logic/services/custom_prompt_internals.py +++ b/products/tasks/backend/logic/services/custom_prompt_internals.py @@ -141,6 +141,7 @@ async def create_task_and_trigger( origin_product: Task.OriginProduct | None = None, signal_report_id: str | None = None, ai_stage: str | None = None, + ai_session_id: str | None = None, internal: bool = False, workflow_id_prefix: str | None = None, ): @@ -163,6 +164,7 @@ async def create_task_and_trigger( branch=branch, signal_report_id=signal_report_id, ai_stage=ai_stage, + ai_session_id=ai_session_id, posthog_mcp_scopes=posthog_mcp_scopes, sandbox_environment_id=context.sandbox_environment_id, model=context.model, diff --git a/products/tasks/backend/logic/services/custom_prompt_multi_turn_runner.py b/products/tasks/backend/logic/services/custom_prompt_multi_turn_runner.py index b5149739ee0c..b05b63a2bebd 100644 --- a/products/tasks/backend/logic/services/custom_prompt_multi_turn_runner.py +++ b/products/tasks/backend/logic/services/custom_prompt_multi_turn_runner.py @@ -62,6 +62,7 @@ async def start( origin_product: Task.OriginProduct | None = None, signal_report_id: str | None = None, ai_stage: str | None = None, + ai_session_id: str | None = None, internal: bool = False, on_task_run_created: Callable[[TaskRun], Awaitable[None]] | None = None, max_poll_seconds: int | None = None, @@ -96,6 +97,7 @@ async def start( origin_product=origin_product, signal_report_id=signal_report_id, ai_stage=ai_stage, + ai_session_id=ai_session_id, internal=internal, on_task_run_created=on_task_run_created, max_poll_seconds=max_poll_seconds, @@ -150,6 +152,7 @@ async def start_raw( origin_product: Task.OriginProduct | None = None, signal_report_id: str | None = None, ai_stage: str | None = None, + ai_session_id: str | None = None, internal: bool = False, on_task_run_created: Callable[[TaskRun], Awaitable[None]] | None = None, max_poll_seconds: int | None = None, @@ -170,6 +173,7 @@ async def start_raw( origin_product=origin_product, signal_report_id=signal_report_id, ai_stage=ai_stage, + ai_session_id=ai_session_id, internal=internal, workflow_id_prefix=workflow_id_prefix, ) diff --git a/products/tasks/backend/models.py b/products/tasks/backend/models.py index 0d3aa5d3737e..53bdace0bbd4 100644 --- a/products/tasks/backend/models.py +++ b/products/tasks/backend/models.py @@ -500,6 +500,7 @@ def _build_task( branch: str | None = None, signal_report_id: str | None = None, ai_stage: str | None = None, + ai_session_id: str | None = None, sandbox_environment_id: str | None = None, internal: bool = False, output_schema: type[BaseModel] | dict | None = None, @@ -654,6 +655,11 @@ def _build_task( if ai_stage: extra_state["ai_stage"] = ai_stage + # Same channel as `ai_stage`, lifted as `$ai_session_id` — groups the run's generations + # under the caller's logical operation (e.g. one ReviewHog review turn) in LLM analytics. + if ai_session_id: + extra_state["ai_session_id"] = ai_session_id + if initial_permission_mode: extra_state["initial_permission_mode"] = initial_permission_mode @@ -780,6 +786,7 @@ def create_and_run( sandbox_timeout_seconds: int | None = None, inactivity_timeout_seconds: int | None = None, ai_stage: str | None = None, + ai_session_id: str | None = None, wizard_config: dict | None = None, wizard_head_branch: str | None = None, pending_user_message: str | None = None, @@ -813,6 +820,7 @@ def create_and_run( sandbox_timeout_seconds=sandbox_timeout_seconds, inactivity_timeout_seconds=inactivity_timeout_seconds, ai_stage=ai_stage, + ai_session_id=ai_session_id, wizard_config=wizard_config, wizard_head_branch=wizard_head_branch, pending_user_message=pending_user_message, diff --git a/products/tasks/backend/tests/test_models.py b/products/tasks/backend/tests/test_models.py index 3c180fba7252..354ed8324ac2 100644 --- a/products/tasks/backend/tests/test_models.py +++ b/products/tasks/backend/tests/test_models.py @@ -152,7 +152,7 @@ def test_create_and_run_threads_initial_permission_mode_into_state(self, mock_ex self.assertEqual(task.origin_product, Task.OriginProduct.SLACK) @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") - def test_create_and_run_threads_ai_stage_into_state(self, mock_execute_workflow): + def test_create_and_run_threads_ai_analytics_stamps_into_state(self, mock_execute_workflow): user = User.objects.create(email="test@test.com") Integration.objects.create(team=self.team, kind="github", config={}) @@ -165,14 +165,16 @@ def test_create_and_run_threads_ai_stage_into_state(self, mock_execute_workflow) user_id=user.id, repository="posthog/posthog", ai_stage="research", + ai_session_id="report-1:r2", ) run_id = mock_execute_workflow.call_args.kwargs["run_id"] task_run = TaskRun.objects.get(id=run_id) self.assertEqual(task_run.state["ai_stage"], "research") + self.assertEqual(task_run.state["ai_session_id"], "report-1:r2") @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") - def test_create_and_run_omits_ai_stage_when_not_provided(self, mock_execute_workflow): + def test_create_and_run_omits_ai_analytics_stamps_when_not_provided(self, mock_execute_workflow): user = User.objects.create(email="test@test.com") Integration.objects.create(team=self.team, kind="github", config={}) @@ -189,6 +191,7 @@ def test_create_and_run_omits_ai_stage_when_not_provided(self, mock_execute_work run_id = mock_execute_workflow.call_args.kwargs["run_id"] task_run = TaskRun.objects.get(id=run_id) self.assertNotIn("ai_stage", task_run.state) + self.assertNotIn("ai_session_id", task_run.state) @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") def test_create_and_run_omits_permission_mode_when_not_provided(self, mock_execute_workflow): From 516ab9e25ec0db01c70b44393efaf00575cb0bcc Mon Sep 17 00:00:00 2001 From: Alex Lebedev Date: Thu, 30 Jul 2026 17:10:50 +0200 Subject: [PATCH 3/3] feat(review_hog): promote ai_session_id to $ai_session_id in the gateway MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reserved $-keys are stripped at the ai-gateway header boundary, so callers transport the session key unreserved (ai_session_id, mirroring ai_stage) and the llm-gateway callback promotes it to the native $ai_session_id — materialized column, bloom-filter index, AIO session grouping — where the event's properties are final, on both success and failure captures. Covered by a new callback test. Generated-By: PostHog Code Task-Id: 77db6432-21d6-400a-ad79-7ce81def58cf --- products/review_hog/ARCHITECTURE.md | 2 +- .../backend/reviewer/sandbox/direct_llm.py | 2 +- products/tasks/backend/models.py | 5 ++-- .../src/llm_gateway/callbacks/posthog.py | 12 +++++++++ .../tests/callbacks/test_posthog.py | 27 +++++++++++++++++++ 5 files changed, 44 insertions(+), 4 deletions(-) diff --git a/products/review_hog/ARCHITECTURE.md b/products/review_hog/ARCHITECTURE.md index 5f72a570756a..e60c0e9b399c 100644 --- a/products/review_hog/ARCHITECTURE.md +++ b/products/review_hog/ARCHITECTURE.md @@ -224,7 +224,7 @@ pr_metadata.head_branch` is threaded (as explicit kwargs, alongside `team_id` / 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` = + 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. diff --git a/products/review_hog/backend/reviewer/sandbox/direct_llm.py b/products/review_hog/backend/reviewer/sandbox/direct_llm.py index cc9573440602..1325412087e9 100644 --- a/products/review_hog/backend/reviewer/sandbox/direct_llm.py +++ b/products/review_hog/backend/reviewer/sandbox/direct_llm.py @@ -70,7 +70,7 @@ class the sandbox path retried on cannot occur). Bedrock fallback is deliberatel "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 {}), + **({"x-posthog-property-ai_session_id": ai_session_id} if ai_session_id else {}), }, timeout=_TIMEOUT_SECONDS, ) diff --git a/products/tasks/backend/models.py b/products/tasks/backend/models.py index 53bdace0bbd4..bed24d539caa 100644 --- a/products/tasks/backend/models.py +++ b/products/tasks/backend/models.py @@ -655,8 +655,9 @@ def _build_task( if ai_stage: extra_state["ai_stage"] = ai_stage - # Same channel as `ai_stage`, lifted as `$ai_session_id` — groups the run's generations - # under the caller's logical operation (e.g. one ReviewHog review turn) in LLM analytics. + # Same channel as `ai_stage`, transported unreserved ($-keys are stripped at gateway + # header boundaries) and promoted to the native `$ai_session_id` by the LLM gateway — + # groups the run's generations under the caller's logical operation in LLM analytics. if ai_session_id: extra_state["ai_session_id"] = ai_session_id diff --git a/services/llm-gateway/src/llm_gateway/callbacks/posthog.py b/services/llm-gateway/src/llm_gateway/callbacks/posthog.py index c64578991d46..941b73de5673 100644 --- a/services/llm-gateway/src/llm_gateway/callbacks/posthog.py +++ b/services/llm-gateway/src/llm_gateway/callbacks/posthog.py @@ -255,6 +255,12 @@ async def _on_success( for key, value in posthog_properties.items(): properties[key] = value + # Reserved $-keys can't travel from callers (the ai-gateway strips them at its header + # boundary), so the session key arrives as `ai_session_id` and is promoted to the native + # `$ai_session_id` here, where the event's properties are final. + if properties.get("ai_session_id") and not properties.get("$ai_session_id"): + properties["$ai_session_id"] = properties["ai_session_id"] + posthog_flags = get_posthog_flags() or {} if isinstance(posthog_flags, dict): for flag_key, variant in posthog_flags.items(): @@ -343,6 +349,12 @@ async def _on_failure( for key, value in posthog_properties.items(): properties[key] = value + # Reserved $-keys can't travel from callers (the ai-gateway strips them at its header + # boundary), so the session key arrives as `ai_session_id` and is promoted to the native + # `$ai_session_id` here, where the event's properties are final. + if properties.get("ai_session_id") and not properties.get("$ai_session_id"): + properties["$ai_session_id"] = properties["ai_session_id"] + posthog_flags = get_posthog_flags() or {} if isinstance(posthog_flags, dict): for flag_key, variant in posthog_flags.items(): diff --git a/services/llm-gateway/tests/callbacks/test_posthog.py b/services/llm-gateway/tests/callbacks/test_posthog.py index 3dc2c69a6bf8..eb089e01d67b 100644 --- a/services/llm-gateway/tests/callbacks/test_posthog.py +++ b/services/llm-gateway/tests/callbacks/test_posthog.py @@ -160,6 +160,33 @@ async def test_effort_is_gateway_owned( else: assert props["$ai_effort"] == expected + @pytest.mark.asyncio + async def test_on_success_promotes_ai_session_id( + self, + callback: PostHogCallback, + auth_user: AuthenticatedUser, + standard_logging_object: dict, + mock_posthog_client: tuple, + ) -> None: + # Callers transport the session key unreserved (the ai-gateway strips $-keys at its + # header boundary); losing this promotion silently breaks session grouping in LLM analytics. + _, mock_client = mock_posthog_client + kwargs = {"standard_logging_object": standard_logging_object, "litellm_params": {}} + + with ( + patch("llm_gateway.callbacks.posthog.get_auth_user", return_value=auth_user), + patch("llm_gateway.callbacks.posthog.get_product", return_value="posthog_code"), + patch( + "llm_gateway.callbacks.posthog.get_posthog_properties", + return_value={"ai_session_id": "report-1:r2"}, + ), + ): + await callback._on_success(kwargs, None, 0.0, 1.0, end_user_id=None) + + props = mock_client.capture.call_args.kwargs["properties"] + assert props["$ai_session_id"] == "report-1:r2" + assert props["ai_session_id"] == "report-1:r2" + @pytest.mark.asyncio async def test_on_success_header_team_id_overrides_auth_user_team( self,