feat(agent): per-turn attribution columns on ai_messages (turn_kind, execution_id, initiator_user_id)

DB unreachable in this worktree (no .env) — migration.sql was hand-written to
match the plan verbatim and validated only via `prisma generate`/`tsc`, not
applied; appendMessage's return type was widened to return the created row
(Promise<AIMessage> in place of Promise<void>) a release early, since Task 13
needs the row id for turnTraceId and the wider return type is a superset any
existing caller can ignore.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018DopjNH21gLaDtWeeMDEHX

Author: GabrielDrapor

frontend/ee/agent/src/__tests__/session-append-message.test.ts

@@ -0,0 +1,58 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const create = vi.fn();
+const findUnique = vi.fn();
+const count = vi.fn();
+vi.mock("@traceroot/core", () => ({
+  prisma: {
+    aIMessage: {
+      create: (...a: unknown[]) => create(...a),
+      count: (...a: unknown[]) => count(...a),
+    },
+    aISession: { findUnique: (...a: unknown[]) => findUnique(...a) },
+  },
+}));
+
+import { SessionManager } from "../session.js";
+
+beforeEach(() => {
+  create.mockReset();
+  findUnique.mockReset();
+  count.mockReset();
+});
+
+describe("appendMessage attribution", () => {
+  it("system session + executionId + first assistant turn → rca_execution", async () => {
+    findUnique.mockResolvedValue({ workspaceId: "w", userId: null, executionId: "exec-1" });
+    count.mockResolvedValue(0);
+    await new SessionManager("s1").appendMessage("assistant", "root cause…");
+    expect(create.mock.calls[0][0].data).toMatchObject({
+      turnKind: "rca_execution",
+      executionId: "exec-1",
+      initiatorUserId: null,
+      kind: "rca",
+    });
+  });
+  it("system session, later turn with a user → rca_followup", async () => {
+    findUnique.mockResolvedValue({ workspaceId: "w", userId: null, executionId: "exec-1" });
+    count.mockResolvedValue(1);
+    await new SessionManager("s1").appendMessage("user", "why?", undefined, undefined, {
+      turnKind: "rca_followup",
+      initiatorUserId: "u9",
+    });
+    expect(create.mock.calls[0][0].data).toMatchObject({
+      turnKind: "rca_followup",
+      initiatorUserId: "u9",
+      kind: "rca",
+    });
+  });
+  it("user session → chat with the session owner as initiator", async () => {
+    findUnique.mockResolvedValue({ workspaceId: "w", userId: "u1", executionId: null });
+    await new SessionManager("s2").appendMessage("user", "hi");
+    expect(create.mock.calls[0][0].data).toMatchObject({
+      turnKind: "chat",
+      initiatorUserId: "u1",
+      kind: "chat",
+    });
+  });
+});

frontend/ee/agent/src/index.ts

@@ -9,6 +9,7 @@ import {
   listSessions,
   deleteSession,
   updateSessionTitle,
+  type TurnAttribution,
 } from "./session.js";
 import { getOrCreateAgent, runAgent, removeAgent, invalidateProviderCache } from "./agent.js";
 import { StreamPersister } from "./stream-persister.js";
@@ -52,13 +53,14 @@ app.post("/api/v1/projects/:projectId/sessions", async (c) => {
   const projectId = c.req.param("projectId");
   const userId = c.req.header("x-user-id") || undefined;
   const workspaceId = c.req.header("x-workspace-id") || "";
-  const body = await c.req.json<{ title?: string }>();
+  const body = await c.req.json<{ title?: string; executionId?: string }>();
 
   const session = await createSession({
     projectId,
     workspaceId,
     userId, // undefined → stored as null for system/RCA sessions
     title: body.title,
+    executionId: body.executionId,
   });
   return c.json(session, 201);
 });
@@ -171,8 +173,26 @@ app.post("/api/v1/projects/:projectId/sessions/:sessionId/messages", async (c) =
 
   console.log(`[Agent] Agent ready, running prompt: "${body.message.slice(0, 50)}"`);
 
+  // Attribution is computed once per turn and applied to every row it
+  // produces (the user message, and every assistant/tool_step row the
+  // persister writes below) so a turn reads as one attributed unit.
+  const attribution: TurnAttribution =
+    ownedSession.userId === null
+      ? userId
+        ? {
+            turnKind: "rca_followup",
+            executionId: ownedSession.executionId,
+            initiatorUserId: userId,
+          }
+        : {
+            turnKind: "rca_execution",
+            executionId: ownedSession.executionId,
+            initiatorUserId: null,
+          }
+      : { turnKind: "chat", initiatorUserId: userId || null };
+
   // Persist user message to DB via SessionManager
-  await sessionManager.appendMessage("user", body.message);
+  await sessionManager.appendMessage("user", body.message, undefined, undefined, attribution);
 
   // Auto-generate session title from first user message (we already have
   // the session loaded above for the auth check — reuse it).
@@ -185,7 +205,7 @@ app.post("/api/v1/projects/:projectId/sessions/:sessionId/messages", async (c) =
     // Mirrors the run into AIMessage rows (text segments, tool steps) so
     // reloaded history matches what the live stream rendered.
     const persister = new StreamPersister((role, content, metadata, tokenUsage) =>
-      sessionManager.appendMessage(role, content, metadata, tokenUsage),
+      sessionManager.appendMessage(role, content, metadata, tokenUsage, attribution),
     );
     // Accumulates token usage across all message_end events (tool-use loops)
     const usageAccumulator = new UsageAccumulator();

frontend/ee/agent/src/session.ts

@@ -20,6 +20,23 @@ export interface TokenUsageData {
   totalTokens?: number;
 }
 
+export type TurnKind = "rca_execution" | "rca_followup" | "chat" | "detector" | "digest";
+
+export interface TurnAttribution {
+  turnKind: TurnKind;
+  executionId?: string | null;
+  initiatorUserId?: string | null;
+}
+
+/** `kind` is kept one release for old readers; derived from turnKind at write time. */
+const LEGACY_KIND: Record<TurnKind, string> = {
+  rca_execution: "rca",
+  rca_followup: "rca",
+  chat: "chat",
+  detector: "detector",
+  digest: "digest-summary",
+};
+
 export class SessionManager {
   constructor(private sessionId: string) {}
 
@@ -87,30 +104,35 @@ export class SessionManager {
    * Like Mom's sessionManager.appendMessage() — persists to DB.
    *
    * `workspaceId` and `kind` are required on every AIMessage row (see schema).
-   * We derive both from the parent AISession: `kind = "chat"` for user sessions
-   * (userId set), `kind = "rca"` for system sessions (userId null). This
-   * mirrors the existing convention in createSession.
+   * `kind` is derived from the turn's attribution (see LEGACY_KIND) and kept
+   * one release for old readers. Returns the created row so callers (e.g. the
+   * turn-trace wrapper) can key off its id.
    */
   async appendMessage(
     role: string,
     content: string,
     metadata?: Record<string, unknown>,
     tokenUsage?: TokenUsageData,
-  ): Promise<void> {
+    attribution?: TurnAttribution,
+  ): Promise<Awaited<ReturnType<typeof prisma.aIMessage.create>>> {
     const session = await prisma.aISession.findUnique({
       where: { id: this.sessionId },
-      select: { workspaceId: true, userId: true },
+      select: { workspaceId: true, userId: true, executionId: true },
     });
     if (!session) {
       throw new Error(`AISession not found: ${this.sessionId}`);
     }
-    const kind = session.userId === null ? "rca" : "chat";
 
-    await prisma.aIMessage.create({
+    const attr = attribution ?? (await this.deriveAttribution(session));
+
+    return prisma.aIMessage.create({
       data: {
         sessionId: this.sessionId,
         workspaceId: session.workspaceId,
-        kind,
+        kind: LEGACY_KIND[attr.turnKind],
+        turnKind: attr.turnKind,
+        executionId: attr.executionId ?? null,
+        initiatorUserId: attr.initiatorUserId ?? null,
         role,
         content,
         metadata: metadata as any,
@@ -125,6 +147,21 @@ export class SessionManager {
       },
     });
   }
+
+  /** Default attribution when the caller did not say: user sessions are chat; a system
+   *  session's turns are the execution until an assistant turn exists, then follow-ups. */
+  private async deriveAttribution(session: {
+    userId: string | null;
+    executionId: string | null;
+  }): Promise<TurnAttribution> {
+    if (session.userId !== null) return { turnKind: "chat", initiatorUserId: session.userId };
+    const priorAssistant = await prisma.aIMessage.count({
+      where: { sessionId: this.sessionId, role: "assistant", content: { not: "" } },
+    });
+    return priorAssistant === 0 && session.executionId
+      ? { turnKind: "rca_execution", executionId: session.executionId, initiatorUserId: null }
+      : { turnKind: "rca_followup", executionId: session.executionId, initiatorUserId: null };
+  }
 }
 
 // ============================================================
@@ -136,13 +173,15 @@ export async function createSession(params: {
   workspaceId: string;
   userId?: string; // optional — null for system/RCA sessions
   title?: string;
+  executionId?: string; // the execution that opened this system session
 }) {
   return prisma.aISession.create({
     data: {
       projectId: params.projectId,
       workspaceId: params.workspaceId,
       userId: params.userId ?? null,
       title: params.title,
+      executionId: params.executionId ?? null,
     },
   });
 }

frontend/ee/agent/src/stream-persister.ts

@@ -1,13 +1,14 @@
 import type { AgentEvent } from "@earendil-works/pi-agent-core";
-import type { TokenUsageData } from "./session.js";
+import type { TokenUsageData, TurnAttribution } from "./session.js";
 
 /** Signature of SessionManager.appendMessage — injected so the persister is testable. */
 export type AppendMessageFn = (
   role: string,
   content: string,
   metadata?: Record<string, unknown>,
   tokenUsage?: TokenUsageData,
-) => Promise<void>;
+  attribution?: TurnAttribution,
+) => Promise<unknown>;
 
 /**
  * Mirrors a run's agent events into durable AIMessage rows so reloaded
@@ -95,7 +96,9 @@ export class StreamPersister {
     tokenUsage?: TokenUsageData,
   ): void {
     this.chain = this.chain
-      .then(() => this.append(role, content, metadata, tokenUsage))
+      .then(async () => {
+        await this.append(role, content, metadata, tokenUsage);
+      })
       .catch((error) => {
         console.error(`[Agent] Failed to persist ${role} message:`, error);
       });

frontend/packages/core/prisma/migrations/20260901000001_ai_message_attribution/migration.sql

@@ -0,0 +1,17 @@
+CREATE TYPE "TurnKind" AS ENUM ('rca_execution', 'rca_followup', 'chat', 'detector', 'digest');
+ALTER TABLE "ai_messages"
+  ADD COLUMN "turn_kind" "TurnKind" NOT NULL DEFAULT 'chat',
+  ADD COLUMN "execution_id" VARCHAR,
+  ADD COLUMN "initiator_user_id" VARCHAR;
+ALTER TABLE "ai_messages" ADD CONSTRAINT "ai_messages_execution_id_fkey"
+  FOREIGN KEY ("execution_id") REFERENCES "detector_rca_executions"("id") ON DELETE SET NULL;
+ALTER TABLE "ai_sessions" ADD COLUMN "execution_id" VARCHAR;
+
+-- Backfill from the legacy kind. Historical system-session turns cannot be split into
+-- execution vs follow-up after the fact; all are attributed to the execution.
+UPDATE "ai_messages" SET "turn_kind" = 'rca_execution' WHERE "kind" = 'rca';
+UPDATE "ai_messages" SET "turn_kind" = 'detector'      WHERE "kind" = 'detector';
+UPDATE "ai_messages" SET "turn_kind" = 'digest'        WHERE "kind" = 'digest-summary';
+UPDATE "ai_messages" m SET "initiator_user_id" = s."user_id"
+  FROM "ai_sessions" s WHERE m."session_id" = s."id" AND s."user_id" IS NOT NULL;
+CREATE INDEX "ix_ai_message_workspace_turnkind_time" ON "ai_messages"("workspace_id", "turn_kind", "create_time");

frontend/packages/core/prisma/schema.prisma

@@ -344,6 +344,9 @@ model AISession {
   title         String?                @db.VarChar
   status        String                 @default("active") @db.VarChar // active, completed, failed
   metadata      Json?                  @db.JsonB
+  // The execution that opened this system session (null for user sessions and
+  // for system sessions predating executions — no backfill by design).
+  executionId   String?                @map("execution_id") @db.VarChar
   createTime    DateTime               @default(now()) @map("create_time") @db.Timestamp(6)
   updateTime    DateTime               @default(now()) @updatedAt @map("update_time") @db.Timestamp(6)
   messages      AIMessage[]
@@ -356,27 +359,42 @@ model AISession {
   @@map("ai_sessions")
 }
 
+enum TurnKind {
+  rca_execution
+  rca_followup
+  chat
+  detector
+  digest
+}
+
 // AI Messages - Individual messages within an AI session
 model AIMessage {
-  id           String     @id @default(cuid()) @db.VarChar
-  workspaceId  String     @map("workspace_id") @db.VarChar // direct workspace pointer for per-kind aggregation (detector scans have no session)
-  sessionId    String?    @map("session_id") @db.VarChar // nullable: detector scans have no chat session
-  kind         String     @default("chat") @db.VarChar // "chat" | "rca" | "detector" | "digest-summary" — categorical tag; usage metering aggregates only chat|rca|detector (digest-summary is recorded, not metered)
-  role         String     @db.VarChar // user, assistant, tool
-  content      String     @db.Text
-  metadata     Json?      @db.JsonB // tool name, tool result, etc.
-  model        String?    @db.VarChar
-  provider     String?    @db.VarChar
-  isByok       Boolean?   @map("is_byok")
-  inputTokens  Int?       @map("input_tokens")
-  outputTokens Int?       @map("output_tokens")
-  cost         Decimal?   @map("cost_usd") @db.Decimal(10, 6)
-  createTime   DateTime   @default(now()) @map("create_time") @db.Timestamp(6)
-  workspace    Workspace  @relation(fields: [workspaceId], references: [id], onDelete: Cascade, onUpdate: NoAction)
-  session      AISession? @relation(fields: [sessionId], references: [id], onDelete: Cascade, onUpdate: NoAction)
+  id          String @id @default(cuid()) @db.VarChar
+  workspaceId String @map("workspace_id") @db.VarChar // direct workspace pointer for per-kind aggregation (detector scans have no session)
+  sessionId   String? @map("session_id") @db.VarChar // nullable: detector scans have no chat session
+  kind        String @default("chat") @db.VarChar // "chat" | "rca" | "detector" | "digest-summary" — categorical tag; usage metering aggregates only chat|rca|detector (digest-summary is recorded, not metered)
+  // Per-turn attribution. `kind` is kept one release for old readers and is
+  // derived from turnKind at write time; metering will move to turnKind (E-09 follow-up).
+  turnKind        TurnKind              @default(chat) @map("turn_kind")
+  executionId     String?               @map("execution_id") @db.VarChar
+  initiatorUserId String?               @map("initiator_user_id") @db.VarChar
+  role            String                @db.VarChar // user, assistant, tool
+  content         String                @db.Text
+  metadata        Json?                 @db.JsonB // tool name, tool result, etc.
+  model           String?               @db.VarChar
+  provider        String?               @db.VarChar
+  isByok          Boolean?              @map("is_byok")
+  inputTokens     Int?                  @map("input_tokens")
+  outputTokens    Int?                  @map("output_tokens")
+  cost            Decimal?              @map("cost_usd") @db.Decimal(10, 6)
+  createTime      DateTime              @default(now()) @map("create_time") @db.Timestamp(6)
+  workspace       Workspace             @relation(fields: [workspaceId], references: [id], onDelete: Cascade, onUpdate: NoAction)
+  session         AISession?            @relation(fields: [sessionId], references: [id], onDelete: Cascade, onUpdate: NoAction)
+  execution       DetectorRcaExecution? @relation(fields: [executionId], references: [id], onDelete: SetNull)
 
   @@index([sessionId], map: "ix_ai_message_session_id")
   @@index([workspaceId, kind, createTime], map: "ix_ai_message_workspace_kind_time")
+  @@index([workspaceId, turnKind, createTime], map: "ix_ai_message_workspace_turnkind_time")
   @@map("ai_messages")
 }
 
@@ -516,6 +534,7 @@ model DetectorRcaExecution {
   rca      DetectorRca  @relation("rcaExecutions", fields: [findingId], references: [findingId], onDelete: Cascade)
   session  AISession?   @relation(fields: [sessionId], references: [id], onDelete: SetNull)
   latestOf DetectorRca? @relation("latestExecution")
+  messages AIMessage[]
 
   @@unique([findingId, attempt], map: "uq_rca_execution_finding_attempt")
   @@index([projectId], map: "ix_rca_execution_project_id")

frontend/worker/src/processors/detector-digest-processor.ts

@@ -111,6 +111,7 @@ export async function flushDigest(job: DigestFlushJob): Promise<void> {
             workspaceId: recipients.workspaceId,
             sessionId: null,
             kind: "digest-summary",
+            turnKind: "digest",
             role: "assistant",
             content: "",
             model: result.usage.model,

frontend/worker/src/processors/detector-run-processor.ts

@@ -429,6 +429,7 @@ async function evaluateTrace(
         workspaceId,
         sessionId: null,
         kind: "detector",
+        turnKind: "detector" as const,
         role: "assistant",
         content: "", // detector scans don't have a chat-like content payload
         model: u.inferenceModel,