From 8c99efa09de3f892823ac14d774e61d3dd6137b6 Mon Sep 17 00:00:00 2001
From: Alex Lebedev <woutut@gmail.com>
Date: Thu, 30 Jul 2026 17:10:54 +0200
Subject: [PATCH] feat(agent): forward ai_session_id run-state key to gateway
 properties
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Products that stash ai_session_id in run state (ReviewHog stamps one per review turn) get it lifted onto their runs' $ai_generation events, grouping a logical operation's generations. Transported unreserved — $-keys are stripped at gateway header boundaries — and promoted to the native $ai_session_id by the llm-gateway at capture. Null-safe: runs without the state key emit no header.

Generated-By: PostHog Code
Task-Id: 77db6432-21d6-400a-ad79-7ce81def58cf
---
 .../src/server/agent-server.configure-environment.test.ts | 4 ++++
 packages/agent/src/server/agent-server.ts                 | 8 ++++++++
 2 files changed, 12 insertions(+)

diff --git a/packages/agent/src/server/agent-server.configure-environment.test.ts b/packages/agent/src/server/agent-server.configure-environment.test.ts
index 0ae96970f6..fc75920c5c 100644
--- a/packages/agent/src/server/agent-server.configure-environment.test.ts
+++ b/packages/agent/src/server/agent-server.configure-environment.test.ts
@@ -207,6 +207,7 @@ describe("AgentServer.configureEnvironment", () => {
       originProduct: "signal_report",
       signalReportId: "report-123",
       aiStage: "research",
+      aiSessionId: "report-1:r2",
       taskId: "task-abc",
       taskRunId: "run-xyz",
       taskUserId: 42,
@@ -218,6 +219,7 @@ describe("AgentServer.configureEnvironment", () => {
       "x-posthog-property-task_internal": "true",
       "x-posthog-property-signal_report_id": "report-123",
       "x-posthog-property-ai_stage": "research",
+      "x-posthog-property-ai_session_id": "report-1:r2",
       "x-posthog-property-task_id": "task-abc",
       "x-posthog-property-task_run_id": "run-xyz",
       "x-posthog-property-task_user_id": "42",
@@ -232,6 +234,7 @@ describe("AgentServer.configureEnvironment", () => {
       originProduct: "signal_report",
       signalReportId: "report-123",
       aiStage: "research",
+      aiSessionId: "report-1:r2",
       taskId: "task-abc",
       taskRunId: "run-xyz",
       taskUserId: 42,
@@ -244,6 +247,7 @@ describe("AgentServer.configureEnvironment", () => {
         "x-posthog-property-task_internal: true",
         "x-posthog-property-signal_report_id: report-123",
         "x-posthog-property-ai_stage: research",
+        "x-posthog-property-ai_session_id: report-1:r2",
         "x-posthog-property-task_id: task-abc",
         "x-posthog-property-task_run_id: run-xyz",
         "x-posthog-property-task_user_id: 42",
diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts
index f53594a274..88b10fb218 100644
--- a/packages/agent/src/server/agent-server.ts
+++ b/packages/agent/src/server/agent-server.ts
@@ -1584,6 +1584,7 @@ export class AgentServer {
       originProduct: preTask?.origin_product,
       signalReportId: preTask?.signal_report,
       aiStage: getTaskRunStateString(preTaskRun, "ai_stage"),
+      aiSessionId: getTaskRunStateString(preTaskRun, "ai_session_id"),
       taskId: payload.task_id,
       taskRunId: payload.run_id,
       taskUserId: payload.user_id || preTask?.created_by?.id || null,
@@ -3912,6 +3913,7 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions}
     originProduct,
     signalReportId,
     aiStage,
+    aiSessionId,
     taskId,
     taskRunId,
     taskUserId,
@@ -3921,6 +3923,7 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions}
     originProduct?: Task["origin_product"] | null;
     signalReportId?: string | null;
     aiStage?: string | null;
+    aiSessionId?: string | null;
     taskId?: string | null;
     taskRunId?: string | null;
     taskUserId?: number | null;
@@ -3946,6 +3949,11 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions}
       task_internal: isInternal,
       signal_report_id: signalReportId,
       ai_stage: aiStage,
+      // Session key for LLM analytics, grouping the run's generations under the caller's
+      // logical operation (e.g. one ReviewHog review turn). Transported unreserved — $-keys
+      // are stripped at gateway header boundaries — and promoted to `$ai_session_id` by the
+      // llm-gateway when it captures the event.
+      ai_session_id: aiSessionId,
       task_id: taskId,
       task_run_id: taskRunId,
       task_user_id: taskUserId,