From e9ff1a8ce37602f3461f1d6c9f9427119fda375b Mon Sep 17 00:00:00 2001 From: sahu-virendra-1908 Date: Thu, 28 May 2026 15:11:09 +0530 Subject: [PATCH 1/2] fix: make agent cost accounting updates transactional Signed-off-by: sahu-virendra-1908 --- server/src/core/heartbeat.ts | 87 +++++++++++++++++++++--------------- 1 file changed, 52 insertions(+), 35 deletions(-) diff --git a/server/src/core/heartbeat.ts b/server/src/core/heartbeat.ts index 43940fe2..1df439a4 100644 --- a/server/src/core/heartbeat.ts +++ b/server/src/core/heartbeat.ts @@ -962,46 +962,63 @@ export function heartbeatService(db: Db) { const additionalCostCents = Math.max(0, Math.round((result.costUsd ?? 0) * 100)); const hasTokenUsage = inputTokens > 0 || outputTokens > 0 || cachedInputTokens > 0; - await db - .update(agentRuntimeState) + +const now = new Date(); + +await db.transaction(async (tx) => { + await tx + .update(agentRuntimeState) + .set({ + adapterType: agent.adapterType, + sessionId: session.legacySessionId, + lastRunId: run.id, + lastRunStatus: run.status, + lastError: result.errorMessage ?? null, + + totalInputTokens: sql`${agentRuntimeState.totalInputTokens} + ${inputTokens}`, + + totalOutputTokens: sql`${agentRuntimeState.totalOutputTokens} + ${outputTokens}`, + + totalCachedInputTokens: sql`${agentRuntimeState.totalCachedInputTokens} + ${cachedInputTokens}`, + + totalCostCents: sql`${agentRuntimeState.totalCostCents} + ${additionalCostCents}`, + + updatedAt: now, + }) + .where(eq(agentRuntimeState.agentId, agent.id)); + + if (additionalCostCents > 0 || hasTokenUsage) { + await tx.insert(costEvents).values({ + projectId: agent.projectId, + agentId: agent.id, + + provider: result.provider ?? "unknown", + model: result.model ?? "unknown", + + inputTokens, + outputTokens, + + costCents: additionalCostCents, + + occurredAt: now, + }); + } + + if (additionalCostCents > 0) { + await tx + .update(agents) .set({ - adapterType: agent.adapterType, - sessionId: session.legacySessionId, - lastRunId: run.id, - lastRunStatus: run.status, - lastError: result.errorMessage ?? null, - totalInputTokens: sql`${agentRuntimeState.totalInputTokens} + ${inputTokens}`, - totalOutputTokens: sql`${agentRuntimeState.totalOutputTokens} + ${outputTokens}`, - totalCachedInputTokens: sql`${agentRuntimeState.totalCachedInputTokens} + ${cachedInputTokens}`, - totalCostCents: sql`${agentRuntimeState.totalCostCents} + ${additionalCostCents}`, - updatedAt: new Date(), - }) - .where(eq(agentRuntimeState.agentId, agent.id)); + spentMonthlyCents: sql`${agents.spentMonthlyCents} + ${additionalCostCents}`, - if (additionalCostCents > 0 || hasTokenUsage) { - await db.insert(costEvents).values({ - projectId: agent.projectId, - agentId: agent.id, - provider: result.provider ?? "unknown", - model: result.model ?? "unknown", - inputTokens, - outputTokens, - costCents: additionalCostCents, - occurredAt: new Date(), - }); - } + updatedAt: now, + }) + .where(eq(agents.id, agent.id)); + } +}); - if (additionalCostCents > 0) { - await db - .update(agents) - .set({ - spentMonthlyCents: sql`${agents.spentMonthlyCents} + ${additionalCostCents}`, - updatedAt: new Date(), - }) - .where(eq(agents.id, agent.id)); - } } + async function startNextQueuedRunForAgent(agentId: string) { return withAgentStartLock(agentId, async () => { const agent = await getAgent(agentId); From e2398a780b32fdcafab47d4e90165c022e7faabd Mon Sep 17 00:00:00 2001 From: sahu-virendra-1908 Date: Thu, 28 May 2026 15:11:53 +0530 Subject: [PATCH 2/2] fix: make agent cost accounting updates transactional Signed-off-by: sahu-virendra-1908 --- server/src/core/heartbeat.ts | 158 +++++++++++++++++------------------ 1 file changed, 79 insertions(+), 79 deletions(-) diff --git a/server/src/core/heartbeat.ts b/server/src/core/heartbeat.ts index 1df439a4..b92f9860 100644 --- a/server/src/core/heartbeat.ts +++ b/server/src/core/heartbeat.ts @@ -384,11 +384,11 @@ function resolveNextSessionState(input: { const displayId = truncateDisplayId( explicitDisplayId ?? - (codec.getDisplayId ? codec.getDisplayId(deserialized) : null) ?? - readNonEmptyString(deserialized?.sessionId) ?? - (shouldUsePrevious ? previousDisplayId : null) ?? - explicitSessionId ?? - (shouldUsePrevious ? previousLegacySessionId : null), + (codec.getDisplayId ? codec.getDisplayId(deserialized) : null) ?? + readNonEmptyString(deserialized?.sessionId) ?? + (shouldUsePrevious ? previousDisplayId : null) ?? + explicitSessionId ?? + (shouldUsePrevious ? previousLegacySessionId : null), ); const legacySessionId = @@ -470,8 +470,8 @@ export function heartbeatService(db: Db) { ); return truncateDisplayId( existingTaskSession?.sessionDisplayId ?? - (codec.getDisplayId ? codec.getDisplayId(parsedParams) : null) ?? - readNonEmptyString(parsedParams?.sessionId), + (codec.getDisplayId ? codec.getDisplayId(parsedParams) : null) ?? + readNonEmptyString(parsedParams?.sessionId), ); } @@ -489,10 +489,10 @@ export function heartbeatService(db: Db) { const contextProjectId = readNonEmptyString(context.projectId); const issueProjectId = issueId ? await db - .select({ projectId: issues.projectId }) - .from(issues) - .where(and(eq(issues.id, issueId), eq(issues.projectId, agent.projectId))) - .then((rows) => rows[0]?.projectId ?? null) + .select({ projectId: issues.projectId }) + .from(issues) + .where(and(eq(issues.id, issueId), eq(issues.projectId, agent.projectId))) + .then((rows) => rows[0]?.projectId ?? null) : null; const resolvedProjectId = issueProjectId ?? contextProjectId; const useProjectWorkspace = opts?.useProjectWorkspace !== false; @@ -500,15 +500,15 @@ export function heartbeatService(db: Db) { const projectWorkspaceRows = workspaceProjectId ? await db - .select() - .from(projectWorkspaces) - .where( - and( - eq(projectWorkspaces.projectId, agent.projectId), - eq(projectWorkspaces.projectId, workspaceProjectId), - ), - ) - .orderBy(asc(projectWorkspaces.createdAt), asc(projectWorkspaces.id)) + .select() + .from(projectWorkspaces) + .where( + and( + eq(projectWorkspaces.projectId, agent.projectId), + eq(projectWorkspaces.projectId, workspaceProjectId), + ), + ) + .orderBy(asc(projectWorkspaces.createdAt), asc(projectWorkspaces.id)) : []; const workspaceHints = projectWorkspaceRows.map((workspace) => ({ @@ -963,58 +963,58 @@ export function heartbeatService(db: Db) { const hasTokenUsage = inputTokens > 0 || outputTokens > 0 || cachedInputTokens > 0; -const now = new Date(); + const now = new Date(); -await db.transaction(async (tx) => { - await tx - .update(agentRuntimeState) - .set({ - adapterType: agent.adapterType, - sessionId: session.legacySessionId, - lastRunId: run.id, - lastRunStatus: run.status, - lastError: result.errorMessage ?? null, + await db.transaction(async (tx) => { + await tx + .update(agentRuntimeState) + .set({ + adapterType: agent.adapterType, + sessionId: session.legacySessionId, + lastRunId: run.id, + lastRunStatus: run.status, + lastError: result.errorMessage ?? null, - totalInputTokens: sql`${agentRuntimeState.totalInputTokens} + ${inputTokens}`, + totalInputTokens: sql`${agentRuntimeState.totalInputTokens} + ${inputTokens}`, - totalOutputTokens: sql`${agentRuntimeState.totalOutputTokens} + ${outputTokens}`, + totalOutputTokens: sql`${agentRuntimeState.totalOutputTokens} + ${outputTokens}`, - totalCachedInputTokens: sql`${agentRuntimeState.totalCachedInputTokens} + ${cachedInputTokens}`, + totalCachedInputTokens: sql`${agentRuntimeState.totalCachedInputTokens} + ${cachedInputTokens}`, - totalCostCents: sql`${agentRuntimeState.totalCostCents} + ${additionalCostCents}`, + totalCostCents: sql`${agentRuntimeState.totalCostCents} + ${additionalCostCents}`, - updatedAt: now, - }) - .where(eq(agentRuntimeState.agentId, agent.id)); + updatedAt: now, + }) + .where(eq(agentRuntimeState.agentId, agent.id)); - if (additionalCostCents > 0 || hasTokenUsage) { - await tx.insert(costEvents).values({ - projectId: agent.projectId, - agentId: agent.id, + if (additionalCostCents > 0 || hasTokenUsage) { + await tx.insert(costEvents).values({ + projectId: agent.projectId, + agentId: agent.id, - provider: result.provider ?? "unknown", - model: result.model ?? "unknown", + provider: result.provider ?? "unknown", + model: result.model ?? "unknown", - inputTokens, - outputTokens, + inputTokens, + outputTokens, - costCents: additionalCostCents, + costCents: additionalCostCents, - occurredAt: now, - }); - } + occurredAt: now, + }); + } - if (additionalCostCents > 0) { - await tx - .update(agents) - .set({ - spentMonthlyCents: sql`${agents.spentMonthlyCents} + ${additionalCostCents}`, + if (additionalCostCents > 0) { + await tx + .update(agents) + .set({ + spentMonthlyCents: sql`${agents.spentMonthlyCents} + ${additionalCostCents}`, - updatedAt: now, - }) - .where(eq(agents.id, agent.id)); - } -}); + updatedAt: now, + }) + .where(eq(agents.id, agent.id)); + } + }); } @@ -1181,19 +1181,19 @@ await db.transaction(async (tx) => { const issueAssigneeConfig = issueId ? await db - .select({ - assigneeAgentId: issues.assigneeAgentId, - assigneeAdapterOverrides: issues.assigneeAdapterOverrides, - }) - .from(issues) - .where(and(eq(issues.id, issueId), eq(issues.projectId, agent.projectId))) - .then((rows) => rows[0] ?? null) + .select({ + assigneeAgentId: issues.assigneeAgentId, + assigneeAdapterOverrides: issues.assigneeAdapterOverrides, + }) + .from(issues) + .where(and(eq(issues.id, issueId), eq(issues.projectId, agent.projectId))) + .then((rows) => rows[0] ?? null) : null; const issueAssigneeOverrides = issueAssigneeConfig && issueAssigneeConfig.assigneeAgentId === agent.id ? parseIssueAssigneeAdapterOverrides( - issueAssigneeConfig.assigneeAdapterOverrides, - ) + issueAssigneeConfig.assigneeAdapterOverrides, + ) : null; const taskSession = taskKey ? await getTaskSession(agent.projectId, agent.id, agent.adapterType, taskKey) @@ -1221,10 +1221,10 @@ await db.transaction(async (tx) => { ...(runtimeSessionResolution.warning ? [runtimeSessionResolution.warning] : []), ...(resetTaskSession && sessionResetReason ? [ - taskKey - ? `Skipping saved session resume for task "${taskKey}" because ${sessionResetReason}.` - : `Skipping saved session resume because ${sessionResetReason}.`, - ] + taskKey + ? `Skipping saved session resume for task "${taskKey}" because ${sessionResetReason}.` + : `Skipping saved session resume because ${sessionResetReason}.`, + ] : []), ]; const workspaceContextValue = { @@ -1245,9 +1245,9 @@ await db.transaction(async (tx) => { const runtimeSessionFallback = taskKey || resetTaskSession ? null : runtime.sessionId; const previousSessionDisplayId = truncateDisplayId( taskSessionForRun?.sessionDisplayId ?? - (sessionCodec.getDisplayId ? sessionCodec.getDisplayId(runtimeSessionParams) : null) ?? - readNonEmptyString(runtimeSessionParams?.sessionId) ?? - runtimeSessionFallback, + (sessionCodec.getDisplayId ? sessionCodec.getDisplayId(runtimeSessionParams) : null) ?? + readNonEmptyString(runtimeSessionParams?.sessionId) ?? + runtimeSessionFallback, ); const runtimeForAdapter = { sessionId: readNonEmptyString(runtimeSessionParams?.sessionId) ?? runtimeSessionFallback, @@ -1435,10 +1435,10 @@ await db.transaction(async (tx) => { const usageJson = adapterResult.usage || adapterResult.costUsd != null ? ({ - ...(adapterResult.usage ?? {}), - ...(adapterResult.costUsd != null ? { costUsd: adapterResult.costUsd } : {}), - ...(adapterResult.billingType ? { billingType: adapterResult.billingType } : {}), - } as Record) + ...(adapterResult.usage ?? {}), + ...(adapterResult.costUsd != null ? { costUsd: adapterResult.costUsd } : {}), + ...(adapterResult.billingType ? { billingType: adapterResult.billingType } : {}), + } as Record) : null; await setRunStatus(run.id, status, {