@rudderhq/server 0.2.1 → 0.2.2-canary.0

package/dist/services/runtime-kernel/heartbeat.js

@@ -1,6 +1,6 @@
 import path from "node:path";
 import { and, asc, desc, eq, gt, gte, inArray, lte, sql } from "drizzle-orm";
-import { AGENT_RUN_CONCURRENCY_DEFAULT, AGENT_RUN_CONCURRENCY_MAX, AGENT_RUN_CONCURRENCY_MIN, } from "@rudderhq/shared";
+import { AGENT_RUN_CONCURRENCY_DEFAULT, AGENT_RUN_CONCURRENCY_MAX, AGENT_RUN_CONCURRENCY_MIN, summarizeTokenUsage, } from "@rudderhq/shared";
 import { agents, agentRuntimeState, agentTaskSessions, agentWakeupRequests, activityLog, heartbeatRunEvents, heartbeatRuns, issueComments, issues, organizations, projects, } from "@rudderhq/db";
 import { conflict, notFound } from "../../errors.js";
 import { createExecutionScores, observeExecutionEvent, updateExecutionObservation, updateExecutionTraceIO, updateExecutionTraceName, updateExecutionTraceSession, withExecutionObservation, } from "../../langfuse.js";
@@ -2481,13 +2481,19 @@ export function heartbeatService(db) {
     async function updateRuntimeState(agent, run, result, session, normalizedUsage) {
         await ensureRuntimeState(agent);
         const usage = normalizedUsage ?? normalizeUsageTotals(result.usage);
-        const inputTokens = usage?.inputTokens ?? 0;
+        const rawInputTokens = usage?.inputTokens ?? 0;
         const outputTokens = usage?.outputTokens ?? 0;
         const cachedInputTokens = usage?.cachedInputTokens ?? 0;
         const billingType = normalizeLedgerBillingType(result.billingType);
         const additionalCostCents = normalizeBilledCostCents(result.costUsd, billingType);
-        const hasTokenUsage = inputTokens > 0 || outputTokens > 0 || cachedInputTokens > 0;
+        const hasTokenUsage = rawInputTokens > 0 || outputTokens > 0 || cachedInputTokens > 0;
         const provider = result.provider ?? "unknown";
+        const tokenSummary = summarizeTokenUsage({
+            provider,
+            inputTokens: rawInputTokens,
+            cachedInputTokens,
+            outputTokens,
+        });
         const biller = resolveLedgerBiller(result);
         const ledgerScope = await resolveLedgerScopeForRun(db, agent.orgId, run);
         await db
@@ -2498,7 +2504,7 @@ export function heartbeatService(db) {
             lastRunId: run.id,
             lastRunStatus: run.status,
             lastError: result.errorMessage ?? null,
-            totalInputTokens: sql `${agentRuntimeState.totalInputTokens} + ${inputTokens}`,
+            totalInputTokens: sql `${agentRuntimeState.totalInputTokens} + ${tokenSummary.promptTokens}`,
             totalOutputTokens: sql `${agentRuntimeState.totalOutputTokens} + ${outputTokens}`,
             totalCachedInputTokens: sql `${agentRuntimeState.totalCachedInputTokens} + ${cachedInputTokens}`,
             totalCostCents: sql `${agentRuntimeState.totalCostCents} + ${additionalCostCents}`,
@@ -2516,7 +2522,7 @@ export function heartbeatService(db) {
                 biller,
                 billingType,
                 model: result.model ?? "unknown",
-                inputTokens,
+                inputTokens: rawInputTokens,
                 cachedInputTokens,
                 outputTokens,
                 costCents: additionalCostCents,