Index agent message consumption analytics

front/lib/analytics/agent_message_consumption/index.ts

@@ -0,0 +1,33 @@
+import { buildAgentMessageConsumptionAnalyticsDocuments } from "@app/lib/analytics/agent_message_consumption/documents";
+import { loadAgentMessageConsumptionAnalyticsInput } from "@app/lib/analytics/agent_message_consumption/load";
+import { upsertAgentMessageConsumptionAnalyticsDocuments } from "@app/lib/analytics/agent_message_consumption/store";
+import type { ElasticsearchError } from "@app/lib/api/elasticsearch";
+import type { Authenticator } from "@app/lib/auth";
+import type { Result } from "@app/types/shared/result";
+import { Ok } from "@app/types/shared/result";
+import assert from "assert";
+
+/**
+ * Loads, projects, and indexes the complete consumption analytics snapshot for one agent message.
+ * Callers only identify the message. This module owns the ordering and completeness requirements
+ * of the indexed snapshot.
+ */
+export async function indexAgentMessageConsumptionAnalytics(
+  auth: Authenticator,
+  { agentMessageId }: { agentMessageId: string }
+): Promise<Result<void, ElasticsearchError>> {
+  const input = await loadAgentMessageConsumptionAnalyticsInput(auth, {
+    agentMessageId,
+  });
+  if (!input) {
+    return new Ok(undefined);
+  }
+
+  const documents = buildAgentMessageConsumptionAnalyticsDocuments(input);
+  assert(
+    documents && documents.length > 0,
+    "Consumption attribution is incomplete for analytics"
+  );
+
+  return upsertAgentMessageConsumptionAnalyticsDocuments(documents);
+}

front/lib/analytics/agent_message_consumption/store.test.ts

@@ -0,0 +1,164 @@
+import { upsertAgentMessageConsumptionAnalyticsDocuments } from "@app/lib/analytics/agent_message_consumption/store";
+import {
+  CONSUMPTION_ANALYTICS_ALIAS_NAME,
+  ElasticsearchError,
+  withEs,
+} from "@app/lib/api/elasticsearch";
+import { USAGE_TYPE_USER } from "@app/lib/metronome/constants";
+import type { AgentMessageConsumptionAnalyticsData } from "@app/types/assistant/analytics";
+import { Err, Ok } from "@app/types/shared/result";
+import { normalizeError } from "@app/types/shared/utils/error_utils";
+import { Client } from "@elastic/elasticsearch";
+import { afterAll, beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("@app/lib/api/elasticsearch", async (importActual) => {
+  const actual =
+    await importActual<typeof import("@app/lib/api/elasticsearch")>();
+  return { ...actual, withEs: vi.fn() };
+});
+
+const client = new Client({ node: "http://localhost:9200" });
+const bulkMock = vi.spyOn(client, "bulk");
+
+function makeDocument(): AgentMessageConsumptionAnalyticsData {
+  return {
+    agent: {
+      id: "agent_1",
+      version: "1",
+      tag_ids: [],
+      parent_ids: [],
+      direct_parent_id: null,
+      root_id: "agent_1",
+      depth: 0,
+    },
+    agent_message_id: "agent_message_1",
+    api_key_name: null,
+    attribution_version: 4,
+    completed_at: "2026-08-07T12:00:00.000Z",
+    consumption_key: "run-usage:1",
+    consumption_type: "llm",
+    context_origin: "web",
+    conversation_id: "conversation_1",
+    credit_micro: 1_000_000,
+    execution_time_ms: null,
+    gross_credit_micro: {
+      system: 0,
+      input: 600_000,
+      result_footprint: null,
+      output: 400_000,
+      reasoning: 0,
+      direct: 0,
+      total: 1_000_000,
+    },
+    message_version: "2",
+    model: null,
+    run_usage_id: "1",
+    space_id: null,
+    status: "succeeded",
+    step_index: 0,
+    tokens: {
+      system: 0,
+      input: 10,
+      result_footprint: null,
+      output: 5,
+      reasoning: 0,
+    },
+    tool: null,
+    trigger_id: null,
+    usage_type: USAGE_TYPE_USER,
+    user: null,
+    workspace_id: "workspace_1",
+  };
+}
+
+describe("upsertAgentMessageConsumptionAnalyticsDocuments", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    bulkMock.mockResolvedValue({ errors: false, items: [], took: 1 });
+    vi.mocked(withEs).mockImplementation(async (fn) => {
+      try {
+        return new Ok(await fn(client));
+      } catch (error) {
+        return new Err(
+          new ElasticsearchError("query_error", normalizeError(error).message)
+        );
+      }
+    });
+  });
+
+  afterAll(async () => {
+    await client.close();
+  });
+
+  it("uses a stable identity for idempotent upserts", async () => {
+    const document = makeDocument();
+
+    const result = await upsertAgentMessageConsumptionAnalyticsDocuments([
+      document,
+    ]);
+
+    expect(result.isOk()).toBe(true);
+    expect(bulkMock).toHaveBeenCalledWith({
+      body: [
+        {
+          index: {
+            _index: CONSUMPTION_ANALYTICS_ALIAS_NAME,
+            _id: "workspace_1_agent_message_1_run-usage:1",
+          },
+        },
+        document,
+      ],
+      refresh: false,
+    });
+  });
+
+  it("does nothing when there are no documents", async () => {
+    const result = await upsertAgentMessageConsumptionAnalyticsDocuments([]);
+
+    expect(result.isOk()).toBe(true);
+    expect(withEs).not.toHaveBeenCalled();
+  });
+
+  it("returns the Elasticsearch error when the request fails", async () => {
+    const error = new ElasticsearchError("connection_error", "write failed");
+    vi.mocked(withEs).mockResolvedValueOnce(new Err(error));
+
+    const result = await upsertAgentMessageConsumptionAnalyticsDocuments([
+      makeDocument(),
+    ]);
+
+    expect(result.isErr()).toBe(true);
+    if (result.isErr()) {
+      expect(result.error).toBe(error);
+    }
+  });
+
+  it("returns the error from a failed bulk item", async () => {
+    bulkMock.mockResolvedValueOnce({
+      errors: true,
+      items: [
+        {
+          index: {
+            _index: CONSUMPTION_ANALYTICS_ALIAS_NAME,
+            status: 429,
+            error: {
+              type: "es_rejected_execution_exception",
+              reason: "queue full",
+            },
+          },
+        },
+      ],
+      took: 1,
+    });
+
+    const result = await upsertAgentMessageConsumptionAnalyticsDocuments([
+      makeDocument(),
+    ]);
+
+    expect(result.isErr()).toBe(true);
+    if (result.isErr()) {
+      expect(result.error.message).toBe("queue full");
+      expect(result.error.statusCode).toBe(429);
+    }
+  });
+});

front/lib/analytics/agent_message_consumption/store.ts

@@ -0,0 +1,62 @@
+import {
+  CONSUMPTION_ANALYTICS_ALIAS_NAME,
+  ElasticsearchError,
+  withEs,
+} from "@app/lib/api/elasticsearch";
+import type { AgentMessageConsumptionAnalyticsData } from "@app/types/assistant/analytics";
+import type { Result } from "@app/types/shared/result";
+import { Err, Ok } from "@app/types/shared/result";
+
+function makeAgentMessageConsumptionAnalyticsDocumentId(
+  document: Pick<
+    AgentMessageConsumptionAnalyticsData,
+    "agent_message_id" | "consumption_key" | "workspace_id"
+  >
+): string {
+  return `${document.workspace_id}_${document.agent_message_id}_${document.consumption_key}`;
+}
+
+/** Upserts every consumption unit using its stable identity. */
+export async function upsertAgentMessageConsumptionAnalyticsDocuments(
+  documents: AgentMessageConsumptionAnalyticsData[]
+): Promise<Result<void, ElasticsearchError>> {
+  if (documents.length === 0) {
+    return new Ok(undefined);
+  }
+
+  const result = await withEs((client) =>
+    client.bulk({
+      body: documents.flatMap((document) => [
+        {
+          index: {
+            _index: CONSUMPTION_ANALYTICS_ALIAS_NAME,
+            _id: makeAgentMessageConsumptionAnalyticsDocumentId(document),
+          },
+        },
+        document,
+      ]),
+      refresh: false,
+    })
+  );
+
+  if (result.isErr()) {
+    return result;
+  }
+
+  if (!result.value.errors) {
+    return new Ok(undefined);
+  }
+
+  const failedItem = result.value.items.find(
+    (item) => item.index?.error
+  )?.index;
+
+  return new Err(
+    new ElasticsearchError(
+      "query_error",
+      failedItem?.error?.reason ??
+        "Elasticsearch bulk response contains failed items",
+      failedItem?.status
+    )
+  );
+}

front/temporal/agent_loop/activities/finalize.ts

@@ -29,19 +29,37 @@ import {
 } from "@app/temporal/agent_loop/activities/usage_tracking";
 import type { AgentLoopArgs } from "@app/types/assistant/agent_run";
 
+async function launchAgentMessageConsumptionAttributionAfterPersistingInputs(
+  auth: Authenticator,
+  agentLoopArgs: AgentLoopArgs,
+  {
+    creditArgs = agentLoopArgs,
+  }: {
+    creditArgs?: { agentMessageId: string; dustRunIds?: string[] };
+  } = {}
+): Promise<void> {
+  // Consumption analytics needs the authoritative bill, usage type, and historical skill snapshot
+  // before its attribution workflow can safely materialize Elasticsearch documents.
+  await snapshotAgentMessageSkills(auth, agentLoopArgs);
+  await computeAndStoreAgentMessageCredits(auth, creditArgs);
+
+  await launchAgentMessageConsumptionAttribution(auth, agentLoopArgs);
+}
+
 export async function finalizeSuccessfulAgentLoopActivity(
   authType: AuthenticatorType,
   agentLoopArgs: AgentLoopArgs
 ): Promise<void> {
   const auth = await Authenticator.fromJsonWithRefrehedGroups(authType);
 
   await Promise.all([
-    snapshotAgentMessageSkills(auth, agentLoopArgs),
     launchAgentMessageAnalytics(auth, agentLoopArgs),
-    launchAgentMessageConsumptionAttribution(auth, agentLoopArgs),
+    launchAgentMessageConsumptionAttributionAfterPersistingInputs(
+      auth,
+      agentLoopArgs
+    ),
     launchTrackProgrammaticUsage(auth, agentLoopArgs),
     launchEmitMetronomeUsageEvents(auth, agentLoopArgs),
-    computeAndStoreAgentMessageCredits(auth, agentLoopArgs),
     conversationUnreadNotification(auth, agentLoopArgs),
     activationNewConversationNotification(auth, agentLoopArgs),
     handleMentions(auth, agentLoopArgs),
@@ -64,12 +82,13 @@ export async function finalizeGracefullyStoppedAgentLoopActivity(
   const auth = await Authenticator.fromJsonWithRefrehedGroups(authType);
 
   await Promise.all([
-    snapshotAgentMessageSkills(auth, agentLoopArgs),
     launchAgentMessageAnalytics(auth, agentLoopArgs),
-    launchAgentMessageConsumptionAttribution(auth, agentLoopArgs),
+    launchAgentMessageConsumptionAttributionAfterPersistingInputs(
+      auth,
+      agentLoopArgs
+    ),
     launchTrackProgrammaticUsage(auth, agentLoopArgs),
     launchEmitMetronomeUsageEvents(auth, agentLoopArgs),
-    computeAndStoreAgentMessageCredits(auth, agentLoopArgs),
     conversationUnreadNotification(auth, agentLoopArgs),
     handleMentions(auth, agentLoopArgs),
   ]);
@@ -92,12 +111,13 @@ export async function finalizeInterruptedAgentLoopActivity(
   const auth = await Authenticator.fromJsonWithRefrehedGroups(authType);
 
   await Promise.all([
-    snapshotAgentMessageSkills(auth, agentLoopArgs),
     launchAgentMessageAnalytics(auth, agentLoopArgs),
-    launchAgentMessageConsumptionAttribution(auth, agentLoopArgs),
+    launchAgentMessageConsumptionAttributionAfterPersistingInputs(
+      auth,
+      agentLoopArgs
+    ),
     launchTrackProgrammaticUsage(auth, agentLoopArgs),
     launchEmitMetronomeUsageEvents(auth, agentLoopArgs),
-    computeAndStoreAgentMessageCredits(auth, agentLoopArgs),
     conversationUnreadNotification(auth, agentLoopArgs),
     handleMentions(auth, agentLoopArgs),
   ]);
@@ -112,12 +132,13 @@ export async function finalizeCancelledAgentLoopActivity(
   const auth = await Authenticator.fromJsonWithRefrehedGroups(authType);
 
   await Promise.all([
-    snapshotAgentMessageSkills(auth, agentLoopArgs),
     launchAgentMessageAnalytics(auth, agentLoopArgs),
-    launchAgentMessageConsumptionAttribution(auth, agentLoopArgs),
+    launchAgentMessageConsumptionAttributionAfterPersistingInputs(
+      auth,
+      agentLoopArgs
+    ),
     launchTrackProgrammaticUsage(auth, agentLoopArgs),
     launchEmitMetronomeUsageEvents(auth, agentLoopArgs),
-    computeAndStoreAgentMessageCredits(auth, agentLoopArgs),
     sendEmailReplyOnError(
       auth,
       agentLoopArgs,
@@ -135,14 +156,16 @@ export async function finalizeCreditStoppedAgentLoopActivity(
   const auth = await Authenticator.fromJsonWithRefrehedGroups(authType);
 
   await Promise.all([
-    snapshotAgentMessageSkills(auth, agentLoopArgs),
     launchAgentMessageAnalytics(auth, agentLoopArgs),
-    launchAgentMessageConsumptionAttribution(auth, agentLoopArgs),
+    launchAgentMessageConsumptionAttributionAfterPersistingInputs(
+      auth,
+      agentLoopArgs,
+      {
+        creditArgs: { agentMessageId: agentLoopArgs.agentMessageId },
+      }
+    ),
     launchTrackProgrammaticUsage(auth, agentLoopArgs),
     launchEmitMetronomeUsageEvents(auth, agentLoopArgs),
-    computeAndStoreAgentMessageCredits(auth, {
-      agentMessageId: agentLoopArgs.agentMessageId,
-    }),
     sendEmailReplyOnError(auth, agentLoopArgs, creditsExhaustedMessage(auth)),
   ]);
 }
@@ -157,12 +180,13 @@ export async function finalizeErroredAgentLoopActivity(
   const auth = await Authenticator.fromJsonWithRefrehedGroups(authType);
 
   await Promise.all([
-    snapshotAgentMessageSkills(auth, agentLoopArgs),
     launchAgentMessageAnalytics(auth, agentLoopArgs),
-    launchAgentMessageConsumptionAttribution(auth, agentLoopArgs),
+    launchAgentMessageConsumptionAttributionAfterPersistingInputs(
+      auth,
+      agentLoopArgs
+    ),
     launchTrackProgrammaticUsage(auth, agentLoopArgs),
     launchEmitMetronomeUsageEvents(auth, agentLoopArgs),
-    computeAndStoreAgentMessageCredits(auth, agentLoopArgs),
     sendEmailReplyOnError(
       auth,
       agentLoopArgs,