sandoichi 0.4.1 → 0.5.0

package/src/provider-usage.mjs

@@ -11,6 +11,8 @@ const LOCK_WAIT_MS = 10;
 const LOCK_ATTEMPTS = 250;
 const STALE_LOCK_MS = 30_000;
 const COST_SCOPES = new Set(['session', 'event']);
+const COST_SOURCES = new Set(['host-reported', 'provider-reported']);
+const AGGREGATIONS = new Set(['session']);
 
 export const PROVIDER_USAGE_SCHEMA = SCHEMA;
 export const PROVIDER_USAGE_VERSION = VERSION;
@@ -71,11 +73,12 @@ function reportedCost(value) {
 function attachReportedCost(records, totalCostUsd) {
   if (!records.length || !usd(totalCostUsd)) return records;
   return records.map((item, index) => index === records.length - 1
-    ? { ...item, totalCostUsd, costScope: 'session' } : item);
+    ? { ...item, totalCostUsd, costScope: 'session', costSource: 'host-reported' } : item);
 }
 
 function usageRecord({ host, source, sourceKey, sessionId, turnId, at, inputTokens, cachedInputTokens = 0,
-  cacheWriteInputTokens = 0, outputTokens, reasoningOutputTokens = 0, totalCostUsd, arm, experimentId, workloadId }) {
+  cacheWriteInputTokens = 0, outputTokens, reasoningOutputTokens = 0, totalCostUsd, costSource,
+  aggregation, turnCount, arm, experimentId, workloadId, identityAt }) {
   if (!text(host) || !text(source) || !text(sourceKey)
     || (sessionId !== null && !text(sessionId)) || (turnId !== null && !text(turnId))
     || !text(at) || !counter(inputTokens) || !counter(cachedInputTokens)
@@ -83,8 +86,10 @@ function usageRecord({ host, source, sourceKey, sessionId, turnId, at, inputToke
     || !counter(outputTokens) || !counter(reasoningOutputTokens) || reasoningOutputTokens > outputTokens) return null;
   const totalTokens = safeSum(inputTokens, outputTokens);
   if (totalTokens === null) return null;
-  const identity = JSON.stringify({ host, source, sourceKey, at, inputTokens, cachedInputTokens,
-    cacheWriteInputTokens, outputTokens, reasoningOutputTokens, totalTokens, arm, experimentId, workloadId });
+  const identity = JSON.stringify({ host, source, sourceKey,
+    ...(identityAt === undefined ? { sessionId, turnId } : { at: identityAt }),
+    inputTokens, cachedInputTokens, cacheWriteInputTokens, outputTokens, reasoningOutputTokens,
+    totalTokens, aggregation, turnCount, arm, experimentId, workloadId });
   const result = {
     eventKey: `usage:${host}:${sha256(identity)}`,
     schema: SCHEMA, version: VERSION, host, source,
@@ -93,15 +98,18 @@ function usageRecord({ host, source, sourceKey, sessionId, turnId, at, inputToke
     reasoningOutputTokens, totalTokens,
   };
   if (totalCostUsd !== undefined && usd(totalCostUsd)) result.totalCostUsd = totalCostUsd;
+  if (costSource !== undefined) result.costSource = costSource;
+  if (aggregation !== undefined) result.aggregation = aggregation;
+  if (turnCount !== undefined) result.turnCount = turnCount;
   if (arm !== undefined) result.arm = arm;
   if (experimentId !== undefined) result.experimentId = experimentId;
   if (workloadId !== undefined) result.workloadId = workloadId;
   return result;
 }
 
-function claudeRecord(value, index, { sessionId = null, turnId = null, now, arm, experimentId, workloadId } = {}) {
-  if (value.type !== 'assistant' || !record(value.message?.usage)) return null;
-  const usage = value.message.usage;
+function claudeUsageRecord(usage, { source, sourceKey, sessionId = null, turnId = null, at,
+  arm, experimentId, workloadId, aggregation, turnCount, identityAt } = {}) {
+  if (!record(usage)) return null;
   const inputTokens = usage.input_tokens;
   const cachedInputTokens = optionalCounter(usage.cache_read_input_tokens);
   const cacheWriteInputTokens = optionalCounter(usage.cache_creation_input_tokens);
@@ -112,13 +120,37 @@ function claudeRecord(value, index, { sessionId = null, turnId = null, now, arm,
   const totalTokens = usage.total_tokens === undefined ? safeSum(totalInputTokens, outputTokens) : usage.total_tokens;
   if (!counter(totalTokens) || totalTokens !== totalInputTokens + outputTokens) return null;
   return usageRecord({
-    host: 'claude', source: 'claude-transcript', sourceKey: value.uuid ?? value.request_id ?? value.timestamp ?? String(index),
-    sessionId, turnId: value.turn_id ?? turnId, at: isoDate(value.timestamp, now),
+    host: 'claude', source, sourceKey, sessionId, turnId, at,
     inputTokens: totalInputTokens, cachedInputTokens, cacheWriteInputTokens, outputTokens,
+    aggregation, turnCount, identityAt,
     arm, experimentId, workloadId,
   });
 }
 
+function claudeRecord(value, index, { sessionId = null, turnId = null, now, arm, experimentId, workloadId } = {}) {
+  if (value.type !== 'assistant' || !record(value.message?.usage)) return null;
+  const identityAt = value.timestamp === undefined ? undefined : isoDate(value.timestamp, now);
+  return claudeUsageRecord(value.message.usage, {
+    source: 'claude-transcript', sourceKey: value.message.id ?? value.uuid ?? value.request_id ?? value.timestamp ?? String(index),
+    sessionId, turnId: value.turn_id ?? turnId, at: isoDate(value.timestamp, now), identityAt,
+    arm, experimentId, workloadId,
+  });
+}
+
+function claudeResultRecord(value, entries, { sessionId = null, arm, experimentId, workloadId, now } = {}) {
+  const identityTimestamp = value.timestamp ?? entries.slice().reverse().map((entry) => entry.value.timestamp).find(text);
+  const at = identityTimestamp ?? now;
+  const messages = new Set(entries
+    .filter(({ value: entry }) => entry.type === 'assistant' && record(entry.message?.usage))
+    .map(({ value: entry }, index) => entry.message.id ?? entry.uuid ?? entry.request_id ?? String(index)));
+  const turnCount = counter(value.num_turns ?? value.numTurns) ? value.num_turns ?? value.numTurns : messages.size || undefined;
+  return claudeUsageRecord(value.usage, {
+    source: 'claude-result', sourceKey: sessionId === null ? value.uuid ?? value.request_id ?? 'final' : 'session',
+    sessionId, turnId: null, at: isoDate(at, now), identityAt: identityTimestamp === undefined ? undefined : isoDate(identityTimestamp, now),
+    aggregation: 'session', turnCount, arm, experimentId, workloadId,
+  });
+}
+
 function codexRecord(value, index, { sessionId = null, turnId = null, now, arm, experimentId, workloadId } = {}) {
   const usage = value.type === 'turn.completed'
     ? value.usage
@@ -135,9 +167,10 @@ function codexRecord(value, index, { sessionId = null, turnId = null, now, arm,
     || !counter(outputTokens) || reasoningOutputTokens === null) return null;
   const totalTokens = usage.total_tokens === undefined ? safeSum(inputTokens, outputTokens) : usage.total_tokens;
   if (!counter(totalTokens) || totalTokens !== inputTokens + outputTokens) return null;
+  const identityAt = value.timestamp === undefined ? undefined : isoDate(value.timestamp, now);
   return usageRecord({
     host: 'codex', source: 'codex-transcript', sourceKey: value.turn_id ?? value.id ?? value.timestamp ?? String(index),
-    sessionId, turnId: value.turn_id ?? value.id ?? (value.timestamp ? `at:${value.timestamp}` : turnId), at: isoDate(value.timestamp, now),
+    sessionId, turnId: value.turn_id ?? value.id ?? (value.timestamp ? `at:${value.timestamp}` : turnId), at: isoDate(value.timestamp, now), identityAt,
     inputTokens, cachedInputTokens, cacheWriteInputTokens, outputTokens, reasoningOutputTokens,
     arm, experimentId, workloadId,
   });
@@ -145,9 +178,18 @@ function codexRecord(value, index, { sessionId = null, turnId = null, now, arm,
 
 export function parseClaudeTranscript(textValue, options = {}) {
   const entries = jsonLines(textValue);
-  const records = entries.map(({ value }, index) => claudeRecord(value, index, options)).filter(Boolean);
+  const unique = new Map();
+  entries.forEach(({ value }, index) => {
+    const result = claudeRecord(value, index, options);
+    if (!result) return;
+    const messageId = text(value.message?.id);
+    const key = messageId ? `message:${messageId}` : result.eventKey;
+    unique.set(key, result);
+  });
+  const resultEntry = entries.findLast(({ value }) => value.type === 'result' && record(value.usage));
   const totalCostUsd = options.totalCostUsd ?? entries.slice().reverse().map(({ value }) => reportedCost(value)).find(usd);
-  return attachReportedCost(records, totalCostUsd);
+  const aggregate = resultEntry ? claudeResultRecord(resultEntry.value, entries, options) : null;
+  return attachReportedCost(aggregate ? [aggregate] : [...unique.values()], totalCostUsd);
 }
 
 export function parseCodexTranscript(textValue, options = {}) {
@@ -188,6 +230,9 @@ function validateUsage(value) {
     || (value.experimentId !== undefined && !text(value.experimentId))
     || (value.workloadId !== undefined && !text(value.workloadId))
     || (value.totalCostUsd !== undefined && !usd(value.totalCostUsd))
+    || (value.costSource !== undefined && !COST_SOURCES.has(value.costSource))
+    || (value.aggregation !== undefined && !AGGREGATIONS.has(value.aggregation))
+    || (value.turnCount !== undefined && !counter(value.turnCount))
     || (value.costScope !== undefined && (!COST_SCOPES.has(value.costScope) || value.totalCostUsd === undefined))) {
     throw new Error('provider usage record is invalid');
   }
@@ -251,11 +296,19 @@ export function appendProviderUsage({ storagePath = defaultProviderUsagePath(),
     const state = readProviderUsage(filePath);
     const existing = new Map(state.records.map((item, index) => [item.eventKey, index]));
     for (const item of records) {
-      const index = existing.get(item.eventKey);
+      const index = item.aggregation === 'session'
+        ? item.sessionId === null
+          ? existing.get(item.eventKey) ?? -1
+          : state.records.findIndex((candidate) => candidate.aggregation === 'session'
+            && candidate.host === item.host && candidate.source === item.source && candidate.sessionId === item.sessionId)
+        : existing.get(item.eventKey);
       if (index === undefined) {
         existing.set(item.eventKey, state.records.length);
         state.records.push(item);
-      } else if (item.costScope === 'session') {
+      } else if (index < 0) {
+        existing.set(item.eventKey, state.records.length);
+        state.records.push(item);
+      } else if (item.aggregation === 'session' || item.costScope === 'session') {
         state.records[index] = item;
       }
     }
@@ -278,9 +331,20 @@ export function collectProviderUsage({ host, transcriptPath, sessionId = null, t
 }
 
 export function buildProviderUsageReport(state, { sessionId, pricing } = {}) {
-  const records = validateState(state).records.filter((item) => sessionId === undefined || item.sessionId === sessionId);
-  const sessions = new Set(records.map((item) => `${item.host}\0${item.sessionId ?? '<unknown>'}`));
-  const turns = new Set(records.map((item, index) => `${item.host}\0${item.sessionId ?? '<unknown>'}\0${item.turnId ?? `record:${index}`}`));
+  const selected = validateState(state).records.filter((item) => sessionId === undefined || item.sessionId === sessionId);
+  const aggregateSessions = new Set(selected
+    .filter((item) => item.aggregation === 'session' && item.sessionId !== null)
+    .map(sessionKey));
+  const records = selected.filter((item) => item.aggregation === 'session'
+    || item.sessionId === null || !aggregateSessions.has(sessionKey(item)));
+  const sessions = new Set(records.map(sessionKey));
+  const turns = new Set(records.map(turnKey));
+  const aggregateRecords = records.filter((item) => item.aggregation === 'session');
+  const aggregateCounts = aggregateRecords.map((item) => item.turnCount);
+  const nonAggregateTurns = new Set(records.filter((item) => item.aggregation !== 'session').map(turnKey));
+  const turnCount = aggregateRecords.length
+    ? aggregateCounts.every(counter) ? safeSum(...aggregateCounts, nonAggregateTurns.size) : null
+    : turns.size;
   const sum = (field) => records.reduce((total, item) => total + item[field], 0);
   const weightedCostUnits = records.reduce((total, item) => total + computeWeightedUsage(item, pricing).costUnits, 0);
   const freshInputTokens = records.reduce((total, item) => total + item.inputTokens - item.cachedInputTokens - item.cacheWriteInputTokens, 0);
@@ -289,7 +353,7 @@ export function buildProviderUsageReport(state, { sessionId, pricing } = {}) {
   const effectiveRate = totalCostUsd !== null && totalTokens(records) > 0
     ? totalCostUsd / totalTokens(records) * 1_000_000 : null;
   const cost = {
-    status: billing.complete ? 'provider-reported' : 'unavailable',
+    status: billing.complete ? billing.source : billing.partial ? billing.source : 'unavailable',
     coverage: billing.complete ? 'complete' : billing.partial ? 'partial' : 'none',
     totalCostUsd,
     effectiveRateUsdPerMillionTokens: effectiveRate,
@@ -299,18 +363,28 @@ export function buildProviderUsageReport(state, { sessionId, pricing } = {}) {
     inputTokens: sum('inputTokens'), cachedInputTokens: sum('cachedInputTokens'),
     cacheWriteInputTokens: sum('cacheWriteInputTokens'), freshInputTokens,
     outputTokens: sum('outputTokens'), reasoningOutputTokens: sum('reasoningOutputTokens'), totalTokens: sum('totalTokens'),
-    turnCount: turns.size, weightedCostUnits,
+    turnCount, weightedCostUnits,
     weightedCost: { source: 'weighted-estimate', costUnits: weightedCostUnits },
     cost,
     totalCostUsd,
-    providerReportedCostUsd: totalCostUsd,
+    providerReportedCostUsd: billing.source === 'provider-reported' && billing.complete ? totalCostUsd : null,
     sessionBlendedEffectiveRateUsdPerMillionTokens: effectiveRate,
     costSource: cost.status,
   };
 }
 
 function billingKey(item) {
-  return `${item.host}\0${item.sessionId ?? '<unknown>'}`;
+  return sessionKey(item);
+}
+
+function sessionKey(item) {
+  return item.sessionId === null
+    ? `${item.host}\0unknown:${item.eventKey}`
+    : `${item.host}\0${item.sessionId}`;
+}
+
+function turnKey(item, index) {
+  return `${sessionKey(item)}\0${item.turnId ?? `record:${index}`}`;
 }
 
 function reportedCostSummary(records) {
@@ -321,18 +395,24 @@ function reportedCostSummary(records) {
     groups.set(billingKey(item), entries);
   });
   const totals = [];
+  const sources = new Set();
   for (const entries of groups.values()) {
     const sessionCosts = entries.filter((item) => item.costScope === 'session' && usd(item.totalCostUsd));
     if (sessionCosts.length) {
       const latest = sessionCosts.reduce((left, right) => right.at >= left.at ? right : left);
       totals.push(latest.totalCostUsd);
+      sources.add(latest.costSource ?? 'unknown');
       continue;
     }
-    if (entries.every((item) => usd(item.totalCostUsd))) totals.push(sumUsd(entries.map((item) => item.totalCostUsd)));
+    if (entries.every((item) => usd(item.totalCostUsd))) {
+      totals.push(sumUsd(entries.map((item) => item.totalCostUsd)));
+      entries.forEach((item) => sources.add(item.costSource ?? 'unknown'));
+    }
   }
   const complete = groups.size > 0 && totals.length === groups.size && !totals.includes(null);
   const totalCostUsd = complete ? sumUsd(totals) : null;
-  return { complete: complete && totalCostUsd !== null, partial: totals.length > 0, totalCostUsd };
+  const source = sources.size === 1 ? [...sources][0] : sources.size ? 'mixed' : null;
+  return { complete: complete && totalCostUsd !== null, partial: totals.length > 0, totalCostUsd, source: source ?? 'unknown' };
 }
 
 function totalTokens(records) {