aws-samples · varunmaggo · Jun 25, 2026
diff --git a/backend/src/lambda/__tests__/app-metrics-handler.test.ts b/backend/src/lambda/__tests__/app-metrics-handler.test.ts
@@ -189,6 +189,66 @@ describe('computeMetrics', () => {
     expect(metrics.p50Latency).toBe(0);
     expect(metrics.p95Latency).toBe(0);
     expect(metrics.p99Latency).toBe(0);
+    expect(metrics.totalInputTokens).toBe(0);
+    expect(metrics.totalOutputTokens).toBe(0);
+    expect(metrics.totalCostMicroUsd).toBe(0);
+  });
+
+  test('aggregates token usage and attributed cost across entries', () => {
+    const entries: AccessLogEntry[] = [
+      makeLogEntry({
+        requestId: 'req-a',
+        status: 200,
+        model: 'anthropic.claude-3-5-sonnet',
+        inputTokens: 1000,
+        outputTokens: 500,
+      }),
+      makeLogEntry({
+        requestId: 'req-b',
+        status: 200,
+        model: 'anthropic.claude-3-5-sonnet',
+        inputTokens: 2000,
+        outputTokens: 0,
+      }),
+    ];
+
+    const metrics = computeMetrics(entries);
+
+    expect(metrics.totalInputTokens).toBe(3000);
+    expect(metrics.totalOutputTokens).toBe(500);
+    // sonnet: $3/1M in, $15/1M out
+    //   req-a: 1000*3 + 500*15 = 10500
+    //   req-b: 2000*3 + 0       =  6000
+    expect(metrics.totalCostMicroUsd).toBe(16500);
+  });
+
+  test('treats entries without token fields as zero tokens and zero cost', () => {
+    const entries: AccessLogEntry[] = [makeLogEntry({ status: 200 })];
+
+    const metrics = computeMetrics(entries);
+
+    expect(metrics.totalInputTokens).toBe(0);
+    expect(metrics.totalOutputTokens).toBe(0);
+    expect(metrics.totalCostMicroUsd).toBe(0);
+  });
+
+  test('parses token fields from structured log entries', () => {
+    const messages = [
+      JSON.stringify(
+        makeLogEntry({
+          requestId: 'req-1',
+          model: 'amazon.nova-lite',
+          inputTokens: 10,
+          outputTokens: 20,
+        }),
+      ),
+    ];
+
+    const parsed = parseAccessLogEntries(messages);
+
+    expect(parsed[0].model).toBe('amazon.nova-lite');
+    expect(parsed[0].inputTokens).toBe(10);
+    expect(parsed[0].outputTokens).toBe(20);
   });
 });
 

diff --git a/backend/src/lambda/__tests__/cost-model.test.ts b/backend/src/lambda/__tests__/cost-model.test.ts
@@ -0,0 +1,77 @@
+/**
+ * Unit tests for the cost model used by per-app token/cost attribution.
+ */
+import {
+  MODEL_RATES,
+  DEFAULT_RATE,
+  resolveRate,
+  computeCost,
+  microUsdToUsd,
+} from '../cost-model';
+
+describe('resolveRate', () => {
+  test('returns the exact rate for a known model', () => {
+    const { rate, usedFallbackRate } = resolveRate('anthropic.claude-3-5-sonnet');
+    expect(rate).toEqual(MODEL_RATES['anthropic.claude-3-5-sonnet']);
+    expect(usedFallbackRate).toBe(false);
+  });
+
+  test('falls back for an unknown model', () => {
+    const { rate, usedFallbackRate } = resolveRate('some.unlisted-model');
+    expect(rate).toEqual(DEFAULT_RATE);
+    expect(usedFallbackRate).toBe(true);
+  });
+
+  test('falls back for undefined model', () => {
+    const { rate, usedFallbackRate } = resolveRate(undefined);
+    expect(rate).toEqual(DEFAULT_RATE);
+    expect(usedFallbackRate).toBe(true);
+  });
+
+  test('does not treat inherited Object properties as known models', () => {
+    const { usedFallbackRate } = resolveRate('toString');
+    expect(usedFallbackRate).toBe(true);
+  });
+});
+
+describe('computeCost', () => {
+  test('computes cost in micro-USD for a known model', () => {
+    // sonnet: $3/1M input, $15/1M output. 1000 in + 500 out:
+    //   1000*3 + 500*15 = 3000 + 7500 = 10500 micro-USD
+    const { costMicroUsd, usedFallbackRate } = computeCost(
+      'anthropic.claude-3-5-sonnet',
+      1000,
+      500,
+    );
+    expect(costMicroUsd).toBe(10500);
+    expect(usedFallbackRate).toBe(false);
+    expect(microUsdToUsd(costMicroUsd)).toBeCloseTo(0.0105, 6);
+  });
+
+  test('uses the fallback rate for an unknown model and flags it', () => {
+    const { costMicroUsd, usedFallbackRate } = computeCost('mystery.model', 1_000_000, 0);
+    // default input rate $3/1M => 1,000,000 * 3 = 3,000,000 micro-USD = $3
+    expect(costMicroUsd).toBe(3_000_000);
+    expect(usedFallbackRate).toBe(true);
+  });
+
+  test('returns zero cost for zero tokens', () => {
+    expect(computeCost('anthropic.claude-3-opus', 0, 0).costMicroUsd).toBe(0);
+  });
+
+  test('clamps negative and non-finite token counts to zero', () => {
+    expect(computeCost('anthropic.claude-3-5-sonnet', -100, -5).costMicroUsd).toBe(0);
+    expect(computeCost('anthropic.claude-3-5-sonnet', NaN, Infinity).costMicroUsd).toBe(0);
+  });
+
+  test('floors fractional token counts before pricing', () => {
+    // haiku: $0.25/1M input. 4 input tokens (from 4.9) => round(4*0.25)=1 micro-USD
+    const { costMicroUsd } = computeCost('anthropic.claude-3-haiku', 4.9, 0);
+    expect(costMicroUsd).toBe(1);
+  });
+
+  test('result is always an integer (safe for DynamoDB ADD counters)', () => {
+    const { costMicroUsd } = computeCost('anthropic.claude-3-haiku', 7, 3);
+    expect(Number.isInteger(costMicroUsd)).toBe(true);
+  });
+});
diff --git a/backend/src/lambda/app-metrics-handler.ts b/backend/src/lambda/app-metrics-handler.ts
@@ -11,6 +11,7 @@
 import { DynamoDBDocumentClient, UpdateCommand, QueryCommand, ScanCommand } from '@aws-sdk/lib-dynamodb';
 import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
 import { gunzipSync } from 'zlib';
+import { computeCost, microUsdToUsd } from './cost-model';
 
 // ── Types ───────────────────────────────────────────────────
 
@@ -21,6 +22,12 @@ export interface AccessLogEntry {
   status: number;
   latency: number;
   timestamp: string;
+  /** LLM model id used to serve the request (optional; older logs omit it). */
+  model?: string;
+  /** Input/prompt tokens consumed (optional; defaults to 0 when absent). */
+  inputTokens?: number;
+  /** Output/completion tokens produced (optional; defaults to 0 when absent). */
+  outputTokens?: number;
 }
 
 export interface AggregatedMetrics {
@@ -31,6 +38,12 @@ export interface AggregatedMetrics {
   p50Latency: number;
   p95Latency: number;
   p99Latency: number;
+  /** Total input/prompt tokens across the entries. */
+  totalInputTokens: number;
+  /** Total output/completion tokens across the entries. */
+  totalOutputTokens: number;
+  /** Attributed cost in integer micro-USD (millionths of a dollar). */
+  totalCostMicroUsd: number;
 }
 
 export interface MetricsDeps {
@@ -63,6 +76,9 @@ export function parseAccessLogEntries(messages: string[]): AccessLogEntry[] {
           status: parsed.status,
           latency: parsed.latency,
           timestamp: parsed.timestamp,
+          model: typeof parsed.model === 'string' ? parsed.model : undefined,
+          inputTokens: typeof parsed.inputTokens === 'number' ? parsed.inputTokens : undefined,
+          outputTokens: typeof parsed.outputTokens === 'number' ? parsed.outputTokens : undefined,
         });
       }
     } catch {
@@ -113,12 +129,18 @@ export function computeMetrics(entries: AccessLogEntry[]): AggregatedMetrics {
       p50Latency: 0,
       p95Latency: 0,
       p99Latency: 0,
+      totalInputTokens: 0,
+      totalOutputTokens: 0,
+      totalCostMicroUsd: 0,
     };
   }
 
   let successCount = 0;
   let clientErrorCount = 0;
   let serverErrorCount = 0;
+  let totalInputTokens = 0;
+  let totalOutputTokens = 0;
+  let totalCostMicroUsd = 0;
   const latencies: number[] = [];
 
   for (const entry of entries) {
@@ -127,6 +149,12 @@ export function computeMetrics(entries: AccessLogEntry[]): AggregatedMetrics {
     else if (statusClass === 4) clientErrorCount++;
     else if (statusClass === 5) serverErrorCount++;
     latencies.push(entry.latency);
+
+    const inTok = entry.inputTokens ?? 0;
+    const outTok = entry.outputTokens ?? 0;
+    totalInputTokens += inTok > 0 ? inTok : 0;
+    totalOutputTokens += outTok > 0 ? outTok : 0;
+    totalCostMicroUsd += computeCost(entry.model, inTok, outTok).costMicroUsd;
   }
 
   latencies.sort((a, b) => a - b);
@@ -139,6 +167,9 @@ export function computeMetrics(entries: AccessLogEntry[]): AggregatedMetrics {
     p50Latency: percentile(latencies, 50),
     p95Latency: percentile(latencies, 95),
     p99Latency: percentile(latencies, 99),
+    totalInputTokens,
+    totalOutputTokens,
+    totalCostMicroUsd,
   };
 }
 
@@ -189,7 +220,10 @@ async function upsertMetrics(
       ADD totalRequests :totalRequests,
           successCount :successCount,
           clientErrorCount :clientErrorCount,
-          serverErrorCount :serverErrorCount
+          serverErrorCount :serverErrorCount,
+          totalInputTokens :totalInputTokens,
+          totalOutputTokens :totalOutputTokens,
+          totalCostMicroUsd :totalCostMicroUsd
       SET groupId = :groupId,
           sortId = :sortId,
           p50Latency = :p50Latency,
@@ -203,6 +237,9 @@ async function upsertMetrics(
       ':successCount': metrics.successCount,
       ':clientErrorCount': metrics.clientErrorCount,
       ':serverErrorCount': metrics.serverErrorCount,
+      ':totalInputTokens': metrics.totalInputTokens,
+      ':totalOutputTokens': metrics.totalOutputTokens,
+      ':totalCostMicroUsd': metrics.totalCostMicroUsd,
       ':groupId': `APP#${appId}`,
       ':sortId': `METRICS#${hourBucket}`,
       ':p50Latency': metrics.p50Latency,
@@ -253,6 +290,12 @@ export interface AppMetricsResult {
   p50Latency: number;
   p95Latency: number;
   p99Latency: number;
+  /** Total input/prompt tokens over the queried range. */
+  totalInputTokens: number;
+  /** Total output/completion tokens over the queried range. */
+  totalOutputTokens: number;
+  /** Attributed cost over the queried range, in USD. */
+  totalCostUsd: number;
   timeSeries: AppMetricsBucket[];
 }
 
@@ -261,6 +304,12 @@ export interface AppMetricsBucket {
   requestCount: number;
   errorCount: number;
   avgLatency: number;
+  /** Input/prompt tokens in this bucket. */
+  inputTokens: number;
+  /** Output/completion tokens in this bucket. */
+  outputTokens: number;
+  /** Attributed cost for this bucket, in USD. */
+  costUsd: number;
 }
 
 /**
@@ -315,6 +364,9 @@ export async function getAppMetrics(
       p50Latency: 0,
       p95Latency: 0,
       p99Latency: 0,
+      totalInputTokens: 0,
+      totalOutputTokens: 0,
+      totalCostUsd: 0,
       timeSeries: [],
     };
   }
@@ -326,6 +378,9 @@ export async function getAppMetrics(
   let weightedP50 = 0;
   let weightedP95 = 0;
   let weightedP99 = 0;
+  let totalInputTokens = 0;
+  let totalOutputTokens = 0;
+  let totalCostMicroUsd = 0;
 
   const timeSeries: AppMetricsBucket[] = [];
 
@@ -339,13 +394,23 @@ export async function getAppMetrics(
     weightedP95 += (item.p95Latency || 0) * bucketRequests;
     weightedP99 += (item.p99Latency || 0) * bucketRequests;
 
+    const bucketInputTokens = item.totalInputTokens || 0;
+    const bucketOutputTokens = item.totalOutputTokens || 0;
+    const bucketCostMicroUsd = item.totalCostMicroUsd || 0;
+    totalInputTokens += bucketInputTokens;
+    totalOutputTokens += bucketOutputTokens;
+    totalCostMicroUsd += bucketCostMicroUsd;
+
     // Extract hour bucket from sortId: "METRICS#yyyy-MM-dd-HH"
     const bucket = item.sortId.replace('METRICS#', '');
     timeSeries.push({
       timestamp: hourBucketToTimestamp(bucket),
       requestCount: bucketRequests,
       errorCount: (item.clientErrorCount || 0) + (item.serverErrorCount || 0),
       avgLatency: item.p50Latency || 0,
+      inputTokens: bucketInputTokens,
+      outputTokens: bucketOutputTokens,
+      costUsd: microUsdToUsd(bucketCostMicroUsd),
     });
   }
 
@@ -357,6 +422,9 @@ export async function getAppMetrics(
     p50Latency: totalRequests > 0 ? weightedP50 / totalRequests : 0,
     p95Latency: totalRequests > 0 ? weightedP95 / totalRequests : 0,
     p99Latency: totalRequests > 0 ? weightedP99 / totalRequests : 0,
+    totalInputTokens,
+    totalOutputTokens,
+    totalCostUsd: microUsdToUsd(totalCostMicroUsd),
     timeSeries,
   };
 }