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
54 changes: 51 additions & 3 deletions apps/sim/app/api/billing/update-cost/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,11 @@ import { type NextRequest, NextResponse } from 'next/server'
import { z } from 'zod'
import { recordUsage } from '@/lib/billing/core/usage-log'
import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing'
import { BillingRouteOutcome } from '@/lib/copilot/generated/trace-attribute-values-v1'
import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
import { checkInternalApiKey } from '@/lib/copilot/request/http'
import { withIncomingGoSpan } from '@/lib/copilot/request/otel'
import { isBillingEnabled } from '@/lib/core/config/feature-flags'
import { type AtomicClaimResult, billingIdempotency } from '@/lib/core/idempotency/service'
import { generateRequestId } from '@/lib/core/utils/request'
Expand All @@ -28,8 +32,28 @@ const UpdateCostSchema = z.object({
/**
* POST /api/billing/update-cost
* Update user cost with a pre-calculated cost value (internal API key auth required)
*
* Parented under the Go-side `sim.update_cost` span via W3C traceparent
* propagation. Every mothership request that bills should therefore show
* the Go client span AND this Sim server span sharing one trace, with
* the actual usage/overage work nested below.
*/
export const POST = withRouteHandler(async (req: NextRequest) => {
export const POST = withRouteHandler((req: NextRequest) =>
withIncomingGoSpan(
req.headers,
TraceSpan.CopilotBillingUpdateCost,
{
[TraceAttr.HttpMethod]: 'POST',
[TraceAttr.HttpRoute]: '/api/billing/update-cost',
},
async (span) => updateCostInner(req, span)
)
)

async function updateCostInner(
req: NextRequest,
span: import('@opentelemetry/api').Span
): Promise<NextResponse> {
const requestId = generateRequestId()
const startTime = Date.now()
let claim: AtomicClaimResult | null = null
Expand All @@ -39,6 +63,8 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
logger.info(`[${requestId}] Update cost request started`)

if (!isBillingEnabled) {
span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.BillingDisabled)
span.setAttribute(TraceAttr.HttpStatusCode, 200)
return NextResponse.json({
success: true,
message: 'Billing disabled, cost update skipped',
Expand All @@ -54,6 +80,8 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
const authResult = checkInternalApiKey(req)
if (!authResult.success) {
logger.warn(`[${requestId}] Authentication failed: ${authResult.error}`)
span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.AuthFailed)
span.setAttribute(TraceAttr.HttpStatusCode, 401)
return NextResponse.json(
{
success: false,
Expand All @@ -69,8 +97,9 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
if (!validation.success) {
logger.warn(`[${requestId}] Invalid request body`, {
errors: validation.error.issues,
body,
})
span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.InvalidBody)
span.setAttribute(TraceAttr.HttpStatusCode, 400)
return NextResponse.json(
{
success: false,
Expand All @@ -85,6 +114,17 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
validation.data
const isMcp = source === 'mcp_copilot'

span.setAttributes({
[TraceAttr.UserId]: userId,
[TraceAttr.GenAiRequestModel]: model,
[TraceAttr.BillingSource]: source,
[TraceAttr.BillingCostUsd]: cost,
[TraceAttr.GenAiUsageInputTokens]: inputTokens,
[TraceAttr.GenAiUsageOutputTokens]: outputTokens,
[TraceAttr.BillingIsMcp]: isMcp,
...(idempotencyKey ? { [TraceAttr.BillingIdempotencyKey]: idempotencyKey } : {}),
})

claim = idempotencyKey
? await billingIdempotency.atomicallyClaim('update-cost', idempotencyKey)
: null
Expand All @@ -95,6 +135,8 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
userId,
source,
})
span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.DuplicateIdempotencyKey)
span.setAttribute(TraceAttr.HttpStatusCode, 409)
return NextResponse.json(
{
success: false,
Expand Down Expand Up @@ -159,6 +201,9 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
cost,
})

span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.Billed)
span.setAttribute(TraceAttr.HttpStatusCode, 200)
span.setAttribute(TraceAttr.BillingDurationMs, duration)
return NextResponse.json({
success: true,
data: {
Expand Down Expand Up @@ -193,6 +238,9 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
)
}

span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.InternalError)
span.setAttribute(TraceAttr.HttpStatusCode, 500)
span.setAttribute(TraceAttr.BillingDurationMs, duration)
return NextResponse.json(
{
success: false,
Expand All @@ -202,4 +250,4 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
{ status: 500 }
)
}
})
}
7 changes: 6 additions & 1 deletion apps/sim/app/api/copilot/api-keys/generate/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ import { type NextRequest, NextResponse } from 'next/server'
import { z } from 'zod'
import { getSession } from '@/lib/auth'
import { SIM_AGENT_API_URL } from '@/lib/copilot/constants'
import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
import { fetchGo } from '@/lib/copilot/request/go/fetch'
import { env } from '@/lib/core/config/env'
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'

Expand Down Expand Up @@ -33,13 +35,16 @@ export const POST = withRouteHandler(async (req: NextRequest) => {

const { name } = validationResult.data

const res = await fetch(`${SIM_AGENT_API_URL}/api/validate-key/generate`, {
const res = await fetchGo(`${SIM_AGENT_API_URL}/api/validate-key/generate`, {
method: 'POST',
headers: {
'Content-Type': 'application/json',
...(env.COPILOT_API_KEY ? { 'x-api-key': env.COPILOT_API_KEY } : {}),
},
body: JSON.stringify({ userId, name }),
spanName: 'sim → go /api/validate-key/generate',
operation: 'generate_api_key',
attributes: { [TraceAttr.UserId]: userId },
})

if (!res.ok) {
Expand Down
130 changes: 84 additions & 46 deletions apps/sim/app/api/copilot/api-keys/route.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,22 @@ vi.mock('@/lib/core/config/env', () => createEnvMock({ COPILOT_API_KEY: 'test-ap

import { DELETE, GET } from '@/app/api/copilot/api-keys/route'

// `fetchGo` reads `response.status` and `response.headers.get('content-length')`
// to stamp span attributes, so mock responses need both fields or the call
// path throws before the route handler sees the body.
function buildMockResponse(init: {
ok: boolean
status?: number
json: () => Promise<unknown>
}): Record<string, unknown> {
return {
ok: init.ok,
status: init.status ?? (init.ok ? 200 : 500),
headers: new Headers(),
json: init.json,
}
}

describe('Copilot API Keys API Route', () => {
beforeEach(() => {
vi.clearAllMocks()
Expand Down Expand Up @@ -60,10 +76,12 @@ describe('Copilot API Keys API Route', () => {
},
]

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.resolve(mockApiKeys),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.resolve(mockApiKeys),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys')
const response = await GET(request)
Expand All @@ -83,10 +101,12 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.resolve([]),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.resolve([]),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys')
const response = await GET(request)
Expand All @@ -101,10 +121,12 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.resolve([]),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.resolve([]),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys')
await GET(request)
Expand All @@ -127,11 +149,13 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: false,
status: 503,
json: () => Promise.resolve({ error: 'Service unavailable' }),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: false,
status: 503,
json: () => Promise.resolve({ error: 'Service unavailable' }),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys')
const response = await GET(request)
Expand All @@ -146,10 +170,12 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.resolve({ invalid: 'response' }),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.resolve({ invalid: 'response' }),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys')
const response = await GET(request)
Expand Down Expand Up @@ -189,10 +215,12 @@ describe('Copilot API Keys API Route', () => {
},
]

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.resolve(mockApiKeys),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.resolve(mockApiKeys),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys')
const response = await GET(request)
Expand All @@ -207,10 +235,12 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.reject(new Error('Invalid JSON')),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.reject(new Error('Invalid JSON')),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys')
const response = await GET(request)
Expand Down Expand Up @@ -251,10 +281,12 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.resolve({ success: true }),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.resolve({ success: true }),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys?id=key-123')
const response = await DELETE(request)
Expand All @@ -281,11 +313,13 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: false,
status: 404,
json: () => Promise.resolve({ error: 'Key not found' }),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: false,
status: 404,
json: () => Promise.resolve({ error: 'Key not found' }),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys?id=non-existent')
const response = await DELETE(request)
Expand All @@ -300,10 +334,12 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.resolve({ success: false }),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.resolve({ success: false }),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys?id=key-123')
const response = await DELETE(request)
Expand Down Expand Up @@ -333,10 +369,12 @@ describe('Copilot API Keys API Route', () => {
user: { id: 'user-123', email: 'test@example.com' },
})

mockFetch.mockResolvedValueOnce({
ok: true,
json: () => Promise.reject(new Error('Invalid JSON')),
})
mockFetch.mockResolvedValueOnce(
buildMockResponse({
ok: true,
json: () => Promise.reject(new Error('Invalid JSON')),
})
)

const request = new NextRequest('http://localhost:3000/api/copilot/api-keys?id=key-123')
const response = await DELETE(request)
Expand Down
Loading
Loading