From 00088af22d574a52eada9f41b53c32fc9abbc735 Mon Sep 17 00:00:00 2001
From: Arthur Knaus <arthur.knaus@sentry.io>
Date: Mon, 2 Jun 2025 10:23:09 +0200
Subject: [PATCH] feat(agent-insights): Handle new keys

---
 .../utils/highlightedSpanAttributes.tsx       | 48 ++++++++++++++++---
 .../insights/agentMonitoring/utils/query.tsx  |  8 ++++
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/static/app/views/insights/agentMonitoring/utils/highlightedSpanAttributes.tsx b/static/app/views/insights/agentMonitoring/utils/highlightedSpanAttributes.tsx
index 309e6d874bcec3..61890b0c36cea0 100644
--- a/static/app/views/insights/agentMonitoring/utils/highlightedSpanAttributes.tsx
+++ b/static/app/views/insights/agentMonitoring/utils/highlightedSpanAttributes.tsx
@@ -2,10 +2,14 @@ import {IconArrow} from 'sentry/icons';
 import {t} from 'sentry/locale';
 import type {EventTransaction} from 'sentry/types/event';
 import type {Organization} from 'sentry/types/organization';
+import {formatAbbreviatedNumberWithDynamicPrecision} from 'sentry/utils/formatters';
 import {prettifyAttributeName} from 'sentry/views/explore/components/traceItemAttributes/utils';
 import type {TraceItemResponseAttribute} from 'sentry/views/explore/hooks/useTraceItemDetails';
 import {hasAgentInsightsFeature} from 'sentry/views/insights/agentMonitoring/utils/features';
-import {getIsAiSpan} from 'sentry/views/insights/agentMonitoring/utils/query';
+import {
+  getIsAiSpan,
+  legacyAttributeKeys,
+} from 'sentry/views/insights/agentMonitoring/utils/query';
 import {
   isEAPSpanNode,
   isSpanNode,
@@ -32,6 +36,29 @@ function ensureAttributeObject(
   return attributes;
 }
 
+/**
+ * Get an attribute from the attribute object, checking both the current and legacy keys.
+ * @param attributeObject - The attribute object.
+ * @param key - The key to check.
+ * @returns The attribute value, or undefined if the attribute is not found.
+ */
+function getAttribute(attributeObject: Record<string, string>, key: string) {
+  if (attributeObject[key]) {
+    return attributeObject[key];
+  }
+  const legacyKeys = legacyAttributeKeys.get(key) ?? [];
+  for (const legacyKey of legacyKeys) {
+    if (attributeObject[legacyKey]) {
+      return attributeObject[legacyKey];
+    }
+  }
+  return undefined;
+}
+
+function formatCost(cost: string) {
+  return `US $${formatAbbreviatedNumberWithDynamicPrecision(cost)}`;
+}
+
 export function getHighlightedSpanAttributes({
   op,
   description,
@@ -51,16 +78,17 @@ export function getHighlightedSpanAttributes({
   const attributeObject = ensureAttributeObject(attributes);
   const highlightedAttributes = [];
 
-  if (attributeObject['ai.model.id']) {
+  const model = getAttribute(attributeObject, 'gen_ai.request.model');
+  if (model) {
     highlightedAttributes.push({
       name: t('Model'),
-      value: attributeObject['ai.model.id'],
+      value: model,
     });
   }
 
-  const promptTokens = attributeObject['ai.prompt_tokens.used'];
-  const completionTokens = attributeObject['ai.completion_tokens.used'];
-  const totalTokens = attributeObject['ai.total_tokens.used'];
+  const promptTokens = getAttribute(attributeObject, 'gen_ai.usage.input_tokens');
+  const completionTokens = getAttribute(attributeObject, 'gen_ai.usage.output_tokens');
+  const totalTokens = getAttribute(attributeObject, 'gen_ai.usage.total_tokens');
   if (promptTokens && completionTokens && totalTokens) {
     highlightedAttributes.push({
       name: t('Tokens'),
@@ -73,6 +101,14 @@ export function getHighlightedSpanAttributes({
     });
   }
 
+  const totalCosts = getAttribute(attributeObject, 'gen_ai.usage.total_cost');
+  if (totalCosts) {
+    highlightedAttributes.push({
+      name: t('Cost'),
+      value: formatCost(totalCosts),
+    });
+  }
+
   if (attributeObject['ai.toolCall.name']) {
     highlightedAttributes.push({
       name: t('Tool Name'),
diff --git a/static/app/views/insights/agentMonitoring/utils/query.tsx b/static/app/views/insights/agentMonitoring/utils/query.tsx
index c3d6113d04fc0c..d916a6a459d904 100644
--- a/static/app/views/insights/agentMonitoring/utils/query.tsx
+++ b/static/app/views/insights/agentMonitoring/utils/query.tsx
@@ -14,6 +14,14 @@ const AI_TOKEN_USAGE_ATTRIBUTE = 'ai.total_tokens.used';
 
 export const AI_TOKEN_USAGE_ATTRIBUTE_SUM = `sum(${AI_TOKEN_USAGE_ATTRIBUTE})`;
 
+export const legacyAttributeKeys = new Map<string, string[]>([
+  ['gen_ai.request.model', ['ai.model.id']],
+  ['gen_ai.usage.input_tokens', ['ai.prompt_tokens.used']],
+  ['gen_ai.usage.output_tokens', ['ai.completion_tokens.used']],
+  ['gen_ai.usage.total_tokens', ['ai.total_tokens.used']],
+  ['gen_ai.usage.total_cost', ['ai.total_cost.used']],
+]);
+
 // TODO: Remove once tool spans have their own op
 function mapMissingSpanOp({
   op = 'default',