Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 35 additions & 12 deletions apps/sim/lib/billing/client/upgrade.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,17 +41,31 @@ export function useSubscriptionUpgrade() {
throw new Error('User not authenticated')
}

let currentSubscriptionId: string | undefined
let currentSubscriptionRowId: string | undefined
let currentStripeSubscriptionId: string | undefined
let allSubscriptions: any[] = []
try {
const listResult = await client.subscription.list()
allSubscriptions = listResult.data || []
const activePersonalSub = allSubscriptions.find(
(sub: any) => hasPaidSubscriptionStatus(sub.status) && sub.referenceId === userId
)
currentSubscriptionId = activePersonalSub?.id
currentSubscriptionRowId = activePersonalSub?.id
currentStripeSubscriptionId = activePersonalSub?.stripeSubscriptionId
Comment thread
icecrasher321 marked this conversation as resolved.
} catch (_e) {
currentSubscriptionId = undefined
currentSubscriptionRowId = undefined
currentStripeSubscriptionId = undefined
}

if (currentSubscriptionRowId && !currentStripeSubscriptionId) {
logger.error('Active paid subscription is missing its Stripe subscription ID', {
userId,
subscriptionRowId: currentSubscriptionRowId,
targetPlan,
})
throw new Error(
'We could not match your current plan with our payment provider. Please contact support before upgrading so you are not charged twice.'
)
}

let referenceId = userId
Expand Down Expand Up @@ -137,36 +151,45 @@ export function useSubscriptionUpgrade() {
...(annual && { annual: true }),
} as const

const finalParams = currentSubscriptionId
? { ...upgradeParams, subscriptionId: currentSubscriptionId }
const finalParams = currentStripeSubscriptionId
? { ...upgradeParams, subscriptionId: currentStripeSubscriptionId }
: upgradeParams

logger.info(
currentSubscriptionId ? 'Upgrading existing subscription' : 'Creating new subscription',
{ targetPlan, planName, annual, currentSubscriptionId, referenceId }
currentStripeSubscriptionId
? 'Upgrading existing subscription'
: 'Creating new subscription',
{
targetPlan,
planName,
annual,
currentStripeSubscriptionId,
currentSubscriptionRowId,
referenceId,
}
)

await betterAuthSubscription.upgrade(finalParams)

if (targetPlan === 'team' && currentSubscriptionId && referenceId !== userId) {
if (targetPlan === 'team' && currentSubscriptionRowId && referenceId !== userId) {
try {
logger.info('Transferring subscription to organization after upgrade', {
subscriptionId: currentSubscriptionId,
subscriptionId: currentSubscriptionRowId,
organizationId: referenceId,
})

try {
await requestJson(subscriptionTransferContract, {
params: { id: currentSubscriptionId },
params: { id: currentSubscriptionRowId },
body: { organizationId: referenceId },
})
logger.info('Successfully transferred subscription to organization', {
subscriptionId: currentSubscriptionId,
subscriptionId: currentSubscriptionRowId,
organizationId: referenceId,
})
} catch (transferError) {
logger.error('Failed to transfer subscription to organization', {
subscriptionId: currentSubscriptionId,
subscriptionId: currentSubscriptionRowId,
organizationId: referenceId,
error:
transferError instanceof ApiClientError
Expand Down
82 changes: 82 additions & 0 deletions apps/sim/lib/logs/execution/logging-factory.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -576,4 +576,86 @@ describe('calculateCostSummary', () => {
expect(Object.keys(result.charges)).toHaveLength(0)
expect(result.totalCost).toBe(BASE_EXECUTION_CHARGE)
})

test('does not double-count the synthetic workflow root (aggregate cost over leaves)', () => {
// buildTraceSpans wraps every run in a synthetic { type: 'workflow' } root
// whose cost.total is the SUM of its leaves. Counting that root in addition
// to the leaves double-charges the run — the root must be a pass-through.
const traceSpans = [
{
id: 'workflow-execution',
name: 'Workflow Execution',
type: 'workflow',
cost: { total: 0.04 }, // == agent(0.03) + exa(0.01)
children: [
{
id: 'agent-1',
name: 'Agent',
type: 'agent',
model: 'gpt-4o',
cost: { input: 0.01, output: 0.02, total: 0.03 },
tokens: { input: 100, output: 200, total: 300 },
},
{
id: 'exa-1',
name: 'Exa Search',
type: 'tool',
cost: { input: 0, output: 0, total: 0.01 },
},
],
},
]

const result = calculateCostSummary(traceSpans)

// The 0.04 root aggregate is NOT added on top of its leaves.
expect(result.charges['Workflow Execution']).toBeUndefined()
expect(result.models['gpt-4o'].total).toBe(0.03)
expect(result.charges['Exa Search'].total).toBe(0.01)
expect(result.totalCost).toBeCloseTo(0.04 + BASE_EXECUTION_CHARGE, 10)
const ledgerSum =
result.baseExecutionCharge +
Object.values(result.models).reduce((s, m) => s + m.total, 0) +
Object.values(result.charges).reduce((s, c) => s + c.total, 0)
expect(ledgerSum).toBeCloseTo(result.totalCost, 10)
})

test('does not double-count nested sub-workflow roots', () => {
// A sub-workflow call nests another synthetic { type: 'workflow' } root
// (captureChildWorkflowLogs runs buildTraceSpans on the child). Both the
// outer root and the inner sub-workflow root carry aggregate costs; only the
// leaf agent inside should be billed.
const traceSpans = [
{
id: 'workflow-execution',
name: 'Workflow Execution',
type: 'workflow',
cost: { total: 0.03 },
children: [
{
id: 'subworkflow-root',
name: 'Workflow Execution',
type: 'workflow',
cost: { total: 0.03 },
children: [
{
id: 'child-agent',
name: 'Agent',
type: 'agent',
model: 'gpt-4o',
cost: { input: 0.01, output: 0.02, total: 0.03 },
tokens: { input: 100, output: 200, total: 300 },
},
],
},
],
},
]

const result = calculateCostSummary(traceSpans)

expect(result.charges['Workflow Execution']).toBeUndefined()
expect(result.models['gpt-4o'].total).toBe(0.03)
expect(result.totalCost).toBeCloseTo(0.03 + BASE_EXECUTION_CHARGE, 10)
})
})
23 changes: 18 additions & 5 deletions apps/sim/lib/logs/execution/logging-factory.ts
Original file line number Diff line number Diff line change
Expand Up @@ -165,19 +165,32 @@ export function calculateCostSummary(traceSpans: CostTraceSpan[] | undefined): C
const costSpans: BillableTraceSpan[] = []

for (const span of spans) {
// `workflow`-typed spans are aggregate containers, not billable units: the
// synthetic "Workflow Execution" root (added to every run by
// buildTraceSpans) and any nested sub-workflow root carry a `cost.total`
// equal to the SUM of their descendants. Counting that aggregate in
// addition to the descendants double-charges the run, so treat these as
// pass-through: never count their own cost, always recurse into all
// children where the real billable leaves (agents, tools) live.
const isAggregateContainer = span.type === 'workflow'
const hasOwnCost = hasBillableCost(span)
if (hasOwnCost) {
const countOwnCost = hasOwnCost && !isAggregateContainer

if (countOwnCost) {
costSpans.push(span)
}

if (span.children && Array.isArray(span.children)) {
if (hasOwnCost) {
// Parent already accounts for its model segments; only recurse into
// non-model children (e.g. nested workflow spans) to find further
// billable units.
if (countOwnCost) {
// Authoritative leaf (e.g. an agent block whose block-level cost is set
// by the provider response and already accounts for its model
// segments): only recurse into non-model children to find further
// standalone billable units, skipping the model-breakdown duplicates.
const nonModelChildren = span.children.filter((child) => !isModelBreakdownSpan(child))
costSpans.push(...collectCostSpans(nonModelChildren))
} else {
// Container (workflow / sub-workflow root) or a no-cost parent: recurse
// into everything so nested billable leaves are counted exactly once.
costSpans.push(...collectCostSpans(span.children))
}
}
Expand Down
Loading