diff --git a/apps/web/src/app/api/profile/cloud-agent-admission/route.test.ts b/apps/web/src/app/api/profile/cloud-agent-admission/route.test.ts new file mode 100644 index 0000000000..7fd0debbdf --- /dev/null +++ b/apps/web/src/app/api/profile/cloud-agent-admission/route.test.ts @@ -0,0 +1,136 @@ +import { beforeEach, describe, expect, test } from '@jest/globals'; +import { NextRequest, NextResponse } from 'next/server'; +import { getUserFromAuth } from '@/lib/user/server'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { classifyCloudAgentModelBilling } from '@/lib/cloud-agent-next/classify-model-billing'; +import { ORGANIZATION_ID_HEADER } from '@/lib/constants'; +import { POST } from './route'; + +jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() })); +jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() })); +jest.mock('@/lib/drizzle', () => ({ db: {} })); +jest.mock('@/lib/organizations/organization-usage', () => ({ + getBalanceAndOrgSettings: jest.fn(), +})); +jest.mock('@/lib/cloud-agent-next/classify-model-billing', () => ({ + classifyCloudAgentModelBilling: jest.fn(), +})); + +const mockedGetUserFromAuth = jest.mocked(getUserFromAuth); +const mockedGetBalance = jest.mocked(getBalanceAndOrgSettings); +const mockedClassify = jest.mocked(classifyCloudAgentModelBilling); + +function request(body: unknown, headers?: HeadersInit) { + return new NextRequest('http://localhost:3000/api/profile/cloud-agent-admission', { + method: 'POST', + headers, + body: JSON.stringify(body), + }); +} + +describe('POST /api/profile/cloud-agent-admission', () => { + beforeEach(() => { + jest.resetAllMocks(); + mockedGetUserFromAuth.mockResolvedValue({ + user: { id: 'user-1' }, + organizationId: undefined, + authFailedResponse: null, + } as never); + mockedClassify.mockResolvedValue('balance-required'); + mockedGetBalance.mockResolvedValue({ balance: 5 } as never); + }); + + test('does not read balance for a non-balance-required model', async () => { + mockedClassify.mockResolvedValue('byok'); + + const response = await POST(request({ modelId: 'anthropic/claude-sonnet-4' })); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual({ + classification: 'byok', + balance: null, + isDepleted: null, + }); + expect(mockedGetBalance).not.toHaveBeenCalled(); + }); + + test('resolves balance for a balance-required model', async () => { + mockedGetBalance.mockResolvedValue({ balance: 5 } as never); + + const response = await POST(request({ modelId: 'anthropic/claude-sonnet-4' })); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual({ + classification: 'balance-required', + balance: 5, + isDepleted: false, + }); + expect(mockedClassify).toHaveBeenCalledWith( + expect.objectContaining({ modelId: 'anthropic/claude-sonnet-4', userId: 'user-1' }) + ); + }); + + test('marks a depleted balance as depleted', async () => { + mockedGetBalance.mockResolvedValue({ balance: 0 } as never); + + const response = await POST(request({ modelId: 'anthropic/claude-sonnet-4' })); + + await expect(response.json()).resolves.toEqual({ + classification: 'balance-required', + balance: 0, + isDepleted: true, + }); + }); + + test('reads balance against the organization when membership is resolved', async () => { + mockedGetUserFromAuth.mockResolvedValue({ + user: { id: 'user-1' }, + organizationId: 'org-1', + authFailedResponse: null, + } as never); + + await POST( + request({ modelId: 'anthropic/claude-sonnet-4' }, { [ORGANIZATION_ID_HEADER]: 'org-1' }) + ); + + expect(mockedClassify).toHaveBeenCalledWith( + expect.objectContaining({ organizationId: 'org-1' }) + ); + expect(mockedGetBalance).toHaveBeenCalledWith( + 'org-1', + expect.objectContaining({ id: 'user-1' }) + ); + }); + + test('propagates the auth failure response without classifying', async () => { + mockedGetUserFromAuth.mockResolvedValue({ + user: null, + organizationId: undefined, + authFailedResponse: NextResponse.json({ error: 'Unauthorized' }, { status: 401 }), + } as never); + + const response = await POST(request({ modelId: 'anthropic/claude-sonnet-4' })); + + expect(response.status).toBe(401); + expect(mockedClassify).not.toHaveBeenCalled(); + }); + + test('rejects an invalid body without authenticating', async () => { + const response = await POST(request({ modelId: '' })); + + expect(response.status).toBe(400); + expect(mockedGetUserFromAuth).not.toHaveBeenCalled(); + expect(mockedClassify).not.toHaveBeenCalled(); + }); + + test('returns a service failure when classification throws', async () => { + mockedClassify.mockRejectedValue(new Error('catalog unavailable')); + + const response = await POST(request({ modelId: 'anthropic/claude-sonnet-4' })); + + expect(response.status).toBe(500); + await expect(response.json()).resolves.toEqual({ + error: 'Failed to check cloud agent admission', + }); + }); +}); diff --git a/apps/web/src/app/api/profile/cloud-agent-admission/route.ts b/apps/web/src/app/api/profile/cloud-agent-admission/route.ts new file mode 100644 index 0000000000..911ed6ffef --- /dev/null +++ b/apps/web/src/app/api/profile/cloud-agent-admission/route.ts @@ -0,0 +1,72 @@ +import type { NextRequest } from 'next/server'; +import { NextResponse } from 'next/server'; +import { captureException } from '@sentry/nextjs'; +import * as z from 'zod'; +import { db } from '@/lib/drizzle'; +import { getUserFromAuth } from '@/lib/user/server'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { + classifyCloudAgentModelBilling, + type CloudAgentModelBilling, +} from '@/lib/cloud-agent-next/classify-model-billing'; + +const BodySchema = z.object({ modelId: z.string().trim().min(1) }); + +type AdmissionResult = { + classification: CloudAgentModelBilling; + /** Owner balance, only resolved when the model is `balance-required`; `null` otherwise. */ + balance: number | null; + isDepleted: boolean | null; +}; + +type AdmissionResponse = { error: string } | AdmissionResult; + +/** + * Single admission check for a Cloud Agent mutation: classifies how the model + * would be billed and, when it is `balance-required`, resolves the owner's + * balance in the same round-trip. Backs the worker's balance middleware, which + * cannot import the model catalog or read balance directly. + * + * The organization owner comes from the `X-KiloCode-OrganizationId` header + * (membership validated by `getUserFromAuth`); otherwise the caller's user is + * the owner. Balance is only computed when needed, so `free`/`byok` sessions + * pay no extra cost. + */ +export async function POST(request: NextRequest): Promise> { + let body: unknown; + try { + body = await request.json(); + } catch { + return NextResponse.json({ error: 'Invalid JSON body' }, { status: 400 }); + } + + const bodyResult = BodySchema.safeParse(body); + if (!bodyResult.success) { + return NextResponse.json({ error: 'Invalid request body' }, { status: 400 }); + } + + const { user, authFailedResponse, organizationId } = await getUserFromAuth({ adminOnly: false }); + if (authFailedResponse) return authFailedResponse; + + try { + const classification = await classifyCloudAgentModelBilling({ + fromDb: db, + modelId: bodyResult.data.modelId, + userId: user.id, + organizationId, + }); + + if (classification !== 'balance-required') { + return NextResponse.json({ classification, balance: null, isDepleted: null }); + } + + const { balance } = await getBalanceAndOrgSettings(organizationId, user); + return NextResponse.json({ classification, balance, isDepleted: balance <= 0 }); + } catch (error) { + captureException(error, { + tags: { endpoint: 'profile/cloud-agent-admission' }, + extra: { userId: user.id, organizationId, modelId: bodyResult.data.modelId }, + }); + return NextResponse.json({ error: 'Failed to check cloud agent admission' }, { status: 500 }); + } +} diff --git a/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx b/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx index 953e8ddf0d..7f9ee7d93a 100644 --- a/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx +++ b/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx @@ -104,6 +104,10 @@ import { shouldCacheBitbucketRepositoryRefreshResult, shouldIncludeBitbucketRepositoryRefresh, } from '@/components/cloud-agent-next/repository-refresh'; +import { + filterModelsForCloudAgentBalance, + getNewSessionBalanceNotice, +} from './new-session-balance'; type Repository = { id: string | number; @@ -161,9 +165,8 @@ export function NewSessionPanel({ organizationId, isDevcontainerAvailable }: New const isEligibilityLoading = organizationId ? orgEligibilityQuery.isPending : personalEligibilityQuery.isPending; - const hasInsufficientBalance = - !isEligibilityLoading && eligibilityData && !eligibilityData.isEligible; - const hasLimitedAccess = !isEligibilityLoading && eligibilityData?.accessLevel === 'limited'; + const hasNoBalance = !isEligibilityLoading && eligibilityData?.accessLevel === 'limited'; + const hasLowBalance = !isEligibilityLoading && eligibilityData?.isLowBalance === true; // --------------------------------------------------------------------------- // Models @@ -185,9 +188,9 @@ export function NewSessionPanel({ organizationId, isDevcontainerAvailable }: New variants: model.opencode?.variants ? Object.keys(model.opencode.variants) : undefined, })); const withLocalTest = appendCloudAgentNextLocalTestModel(options); - if (!hasLimitedAccess) return withLocalTest; - return withLocalTest.filter(option => option.isFree || option.hasUserByokAvailable); - }, [allModels, hasLimitedAccess]); + return filterModelsForCloudAgentBalance(withLocalTest, hasNoBalance); + }, [allModels, hasNoBalance]); + const balanceNotice = getNewSessionBalanceNotice(hasNoBalance, hasLowBalance); // --------------------------------------------------------------------------- // Form state @@ -888,21 +891,11 @@ export function NewSessionPanel({ organizationId, isDevcontainerAvailable }: New // --------------------------------------------------------------------------- const isPromptTooLong = prompt.length > CLOUD_AGENT_PROMPT_MAX_LENGTH; - const selectedModelOption = modelOptions.find(m => m.id === model); - // Limited-access users can submit when they've picked a free or BYOK-capable - // model from the filtered picker; the server still gates paid models behind - // the minimum balance, but the submit button shouldn't pretend otherwise. - const limitedAccessModelIsAllowed = - hasLimitedAccess && - !!selectedModelOption && - (selectedModelOption.isFree || selectedModelOption.hasUserByokAvailable); - const isFormValid = prompt.trim().length > 0 && !isPromptTooLong && model.length > 0 && !isPreparing && - (!hasInsufficientBalance || limitedAccessModelIsAllowed) && !attachmentUpload.hasUploadingAttachments; const handleStartSession = useCallback(async () => { @@ -1182,27 +1175,31 @@ export function NewSessionPanel({ organizationId, isDevcontainerAvailable }: New
- {/* Insufficient balance banner */} - {hasInsufficientBalance && eligibilityData && !hasLimitedAccess && ( + {balanceNotice === 'zero-credit' && eligibilityData && ( )} - {/* Free-models-available banner when balance is low but free models are usable */} - {hasLimitedAccess && eligibilityData && ( + {balanceNotice === 'low-funds' && eligibilityData && ( )} diff --git a/apps/web/src/components/cloud-agent-next/new-session-balance.test.ts b/apps/web/src/components/cloud-agent-next/new-session-balance.test.ts new file mode 100644 index 0000000000..cfee2a0121 --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/new-session-balance.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from '@jest/globals'; +import { + filterModelsForCloudAgentBalance, + getNewSessionBalanceNotice, +} from './new-session-balance'; + +const models = [ + { id: 'paid', isFree: false, hasUserByokAvailable: false }, + { id: 'free', isFree: true, hasUserByokAvailable: false }, + { id: 'byok', isFree: false, hasUserByokAvailable: true }, +]; + +describe('NewSessionPanel balance presentation', () => { + it('filters a stale paid selection out at zero Credits while retaining free and BYOK models', () => { + const availableModels = filterModelsForCloudAgentBalance(models, true); + + expect(availableModels.map(model => model.id)).toEqual(['free', 'byok']); + expect(availableModels.some(model => model.id === 'paid')).toBe(false); + expect(getNewSessionBalanceNotice(true, true)).toBe('zero-credit'); + }); + + it('keeps every model available and shows only a low-funds notice for a positive sub-dollar balance', () => { + expect(filterModelsForCloudAgentBalance(models, false)).toEqual(models); + expect(getNewSessionBalanceNotice(false, true)).toBe('low-funds'); + }); + + it('shows no balance notice at one dollar or more', () => { + expect(filterModelsForCloudAgentBalance(models, false)).toEqual(models); + expect(getNewSessionBalanceNotice(false, false)).toBe('none'); + }); +}); diff --git a/apps/web/src/components/cloud-agent-next/new-session-balance.ts b/apps/web/src/components/cloud-agent-next/new-session-balance.ts new file mode 100644 index 0000000000..9458f55919 --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/new-session-balance.ts @@ -0,0 +1,22 @@ +export type BalanceRestrictedModelOption = { + isFree?: boolean; + hasUserByokAvailable?: boolean; +}; + +export type NewSessionBalanceNotice = 'none' | 'zero-credit' | 'low-funds'; + +export function filterModelsForCloudAgentBalance( + models: T[], + hasNoBalance: boolean +): T[] { + if (!hasNoBalance) return models; + return models.filter(model => model.isFree === true || model.hasUserByokAvailable === true); +} + +export function getNewSessionBalanceNotice( + hasNoBalance: boolean, + hasLowBalance: boolean +): NewSessionBalanceNotice { + if (hasNoBalance) return 'zero-credit'; + return hasLowBalance ? 'low-funds' : 'none'; +} diff --git a/apps/web/src/lib/app-builder/app-builder-service.ts b/apps/web/src/lib/app-builder/app-builder-service.ts index 779d4d2d56..609fc8007f 100644 --- a/apps/web/src/lib/app-builder/app-builder-service.ts +++ b/apps/web/src/lib/app-builder/app-builder-service.ts @@ -1,7 +1,7 @@ import 'server-only'; import type { Owner } from '@/lib/integrations/core/types'; import { - createAppBuilderCloudAgentNextClient, + createCloudAgentNextClient, type InterruptResult, type InitiateSessionOutput, } from '@/lib/cloud-agent-next/cloud-agent-client'; @@ -165,8 +165,7 @@ async function shouldCreateNewSession( } if (project.git_repo_full_name) { - const session = - await createAppBuilderCloudAgentNextClient(authToken).getSession(currentSessionId); + const session = await createCloudAgentNextClient(authToken).getSession(currentSessionId); if (session.gitUrl && !session.githubRepo) { return { @@ -261,7 +260,7 @@ async function createCloudAgentNextSession( images, reason, } = params; - const client = createAppBuilderCloudAgentNextClient(authToken); + const client = createCloudAgentNextClient(authToken); const augmentedMessage = message + buildImageContextFromAttachments(images); @@ -358,7 +357,7 @@ async function sendToExistingCloudAgentNextSession( gitToken = tokenResult.token; } - const client = createAppBuilderCloudAgentNextClient(authToken); + const client = createCloudAgentNextClient(authToken); const result = await client.sendMessage({ cloudAgentSessionId: sessionId, payload: { @@ -436,7 +435,7 @@ export async function createProject(input: CreateProjectInput): Promise { for (const field of profileDerivedInlineFields) { expect(prepareInput).not.toHaveProperty(field); } - expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('auth-token', { - skipBalanceCheck: true, - }); + expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('auth-token'); expect(mockInitiateFromPreparedSession).toHaveBeenCalledWith({ cloudAgentSessionId: 'cloud-session-1', }); diff --git a/apps/web/src/lib/bot/tools/spawn-cloud-agent-session.ts b/apps/web/src/lib/bot/tools/spawn-cloud-agent-session.ts index b7caa7814e..c11d8e190d 100644 --- a/apps/web/src/lib/bot/tools/spawn-cloud-agent-session.ts +++ b/apps/web/src/lib/bot/tools/spawn-cloud-agent-session.ts @@ -203,7 +203,7 @@ export default async function spawnCloudAgentSession( }; } - const client = createCloudAgentNextClient(authToken, { skipBalanceCheck: true }); + const client = createCloudAgentNextClient(authToken); let cloudAgentSessionId: string; let kiloSessionId: string; diff --git a/apps/web/src/lib/cloud-agent-next/balance-check-eligibility.test.ts b/apps/web/src/lib/cloud-agent-next/balance-check-eligibility.test.ts deleted file mode 100644 index 0f4dcdba23..0000000000 --- a/apps/web/src/lib/cloud-agent-next/balance-check-eligibility.test.ts +++ /dev/null @@ -1,143 +0,0 @@ -const mockIsFreeModel = jest.fn(); -const mockGetModelUserByokProviders = jest.fn(); -const mockGetUserByokProviderIds = jest.fn(); -const mockGetOrganizationByokProviderIds = jest.fn(); - -jest.mock('@/lib/ai-gateway/is-free-model', () => ({ - isFreeModel: (...args: unknown[]) => mockIsFreeModel(...args), -})); - -jest.mock('@/lib/ai-gateway/byok', () => ({ - getModelUserByokProviders: (...args: unknown[]) => mockGetModelUserByokProviders(...args), - getUserByokProviderIds: (...args: unknown[]) => mockGetUserByokProviderIds(...args), - getOrganizationByokProviderIds: (...args: unknown[]) => - mockGetOrganizationByokProviderIds(...args), -})); - -import { computeCloudAgentNextBalanceCheckEligibility } from './balance-check-eligibility'; - -const KILO_EXCLUSIVE_MODEL = 'deepseek/deepseek-v4-pro:discounted'; -const NON_EXCLUSIVE_MODEL = 'anthropic/claude-sonnet-4'; - -const fakeDb = {} as never; -const fakeUser = { id: 'user-1' }; - -beforeEach(() => { - jest.resetAllMocks(); - mockIsFreeModel.mockResolvedValue(false); - mockGetModelUserByokProviders.mockResolvedValue([]); - mockGetUserByokProviderIds.mockResolvedValue([]); - mockGetOrganizationByokProviderIds.mockResolvedValue([]); -}); - -describe('computeCloudAgentNextBalanceCheckEligibility', () => { - it('returns isFree and skips BYOK when the model is free', async () => { - mockIsFreeModel.mockResolvedValueOnce(true); - - const result = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: fakeDb, - user: fakeUser, - modelId: 'kilo/free-model', - }); - - expect(result).toEqual({ isFree: true, hasUserByokAvailable: false }); - expect(mockGetModelUserByokProviders).not.toHaveBeenCalled(); - }); - - it('returns hasUserByokAvailable: false for a Kilo-exclusive model even when BYOK providers can serve it', async () => { - const result = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: fakeDb, - user: fakeUser, - modelId: KILO_EXCLUSIVE_MODEL, - }); - - expect(result).toEqual({ isFree: false, hasUserByokAvailable: false }); - expect(mockGetModelUserByokProviders).not.toHaveBeenCalled(); - expect(mockGetUserByokProviderIds).not.toHaveBeenCalled(); - expect(mockGetOrganizationByokProviderIds).not.toHaveBeenCalled(); - }); - - it('returns hasUserByokAvailable: false for a Kilo-exclusive model even when the user has an enabled matching BYOK provider', async () => { - mockGetUserByokProviderIds.mockResolvedValueOnce(['openrouter']); - - const result = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: fakeDb, - user: fakeUser, - modelId: KILO_EXCLUSIVE_MODEL, - }); - - expect(result).toEqual({ isFree: false, hasUserByokAvailable: false }); - expect(mockGetModelUserByokProviders).not.toHaveBeenCalled(); - expect(mockGetUserByokProviderIds).not.toHaveBeenCalled(); - }); - - it('returns hasUserByokAvailable: false for a Kilo-exclusive model in an organization context', async () => { - mockGetOrganizationByokProviderIds.mockResolvedValueOnce(['openrouter']); - - const result = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: fakeDb, - user: fakeUser, - modelId: KILO_EXCLUSIVE_MODEL, - organizationId: 'org-1', - }); - - expect(result).toEqual({ isFree: false, hasUserByokAvailable: false }); - expect(mockGetModelUserByokProviders).not.toHaveBeenCalled(); - expect(mockGetOrganizationByokProviderIds).not.toHaveBeenCalled(); - }); - - it('returns hasUserByokAvailable: true for a non-Kilo-exclusive paid model with a matching enabled user BYOK provider', async () => { - mockGetModelUserByokProviders.mockResolvedValueOnce(['openrouter']); - mockGetUserByokProviderIds.mockResolvedValueOnce(['openrouter']); - - const result = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: fakeDb, - user: fakeUser, - modelId: NON_EXCLUSIVE_MODEL, - }); - - expect(result).toEqual({ isFree: false, hasUserByokAvailable: true }); - }); - - it('returns hasUserByokAvailable: false for a non-Kilo-exclusive paid model with no matching BYOK provider', async () => { - mockGetModelUserByokProviders.mockResolvedValueOnce(['openrouter']); - mockGetUserByokProviderIds.mockResolvedValueOnce(['anthropic']); - - const result = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: fakeDb, - user: fakeUser, - modelId: NON_EXCLUSIVE_MODEL, - }); - - expect(result).toEqual({ isFree: false, hasUserByokAvailable: false }); - }); - - it('returns hasUserByokAvailable: false for a non-Kilo-exclusive paid model with no resolvable providers', async () => { - mockGetModelUserByokProviders.mockResolvedValueOnce([]); - - const result = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: fakeDb, - user: fakeUser, - modelId: NON_EXCLUSIVE_MODEL, - }); - - expect(result).toEqual({ isFree: false, hasUserByokAvailable: false }); - expect(mockGetUserByokProviderIds).not.toHaveBeenCalled(); - }); - - it('uses organization BYOK providers for a non-Kilo-exclusive paid model when organizationId is provided', async () => { - mockGetModelUserByokProviders.mockResolvedValueOnce(['openrouter']); - mockGetOrganizationByokProviderIds.mockResolvedValueOnce(['openrouter']); - - const result = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: fakeDb, - user: fakeUser, - modelId: NON_EXCLUSIVE_MODEL, - organizationId: 'org-1', - }); - - expect(result).toEqual({ isFree: false, hasUserByokAvailable: true }); - expect(mockGetOrganizationByokProviderIds).toHaveBeenCalledWith(fakeDb, 'org-1'); - expect(mockGetUserByokProviderIds).not.toHaveBeenCalled(); - }); -}); diff --git a/apps/web/src/lib/cloud-agent-next/balance-check-eligibility.ts b/apps/web/src/lib/cloud-agent-next/balance-check-eligibility.ts deleted file mode 100644 index fd3d6c0ac5..0000000000 --- a/apps/web/src/lib/cloud-agent-next/balance-check-eligibility.ts +++ /dev/null @@ -1,66 +0,0 @@ -import 'server-only'; -import { type db } from '@/lib/drizzle'; -import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; -import { isKiloExclusiveModel } from '@/lib/ai-gateway/models'; -import { - getModelUserByokProviders, - getOrganizationByokProviderIds, - getUserByokProviderIds, -} from '@/lib/ai-gateway/byok'; -import type { User } from '@kilocode/db/schema'; - -export type BalanceCheckModelEligibility = { - isFree: boolean; - hasUserByokAvailable: boolean; -}; - -/** - * Decide whether `prepareSession` should skip the worker-side $1 balance - * minimum for the chosen model. - * - * Skips the check when either: - * - the model is Kilo-funded (free for the user), or - * - the model is not Kilo-exclusive AND the user has a BYOK provider - * configured that can serve it, so the session is billed against the - * user's own key rather than their balance. - * - * Kilo-exclusive models (e.g. `deepseek/deepseek-v4-pro:discounted`) are - * always excluded from the BYOK bypass: they are Kilo-funded and platform - * billed, so even when `getModelUserByokProviders` reports a provider that - * can route the model, they must still go through the worker-side balance - * check and cannot be legitimately served via a user's own BYOK key. - * - * The resulting predicate is a strict subset of the - * `isFree || hasUserByokAvailable` predicate used by the NewSessionPanel - * model picker to filter `hasLimitedAccess` users: this router additionally - * forces Kilo-exclusive models through the balance check, so the picker - * may offer a model as free while the router still requires a balance. - */ -export async function computeCloudAgentNextBalanceCheckEligibility(params: { - fromDb: typeof db; - user: Pick; - modelId: string; - organizationId?: string; -}): Promise { - const isFree = await isFreeModel(params.modelId); - if (isFree) { - return { isFree: true, hasUserByokAvailable: false }; - } - - if (isKiloExclusiveModel(params.modelId)) { - return { isFree: false, hasUserByokAvailable: false }; - } - - const modelProviders = await getModelUserByokProviders(params.modelId); - if (modelProviders.length === 0) { - return { isFree: false, hasUserByokAvailable: false }; - } - - const enabledProviderIds = params.organizationId - ? await getOrganizationByokProviderIds(params.fromDb, params.organizationId) - : await getUserByokProviderIds(params.fromDb, params.user.id); - - const enabled = new Set(enabledProviderIds); - const hasUserByokAvailable = modelProviders.some(provider => enabled.has(provider)); - return { isFree: false, hasUserByokAvailable }; -} diff --git a/apps/web/src/lib/cloud-agent-next/classify-model-billing.test.ts b/apps/web/src/lib/cloud-agent-next/classify-model-billing.test.ts new file mode 100644 index 0000000000..f93d51a997 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/classify-model-billing.test.ts @@ -0,0 +1,139 @@ +const mockIsFreeModel = jest.fn(); +const mockGetModelUserByokProviders = jest.fn(); +const mockGetUserByokProviderIds = jest.fn(); +const mockGetOrganizationByokProviderIds = jest.fn(); +const mockGetDirectByokModel = jest.fn(); + +jest.mock('@/lib/ai-gateway/is-free-model', () => ({ + isFreeModel: (...args: unknown[]) => mockIsFreeModel(...args), +})); + +jest.mock('@/lib/ai-gateway/byok', () => ({ + getModelUserByokProviders: (...args: unknown[]) => mockGetModelUserByokProviders(...args), + getUserByokProviderIds: (...args: unknown[]) => mockGetUserByokProviderIds(...args), + getOrganizationByokProviderIds: (...args: unknown[]) => + mockGetOrganizationByokProviderIds(...args), +})); + +jest.mock('@/lib/ai-gateway/providers/direct-byok', () => ({ + getDirectByokModel: (...args: unknown[]) => mockGetDirectByokModel(...args), +})); + +import { classifyCloudAgentModelBilling } from './classify-model-billing'; + +// A real Kilo-exclusive model id so the (unmocked) catalog lookup resolves. +const KILO_EXCLUSIVE_MODEL = 'deepseek/deepseek-v4-pro:discounted'; +const NON_EXCLUSIVE_MODEL = 'anthropic/claude-sonnet-4'; + +const fakeDb = {} as never; + +beforeEach(() => { + jest.resetAllMocks(); + mockIsFreeModel.mockResolvedValue(false); + mockGetDirectByokModel.mockResolvedValue({ provider: null, model: null }); + mockGetModelUserByokProviders.mockResolvedValue([]); + mockGetUserByokProviderIds.mockResolvedValue([]); + mockGetOrganizationByokProviderIds.mockResolvedValue([]); +}); + +describe('classifyCloudAgentModelBilling', () => { + it('classifies a free model as free without reading BYOK configuration', async () => { + mockIsFreeModel.mockResolvedValueOnce(true); + + const result = await classifyCloudAgentModelBilling({ + fromDb: fakeDb, + userId: 'user-1', + modelId: 'kilo/free-model', + }); + + expect(result).toBe('free'); + expect(mockGetModelUserByokProviders).not.toHaveBeenCalled(); + }); + + it('classifies a Kilo-exclusive model as balance-required even when a BYOK provider could serve it', async () => { + mockGetUserByokProviderIds.mockResolvedValueOnce(['openrouter']); + + const result = await classifyCloudAgentModelBilling({ + fromDb: fakeDb, + userId: 'user-1', + modelId: KILO_EXCLUSIVE_MODEL, + }); + + expect(result).toBe('balance-required'); + expect(mockGetModelUserByokProviders).not.toHaveBeenCalled(); + expect(mockGetUserByokProviderIds).not.toHaveBeenCalled(); + expect(mockGetOrganizationByokProviderIds).not.toHaveBeenCalled(); + }); + + it('classifies a paid model as byok when the user has a matching enabled provider', async () => { + mockGetModelUserByokProviders.mockResolvedValueOnce(['openrouter']); + mockGetUserByokProviderIds.mockResolvedValueOnce(['openrouter']); + + const result = await classifyCloudAgentModelBilling({ + fromDb: fakeDb, + userId: 'user-1', + modelId: NON_EXCLUSIVE_MODEL, + }); + + expect(result).toBe('byok'); + }); + + it('classifies a direct BYOK model as byok when the user has that provider enabled', async () => { + mockGetDirectByokModel.mockResolvedValueOnce({ + provider: { id: 'chutes-byok' }, + model: { id: 'supported-model' }, + }); + mockGetUserByokProviderIds.mockResolvedValueOnce(['chutes-byok']); + + const result = await classifyCloudAgentModelBilling({ + fromDb: fakeDb, + userId: 'user-1', + modelId: 'chutes-byok/supported-model', + }); + + expect(result).toBe('byok'); + expect(mockGetModelUserByokProviders).not.toHaveBeenCalled(); + }); + + it('classifies a paid model as balance-required when no enabled provider matches', async () => { + mockGetModelUserByokProviders.mockResolvedValueOnce(['openrouter']); + mockGetUserByokProviderIds.mockResolvedValueOnce(['anthropic']); + + const result = await classifyCloudAgentModelBilling({ + fromDb: fakeDb, + userId: 'user-1', + modelId: NON_EXCLUSIVE_MODEL, + }); + + expect(result).toBe('balance-required'); + }); + + it('classifies a paid model as balance-required when no provider can route it', async () => { + mockGetModelUserByokProviders.mockResolvedValueOnce([]); + + const result = await classifyCloudAgentModelBilling({ + fromDb: fakeDb, + userId: 'user-1', + modelId: NON_EXCLUSIVE_MODEL, + }); + + expect(result).toBe('balance-required'); + expect(mockGetUserByokProviderIds).not.toHaveBeenCalled(); + }); + + it('uses organization BYOK providers when organizationId is provided', async () => { + mockGetModelUserByokProviders.mockResolvedValueOnce(['openrouter']); + mockGetOrganizationByokProviderIds.mockResolvedValueOnce(['openrouter']); + + const result = await classifyCloudAgentModelBilling({ + fromDb: fakeDb, + userId: 'user-1', + modelId: NON_EXCLUSIVE_MODEL, + organizationId: 'org-1', + }); + + expect(result).toBe('byok'); + expect(mockGetOrganizationByokProviderIds).toHaveBeenCalledWith(fakeDb, 'org-1'); + expect(mockGetUserByokProviderIds).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/lib/cloud-agent-next/classify-model-billing.ts b/apps/web/src/lib/cloud-agent-next/classify-model-billing.ts new file mode 100644 index 0000000000..182328b251 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/classify-model-billing.ts @@ -0,0 +1,74 @@ +import 'server-only'; +import { type db } from '@/lib/drizzle'; +import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; +import { isKiloExclusiveModel } from '@/lib/ai-gateway/models'; +import { + getModelUserByokProviders, + getOrganizationByokProviderIds, + getUserByokProviderIds, +} from '@/lib/ai-gateway/byok'; +import { getDirectByokModel } from '@/lib/ai-gateway/providers/direct-byok'; +import type { UserByokProviderId } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id'; + +/** + * How a Cloud Agent session for a given model should be billed: + * - `free`: Kilo-funded, no balance required. + * - `byok`: served against the owner's own provider key, no balance required. + * - `balance-required`: platform-billed, needs a positive credit balance. + */ +export type CloudAgentModelBilling = 'free' | 'byok' | 'balance-required'; + +/** + * Classifies how a model would be billed for a given owner (a user, or an + * organization when `organizationId` is set). + * + * This is the source-of-truth classifier that backs the worker-side balance + * middleware: the cloud-agent-next worker runs in a Cloudflare Worker that + * cannot import `apps/web`'s model catalog, so it asks this module over HTTP + * (`POST /api/profile/cloud-agent-admission`) rather than re-deriving the catalog. + * + * Kilo-exclusive models (e.g. `deepseek/deepseek-v4-pro:discounted`) are always + * `balance-required`: they are Kilo-funded and platform billed, so even when + * `getModelUserByokProviders` reports a provider that could route the model, + * they cannot be legitimately served via an owner's own BYOK key. + */ +export async function classifyCloudAgentModelBilling(params: { + fromDb: typeof db; + modelId: string; + userId: string; + organizationId?: string; +}): Promise { + if (await isFreeModel(params.modelId)) { + return 'free'; + } + + if (isKiloExclusiveModel(params.modelId)) { + return 'balance-required'; + } + + const directByok = await getDirectByokModel(params.modelId); + if (directByok.provider && directByok.model) { + const enabledProviderIds = await getEnabledProviderIds(params); + return enabledProviderIds.includes(directByok.provider.id) ? 'byok' : 'balance-required'; + } + + const modelProviders = await getModelUserByokProviders(params.modelId); + if (modelProviders.length === 0) { + return 'balance-required'; + } + + const enabledProviderIds = await getEnabledProviderIds(params); + + const enabled = new Set(enabledProviderIds); + return modelProviders.some(provider => enabled.has(provider)) ? 'byok' : 'balance-required'; +} + +function getEnabledProviderIds(params: { + fromDb: typeof db; + userId: string; + organizationId?: string; +}): Promise { + return params.organizationId + ? getOrganizationByokProviderIds(params.fromDb, params.organizationId) + : getUserByokProviderIds(params.fromDb, params.userId); +} diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts index 6e046ac84f..b8c074fcda 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts @@ -1,89 +1,51 @@ -import { describe, expect, it, jest, beforeEach } from '@jest/globals'; +import { beforeAll, describe, expect, it, jest } from '@jest/globals'; -jest.mock('@/lib/dotenvx', () => ({ - getEnvVariable: jest.fn(() => 'http://cloud-agent-next'), -})); +import type * as ClientModule from './cloud-agent-client'; -jest.mock('@/lib/config.server', () => ({ - INTERNAL_API_SECRET: 'test-secret', -})); +const createTRPCClient = jest.fn(); +const httpLink = jest.fn(); jest.mock('@trpc/client', () => ({ - createTRPCClient: jest.fn(() => ({})), - httpLink: jest.fn(), + createTRPCClient, + httpLink, TRPCClientError: class TRPCClientError extends Error {}, })); -jest.mock('@sentry/nextjs', () => ({ - captureException: jest.fn(), -})); - -jest.mock('./cloud-agent-client', () => { - const createCloudAgentNextClient = jest.fn((_token: string) => ({ marker: 'default' })); - const createAppBuilderCloudAgentNextClient = jest.fn((_token: string) => ({ - marker: 'appbuilder', - })); - return { - createCloudAgentNextClient, - createAppBuilderCloudAgentNextClient, - createCloudAgentNextClientForModel: jest.fn( - (token: string, model: { isFree: boolean; hasUserByokAvailable: boolean }) => - model.isFree || model.hasUserByokAvailable - ? createAppBuilderCloudAgentNextClient(token) - : createCloudAgentNextClient(token) - ), - rethrowAsPaymentRequired: jest.fn(), - }; -}); - -const clientModule: { - createCloudAgentNextClient: jest.Mock; - createAppBuilderCloudAgentNextClient: jest.Mock; - createCloudAgentNextClientForModel: ( - token: string, - model: { isFree: boolean; hasUserByokAvailable: boolean } - ) => unknown; -} = jest.requireMock('./cloud-agent-client'); +jest.mock('@/lib/dotenvx', () => ({ getEnvVariable: jest.fn(() => 'https://agent.example.com') })); +jest.mock('@/lib/config.server', () => ({ INTERNAL_API_SECRET: 'internal-secret' })); +jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() })); -const { - createCloudAgentNextClient: mockCreateCloudAgentNextClient, - createAppBuilderCloudAgentNextClient: mockCreateAppBuilderCloudAgentNextClient, - createCloudAgentNextClientForModel, -} = clientModule; +let clientModule: typeof ClientModule; -beforeEach(() => { - mockCreateCloudAgentNextClient.mockClear(); - mockCreateAppBuilderCloudAgentNextClient.mockClear(); +beforeAll(async () => { + clientModule = await import('./cloud-agent-client'); }); -describe('createCloudAgentNextClientForModel', () => { - it('returns the default client when the model is paid and has no BYOK', () => { - const result = createCloudAgentNextClientForModel('token', { - isFree: false, - hasUserByokAvailable: false, - }); - expect(result).toEqual({ marker: 'default' }); - expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('token'); - expect(mockCreateAppBuilderCloudAgentNextClient).not.toHaveBeenCalled(); - }); +describe('CloudAgentNextClient', () => { + it('uses ordinary authentication headers without a balance bypass header', () => { + httpLink.mockReturnValue({}); + createTRPCClient.mockReturnValue({}); - it('returns the AppBuilder client when the model is free', () => { - const result = createCloudAgentNextClientForModel('token', { - isFree: true, - hasUserByokAvailable: false, + new clientModule.CloudAgentNextClient('user-token'); + + const options = httpLink.mock.calls[0]?.[0] as { headers: () => Record }; + expect(options.headers()).toEqual({ + Authorization: 'Bearer user-token', + 'x-internal-api-key': 'internal-secret', }); - expect(result).toEqual({ marker: 'appbuilder' }); - expect(mockCreateAppBuilderCloudAgentNextClient).toHaveBeenCalledWith('token'); - expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + expect(options.headers()).not.toHaveProperty('x-skip-balance-check'); }); - it('returns the AppBuilder client when the model is BYOK-capable, even if it is not free', () => { - const result = createCloudAgentNextClientForModel('token', { - isFree: false, - hasUserByokAvailable: true, - }); - expect(result).toEqual({ marker: 'appbuilder' }); - expect(mockCreateAppBuilderCloudAgentNextClient).toHaveBeenCalledWith('token'); - expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + it('normalizes an insufficient-credit error for web router procedures', () => { + try { + clientModule.rethrowAsPaymentRequired(new clientModule.InsufficientCreditsError()); + } catch (error) { + expect(error).toMatchObject({ + code: 'PAYMENT_REQUIRED', + message: 'Insufficient credits: a positive credit balance is required', + }); + return; + } + throw new Error('Expected a PAYMENT_REQUIRED error'); }); }); diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts index e04c84fba0..71101326c6 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts @@ -367,7 +367,7 @@ export class InsufficientCreditsError extends Error { readonly httpStatus = 402; readonly code = 'PAYMENT_REQUIRED'; - constructor(message = 'Insufficient credits: $1 minimum required') { + constructor(message = 'Insufficient credits: a positive credit balance is required') { super(message); this.name = 'InsufficientCreditsError'; } @@ -470,16 +470,6 @@ type CloudAgentNextTRPCClient = { }; }; -/** - * Options for configuring the cloud agent client - */ -export type CloudAgentNextClientOptions = { - /** - * Skip balance validation in cloud-agent (used by App Builder which handles its own billing). - */ - skipBalanceCheck?: boolean; -}; - /** * Client for communicating with the cloud-agent-next TRPC API * @@ -489,19 +479,14 @@ export type CloudAgentNextClientOptions = { export class CloudAgentNextClient { private client: CloudAgentNextTRPCClient; private authToken: string; - private options: CloudAgentNextClientOptions; - constructor(authToken: string, options: CloudAgentNextClientOptions = {}) { + constructor(authToken: string) { this.authToken = authToken; - this.options = options; // Build common headers const baseHeaders: Record = { Authorization: `Bearer ${this.authToken}`, }; - if (this.options.skipBalanceCheck) { - baseHeaders['x-skip-balance-check'] = 'true'; - } // Create TRPC client - only uses httpLink for mutations/queries // (streaming is handled via WebSocketManager) @@ -794,43 +779,6 @@ export class CloudAgentNextClient { /** * Create a cloud agent next client instance with the provided auth token */ -export function createCloudAgentNextClient( - authToken: string, - options?: CloudAgentNextClientOptions -): CloudAgentNextClient { - return new CloudAgentNextClient(authToken, options); -} - -/** - * Create a cloud agent next client instance configured for App Builder. - */ -export function createAppBuilderCloudAgentNextClient(authToken: string): CloudAgentNextClient { - return new CloudAgentNextClient(authToken, { - skipBalanceCheck: true, - }); -} - -/** - * Pick a Cloud Agent Next client for a session start based on whether the - * chosen model is free or BYOK-billable for the caller. - * - * Free models (and BYOK-capable models, where the user provides their own key - * and the model is therefore not billed against their balance) cost the user - * nothing, so we mirror AppBuilder and set `x-skip-balance-check: true` to - * allow $0-balance users to still create sessions when their selected model - * is free. Paid models stay on the default client so the worker balance - * middleware can still enforce the $1 minimum on paid-model sessions. - * - * The decision is made from caller-supplied booleans rather than re-querying - * model metadata, so the helper can be unit-tested without a database and - * matches the values the NewSessionPanel model picker already filters on. - */ -export function createCloudAgentNextClientForModel( - authToken: string, - model: { isFree: boolean; hasUserByokAvailable: boolean } -): CloudAgentNextClient { - if (model.isFree || model.hasUserByokAvailable) { - return createAppBuilderCloudAgentNextClient(authToken); - } - return createCloudAgentNextClient(authToken); +export function createCloudAgentNextClient(authToken: string): CloudAgentNextClient { + return new CloudAgentNextClient(authToken); } diff --git a/apps/web/src/lib/code-reviews/dispatch/dispatch-pending-reviews.test.ts b/apps/web/src/lib/code-reviews/dispatch/dispatch-pending-reviews.test.ts index 9b1db2310a..357c562d28 100644 --- a/apps/web/src/lib/code-reviews/dispatch/dispatch-pending-reviews.test.ts +++ b/apps/web/src/lib/code-reviews/dispatch/dispatch-pending-reviews.test.ts @@ -297,7 +297,10 @@ describe('tryDispatchPendingReviews', () => { expect.objectContaining({ reviewId: review.id, platform: 'bitbucket' }) ); expect(mockDispatchReview).toHaveBeenCalledWith( - expect.objectContaining({ reviewId: review.id, skipBalanceCheck: true }) + expect.objectContaining({ reviewId: review.id }) + ); + expect(mockDispatchReview).not.toHaveBeenCalledWith( + expect.objectContaining({ skipBalanceCheck: expect.anything() }) ); }); diff --git a/apps/web/src/lib/code-reviews/dispatch/dispatch-pending-reviews.ts b/apps/web/src/lib/code-reviews/dispatch/dispatch-pending-reviews.ts index 54eec4f515..fdc27c4bf0 100644 --- a/apps/web/src/lib/code-reviews/dispatch/dispatch-pending-reviews.ts +++ b/apps/web/src/lib/code-reviews/dispatch/dispatch-pending-reviews.ts @@ -575,7 +575,6 @@ async function dispatchReservedReview(reservation: ReservedReview, owner: Owner) await codeReviewWorkerClient.dispatchReview({ ...dispatchPayload, attemptId: attempt.id, - skipBalanceCheck: true, }); } catch (dispatchError) { errorExceptInTest('[dispatchReview] Worker dispatch failed, leaving review queued', { diff --git a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts index 7ec4496477..185c3edc29 100644 --- a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts +++ b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts @@ -128,7 +128,6 @@ export type CodeReviewPayload = { authToken: string; sessionInput: SessionInput; owner: Owner; - skipBalanceCheck?: boolean; /** Cloud-agent session ID from a previous completed review, for session continuation */ previousCloudAgentSessionId?: string; /** Provider-reported repository storage size, formatted for log correlation. */ diff --git a/apps/web/src/lib/discord-bot.ts b/apps/web/src/lib/discord-bot.ts index 2f1ab929b2..6ef71b3cd1 100644 --- a/apps/web/src/lib/discord-bot.ts +++ b/apps/web/src/lib/discord-bot.ts @@ -183,7 +183,7 @@ async function spawnCloudAgentSession( : args.prompt; const result = await runSessionToCompletion({ - client: createCloudAgentNextClient(authToken, { skipBalanceCheck: true }), + client: createCloudAgentNextClient(authToken), prepareInput: { githubRepo: args.githubRepo, prompt: promptWithSignature, diff --git a/apps/web/src/routers/cloud-agent-next-eligibility.ts b/apps/web/src/routers/cloud-agent-next-eligibility.ts index 6d705bcf64..2665dcecb2 100644 --- a/apps/web/src/routers/cloud-agent-next-eligibility.ts +++ b/apps/web/src/routers/cloud-agent-next-eligibility.ts @@ -1,15 +1,22 @@ -import { - buildAccessLevelEligibility, - type AccessLevel, - type AccessLevelEligibility, -} from '@/lib/access-level-eligibility'; - export const CLOUD_AGENT_NEXT_MIN_BALANCE_DOLLARS = 1; -export type CloudAgentNextAccessLevel = AccessLevel; +export type CloudAgentNextAccessLevel = 'full' | 'limited'; -export type CloudAgentNextEligibility = AccessLevelEligibility; +export type CloudAgentNextEligibility = { + balance: number; + minBalance: number; + accessLevel: CloudAgentNextAccessLevel; + isEligible: boolean; + isLowBalance: boolean; +}; export function buildCloudAgentNextEligibility(balance: number): CloudAgentNextEligibility { - return buildAccessLevelEligibility(balance, CLOUD_AGENT_NEXT_MIN_BALANCE_DOLLARS); + const hasNoBalance = balance <= 0; + return { + balance, + minBalance: CLOUD_AGENT_NEXT_MIN_BALANCE_DOLLARS, + accessLevel: hasNoBalance ? 'limited' : 'full', + isEligible: !hasNoBalance, + isLowBalance: balance < CLOUD_AGENT_NEXT_MIN_BALANCE_DOLLARS, + }; } diff --git a/apps/web/src/routers/cloud-agent-next-router.test.ts b/apps/web/src/routers/cloud-agent-next-router.test.ts index 97fcc6a793..41de305761 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -42,24 +42,13 @@ const mockGenerateCloudAgentAttachmentUploadUrl = jest.fn< }) => Promise<{ signedUrl: string; key: string; expiresAt: string }> >(() => Promise.resolve({ signedUrl: 'signed', key: 'key', expiresAt: 'expires' })); -const mockCreateCloudAgentNextClient = jest.fn(() => ({ - prepareSession: mockPrepareSession, - sendMessage: mockSendMessage, -})); - -const mockCreateCloudAgentNextClientForModel = jest.fn( - (_authToken: string, _eligibility: unknown) => ({ +const mockCreateCloudAgentNextClient = jest.fn((authToken: string) => { + void authToken; + return { prepareSession: mockPrepareSession, sendMessage: mockSendMessage, - }) -); - -const mockComputeCloudAgentNextBalanceCheckEligibility = jest.fn< - (...args: unknown[]) => Promise<{ - isFree: boolean; - hasUserByokAvailable: boolean; - }> ->(); + }; +}); const mockIsFeatureFlagEnabledOrDevelopment = jest.fn<(flagName: string, distinctId: string) => Promise>(); @@ -93,14 +82,9 @@ jest.mock('@/lib/tokens', () => ({ jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ createCloudAgentNextClient: mockCreateCloudAgentNextClient, - createCloudAgentNextClientForModel: mockCreateCloudAgentNextClientForModel, rethrowAsPaymentRequired: jest.fn(), })); -jest.mock('@/lib/cloud-agent-next/balance-check-eligibility', () => ({ - computeCloudAgentNextBalanceCheckEligibility: mockComputeCloudAgentNextBalanceCheckEligibility, -})); - jest.mock('@/lib/posthog-feature-flags', () => ({ isFeatureFlagEnabledOrDevelopment: mockIsFeatureFlagEnabledOrDevelopment, })); @@ -162,7 +146,8 @@ let createCaller: (ctx: { user: User }) => { balance: number; minBalance: number; isEligible: boolean; - accessLevel: 'full' | 'limited' | 'blocked'; + accessLevel: 'full' | 'limited'; + isLowBalance: boolean; }>; listGitHubRepositories: (input: { forceRefresh: boolean }) => Promise; listGitLabRepositories: (input: { forceRefresh: boolean }) => Promise; @@ -253,22 +238,27 @@ describe('cloudAgentNextRouter helper procedures', () => { }); it.each([ - { balance: 1, isEligible: true, accessLevel: 'full' as const }, - { balance: 0.99, isEligible: false, accessLevel: 'limited' as const }, - ])('reports eligibility for a $balance balance', async ({ balance, isEligible, accessLevel }) => { - mockGetBalanceForUser.mockResolvedValue({ balance }); - const user = { id: 'user-eligibility', is_admin: false } as User; - const caller = createCaller({ user }); - - await expect(caller.checkEligibility()).resolves.toEqual({ - balance, - minBalance: 1, - isEligible, - accessLevel, - }); - expect(mockGetBalanceForUser).toHaveBeenCalledWith(user); - expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); - }); + { balance: 0, isEligible: false, accessLevel: 'limited' as const, isLowBalance: true }, + { balance: 0.01, isEligible: true, accessLevel: 'full' as const, isLowBalance: true }, + { balance: 1, isEligible: true, accessLevel: 'full' as const, isLowBalance: false }, + ])( + 'reports eligibility for a $balance balance', + async ({ balance, isEligible, accessLevel, isLowBalance }) => { + mockGetBalanceForUser.mockResolvedValue({ balance }); + const user = { id: 'user-eligibility', is_admin: false } as User; + const caller = createCaller({ user }); + + await expect(caller.checkEligibility()).resolves.toEqual({ + balance, + minBalance: 1, + isEligible, + accessLevel, + isLowBalance, + }); + expect(mockGetBalanceForUser).toHaveBeenCalledWith(user); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + } + ); it.each([ ['GitHub', 'listGitHubRepositories', mockFetchGitHubRepositoriesForUser], @@ -298,10 +288,6 @@ describe('cloudAgentNextRouter.prepareSession', () => { cloudAgentSessionId: 'agent_123', kiloSessionId: 'ses_12345678901234567890123456', }); - mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValue({ - isFree: false, - hasUserByokAvailable: false, - }); }); it('rejects devcontainer sessions when the feature flag is disabled', async () => { @@ -406,63 +392,8 @@ describe('cloudAgentNextRouter.prepareSession', () => { ); }); - it('routes free models through the AppBuilder client so the worker skips the balance minimum', async () => { - mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValueOnce({ - isFree: true, - hasUserByokAvailable: false, - }); - const caller = createCaller({ - user: { id: 'user-free', is_admin: false } as User, - }); - - await caller.prepareSession({ - prompt: 'Test prompt', - mode: 'code', - model: 'kilo/test-model', - githubRepo: 'acme/repo', - autoInitiate: true, - devcontainer: false, - }); - - expect(mockComputeCloudAgentNextBalanceCheckEligibility).toHaveBeenCalledWith( - expect.objectContaining({ modelId: 'kilo/test-model' }) - ); - expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', { - isFree: true, - hasUserByokAvailable: false, - }); - expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); - }); - - it('routes BYOK-capable paid models through the AppBuilder client so the worker skips the balance minimum', async () => { - mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValueOnce({ - isFree: false, - hasUserByokAvailable: true, - }); - const caller = createCaller({ - user: { id: 'user-byok', is_admin: false } as User, - }); - - await caller.prepareSession({ - prompt: 'Test prompt', - mode: 'code', - model: 'kilo/paid-byok-model', - githubRepo: 'acme/repo', - autoInitiate: true, - devcontainer: false, - }); - - expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', { - isFree: false, - hasUserByokAvailable: true, - }); - expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); - }); - - it('routes paid models the user has no BYOK key for through the model-aware helper with a paid eligibility', async () => { - const caller = createCaller({ - user: { id: 'user-paid', is_admin: false } as User, - }); + it('uses the ordinary client without local model eligibility', async () => { + const caller = createCaller({ user: { id: 'user-paid', is_admin: false } as User }); await caller.prepareSession({ prompt: 'Test prompt', @@ -473,9 +404,6 @@ describe('cloudAgentNextRouter.prepareSession', () => { devcontainer: false, }); - expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', { - isFree: false, - hasUserByokAvailable: false, - }); + expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('cloud-agent-token'); }); }); diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index d02ee9579e..beb2e662d3 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -2,10 +2,8 @@ import 'server-only'; import { baseProcedure, createTRPCRouter } from '@/lib/trpc/init'; import { createCloudAgentNextClient, - createCloudAgentNextClientForModel, rethrowAsPaymentRequired, } from '@/lib/cloud-agent-next/cloud-agent-client'; -import { computeCloudAgentNextBalanceCheckEligibility } from '@/lib/cloud-agent-next/balance-check-eligibility'; import { rethrowAsTerminalError } from '@/lib/cloud-agent-next/terminal-errors'; import { generateCloudAgentToken } from '@/lib/tokens'; import { isFeatureFlagEnabledOrDevelopment } from '@/lib/posthog-feature-flags'; @@ -137,12 +135,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ } const authToken = generateCloudAgentToken(ctx.user); - const eligibility = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: db, - user: ctx.user, - modelId: input.model, - }); - const client = createCloudAgentNextClientForModel(authToken, eligibility); + const client = createCloudAgentNextClient(authToken); const { gitlabProject, githubRepo, attachments, images, ...restInput } = input; diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts index 554b016234..ddee3b49fd 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts @@ -54,24 +54,13 @@ const mockGenerateCloudAgentAttachmentUploadUrl = jest.fn< }) => Promise<{ signedUrl: string; key: string; expiresAt: string }> >(() => Promise.resolve({ signedUrl: 'signed', key: 'key', expiresAt: 'expires' })); -const mockCreateCloudAgentNextClient = jest.fn(() => ({ - prepareSession: mockPrepareSession, - sendMessage: mockSendMessage, -})); - -const mockCreateCloudAgentNextClientForModel = jest.fn( - (_authToken: string, _eligibility: unknown) => ({ +const mockCreateCloudAgentNextClient = jest.fn((authToken: string) => { + void authToken; + return { prepareSession: mockPrepareSession, sendMessage: mockSendMessage, - }) -); - -const mockComputeCloudAgentNextBalanceCheckEligibility = jest.fn< - (...args: unknown[]) => Promise<{ - isFree: boolean; - hasUserByokAvailable: boolean; - }> ->(); + }; +}); const mockIsFeatureFlagEnabledOrDevelopment = jest.fn<(flagName: string, distinctId: string) => Promise>(); @@ -115,14 +104,9 @@ jest.mock('@/lib/tokens', () => ({ jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ createCloudAgentNextClient: mockCreateCloudAgentNextClient, - createCloudAgentNextClientForModel: mockCreateCloudAgentNextClientForModel, rethrowAsPaymentRequired: jest.fn(), })); -jest.mock('@/lib/cloud-agent-next/balance-check-eligibility', () => ({ - computeCloudAgentNextBalanceCheckEligibility: mockComputeCloudAgentNextBalanceCheckEligibility, -})); - jest.mock('@/lib/posthog-feature-flags', () => ({ isFeatureFlagEnabledOrDevelopment: mockIsFeatureFlagEnabledOrDevelopment, })); @@ -214,7 +198,8 @@ let createCaller: (ctx: { user: User }) => { balance: number; minBalance: number; isEligible: boolean; - accessLevel: 'full' | 'limited' | 'blocked'; + accessLevel: 'full' | 'limited'; + isLowBalance: boolean; }>; listGitHubRepositories: (input: { organizationId: string; @@ -344,11 +329,12 @@ describe('organizationCloudAgentNextRouter helper procedures', () => { }); it.each([ - { balance: 1, isEligible: true, accessLevel: 'full' as const }, - { balance: 0.99, isEligible: false, accessLevel: 'limited' as const }, + { balance: 0, isEligible: false, accessLevel: 'limited' as const, isLowBalance: true }, + { balance: 0.01, isEligible: true, accessLevel: 'full' as const, isLowBalance: true }, + { balance: 1, isEligible: true, accessLevel: 'full' as const, isLowBalance: false }, ])( 'reports organization eligibility for a $balance balance', - async ({ balance, isEligible, accessLevel }) => { + async ({ balance, isEligible, accessLevel, isLowBalance }) => { mockGetBalanceForOrganizationUser.mockResolvedValue({ balance }); const caller = createCaller({ user: { id: 'member-user', is_admin: false } as User }); @@ -357,6 +343,7 @@ describe('organizationCloudAgentNextRouter helper procedures', () => { minBalance: 1, isEligible, accessLevel, + isLowBalance, }); expect(mockEnsureOrganizationAccess).toHaveBeenCalledWith('member-user', ORGANIZATION_ID); expect(mockGetBalanceForOrganizationUser).toHaveBeenCalledWith( @@ -469,10 +456,6 @@ describe('organizationCloudAgentNextRouter.prepareSession', () => { cloudAgentSessionId: 'agent_123', kiloSessionId: 'ses_12345678901234567890123456', }); - mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValue({ - isFree: false, - hasUserByokAvailable: false, - }); }); it('rejects devcontainer sessions when the feature flag is disabled', async () => { @@ -590,61 +573,7 @@ describe('organizationCloudAgentNextRouter.prepareSession', () => { ); }); - it('routes free models through the AppBuilder client so the worker skips the balance minimum', async () => { - mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValueOnce({ - isFree: true, - hasUserByokAvailable: false, - }); - const caller = createCaller({ user: { id: 'user-free', is_admin: false } as User }); - - await caller.prepareSession({ - organizationId: ORGANIZATION_ID, - prompt: 'Test prompt', - mode: 'code', - model: 'kilo/test-model', - githubRepo: 'acme/repo', - autoInitiate: true, - devcontainer: false, - }); - - expect(mockComputeCloudAgentNextBalanceCheckEligibility).toHaveBeenCalledWith( - expect.objectContaining({ - modelId: 'kilo/test-model', - organizationId: ORGANIZATION_ID, - }) - ); - expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', { - isFree: true, - hasUserByokAvailable: false, - }); - expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); - }); - - it('routes BYOK-capable paid models through the AppBuilder client so the worker skips the balance minimum', async () => { - mockComputeCloudAgentNextBalanceCheckEligibility.mockResolvedValueOnce({ - isFree: false, - hasUserByokAvailable: true, - }); - const caller = createCaller({ user: { id: 'user-byok', is_admin: false } as User }); - - await caller.prepareSession({ - organizationId: ORGANIZATION_ID, - prompt: 'Test prompt', - mode: 'code', - model: 'kilo/paid-byok-model', - githubRepo: 'acme/repo', - autoInitiate: true, - devcontainer: false, - }); - - expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', { - isFree: false, - hasUserByokAvailable: true, - }); - expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); - }); - - it('routes paid models the org has no BYOK key for through the model-aware helper with a paid eligibility', async () => { + it('uses the ordinary client without local model eligibility', async () => { const caller = createCaller({ user: { id: 'user-paid', is_admin: false } as User }); await caller.prepareSession({ @@ -657,10 +586,7 @@ describe('organizationCloudAgentNextRouter.prepareSession', () => { devcontainer: false, }); - expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', { - isFree: false, - hasUserByokAvailable: false, - }); + expect(mockCreateCloudAgentNextClient).toHaveBeenCalledWith('cloud-agent-token'); }); }); diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index ce672ff0e7..2e4ab5fa90 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -2,10 +2,8 @@ import 'server-only'; import { createTRPCRouter } from '@/lib/trpc/init'; import { createCloudAgentNextClient, - createCloudAgentNextClientForModel, rethrowAsPaymentRequired, } from '@/lib/cloud-agent-next/cloud-agent-client'; -import { computeCloudAgentNextBalanceCheckEligibility } from '@/lib/cloud-agent-next/balance-check-eligibility'; import { rethrowAsTerminalError } from '@/lib/cloud-agent-next/terminal-errors'; import { generateCloudAgentToken } from '@/lib/tokens'; import { isFeatureFlagEnabledOrDevelopment } from '@/lib/posthog-feature-flags'; @@ -226,13 +224,7 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ } const authToken = generateCloudAgentToken(ctx.user); - const eligibility = await computeCloudAgentNextBalanceCheckEligibility({ - fromDb: db, - user: ctx.user, - modelId: input.model, - organizationId: input.organizationId, - }); - const client = createCloudAgentNextClientForModel(authToken, eligibility); + const client = createCloudAgentNextClient(authToken); const { gitlabProject, diff --git a/services/cloud-agent-next/AGENTS.md b/services/cloud-agent-next/AGENTS.md index 562b016820..e7e83493d1 100644 --- a/services/cloud-agent-next/AGENTS.md +++ b/services/cloud-agent-next/AGENTS.md @@ -123,7 +123,7 @@ This pattern blocks API endpoints from running for external contributors who don - Callback delivery retry policy is paired with `wrangler.jsonc`: `CALLBACK_DELIVERY_MAX_ATTEMPTS` includes the initial attempt, and each Cloud Agent Next callback queue consumer must configure `max_retries` for the remaining redeliveries. - Queue/drain emits unfenced `MessageDeliveryRequest`; only `AgentRuntime` may allocate/reuse current identity and construct `FencedWrapperDispatchRequest` with complete `WrapperRunFence` for downstream dispatch. - Session creation selects an explicit `ProfileResolutionPolicy` at the handler boundary. Implicit repository/default profile resolution is limited to the closed set of approved session origins; omitted, unknown, and non-approved automation origins fail closed unless they supply an explicit profile id. -- Public `start` must authorize any supplied `kilocodeOrganizationId` against `organization_memberships` before resolving profile layers or creating session ownership state. Balance validation is billing-only and `x-skip-balance-check` must never bypass organization authorization. +- Public `start` must authorize any supplied `kilocodeOrganizationId` against `organization_memberships` before resolving profile layers or creating session ownership state. Balance validation is billing-only and must never bypass organization authorization. - Current wrapper identity is fenced `wrapperRunId` plus generation/connection; do not reintroduce execution-ID-only reconnect, supervision, or pending-drain blocking. Legacy endpoint/result/callback `executionId` fields remain boundary compatibility aliases only. - The DO should remain a durable coordinator: queue messages, persist metadata/events, fence wrapper connections, schedule alarms, prepare/restore sandbox state, and hand work to the wrapper. - Put Kilo/job behavior in `wrapper/` or Kilo SDK integration code when it does not require durable DO coordination. diff --git a/services/cloud-agent-next/src/balance-validation.test.ts b/services/cloud-agent-next/src/balance-validation.test.ts index 4869d5f6ed..000e3aa2f9 100644 --- a/services/cloud-agent-next/src/balance-validation.test.ts +++ b/services/cloud-agent-next/src/balance-validation.test.ts @@ -1,263 +1,11 @@ -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { describe, expect, it } from 'vitest'; import { - validateBalanceOnly, extractProcedureName, extractOrgIdFromUrl, - fetchOrgIdForSession, BALANCE_REQUIRED_MUTATIONS, } from './balance-validation.js'; -import type { PersistenceEnv } from './persistence/types.js'; -import type { Env } from './types.js'; -import { parseSessionMetadata } from './persistence/session-metadata.js'; - -// Mock the session-service module -vi.mock('./session-service.js', () => ({ - fetchSessionMetadata: vi.fn(), -})); - -vi.mock('./logger.js', () => ({ - logger: { - withFields: () => ({ error: vi.fn(), warn: vi.fn() }), - error: vi.fn(), - warn: vi.fn(), - }, -})); - -import { fetchSessionMetadata } from './session-service.js'; describe('balance-validation', () => { - const originalFetch = global.fetch; - let fetchMock: ReturnType; - - const mockEnv = { - NEXTAUTH_SECRET: 'test-secret', - KILOCODE_BACKEND_BASE_URL: 'https://app.kilo.ai', - Sandbox: {}, - CLOUD_AGENT_SESSION: {}, - } as unknown as Env; - - beforeEach(() => { - vi.clearAllMocks(); - fetchMock = vi.fn(); - global.fetch = fetchMock as unknown as typeof fetch; - }); - - afterEach(() => { - global.fetch = originalFetch; - }); - - describe('validateBalanceOnly', () => { - describe('balance validation', () => { - it('returns 402 when balance is depleted', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 0.5, isDepleted: true }), - } as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ - success: false, - status: 402, - message: 'Insufficient credits: $1 minimum required', - }); - }); - - it('returns 402 when balance is below $1', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 0.99, isDepleted: false }), // Just under $1 - } as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ - success: false, - status: 402, - message: 'Insufficient credits: $1 minimum required', - }); - }); - - it('returns 402 when balance is zero', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 0, isDepleted: false }), - } as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ - success: false, - status: 402, - message: 'Insufficient credits: $1 minimum required', - }); - }); - - it('returns 402 when balance is negative', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: -5, isDepleted: true }), - } as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ - success: false, - status: 402, - message: 'Insufficient credits: $1 minimum required', - }); - }); - }); - - describe('error handling', () => { - it('returns 500 when balance API returns error', async () => { - fetchMock.mockResolvedValue({ - ok: false, - status: 500, - } as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ - success: false, - status: 500, - message: 'Failed to verify balance', - }); - }); - - it('returns 500 when fetch throws', async () => { - fetchMock.mockRejectedValue(new Error('Network error')); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ - success: false, - status: 500, - message: 'Failed to verify balance', - }); - }); - - it('returns 500 when response JSON is invalid', async () => { - fetchMock.mockResolvedValue({ - ok: true, - status: 200, - json: async () => { - throw new Error('Invalid JSON'); - }, - } as unknown as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ - success: false, - status: 500, - message: 'Invalid balance response', - }); - }); - - it('returns 402 when balance is not a number', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 'not-a-number', isDepleted: false }), - } as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ - success: false, - status: 402, - message: 'Insufficient credits: $1 minimum required', - }); - }); - }); - - describe('successful validation', () => { - it('returns success when balance is sufficient', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 5, isDepleted: false }), // $5 - } as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ success: true }); - }); - - it('returns success with exactly $1 balance', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 1, isDepleted: false }), // Exactly $1 - } as Response); - - const result = await validateBalanceOnly('valid-token', undefined, mockEnv); - - expect(result).toEqual({ success: true }); - }); - - it('uses default API URL when KILOCODE_BACKEND_BASE_URL not configured', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 5, isDepleted: false }), - } as Response); - const envWithoutBackendUrl = { ...mockEnv, KILOCODE_BACKEND_BASE_URL: undefined }; - - const result = await validateBalanceOnly( - 'valid-token', - undefined, - envWithoutBackendUrl as Env - ); - - expect(result).toEqual({ success: true }); - // Should use default URL https://api.kilo.ai - expect(fetchMock).toHaveBeenCalledWith( - 'https://api.kilo.ai/api/profile/balance', - expect.objectContaining({ - method: 'GET', - }) - ); - }); - }); - - describe('organization header', () => { - it('includes X-KiloCode-OrganizationId header when orgId provided', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 5, isDepleted: false }), - } as Response); - - const orgId = '11111111-2222-3333-4444-555555555555'; - await validateBalanceOnly('valid-token', orgId, mockEnv); - - expect(fetchMock).toHaveBeenCalledWith( - `${mockEnv.KILOCODE_BACKEND_BASE_URL}/api/profile/balance`, - expect.objectContaining({ - method: 'GET', - headers: expect.any(Headers) as unknown, - }) - ); - - const [, init] = fetchMock.mock.calls[0] as [string, RequestInit]; - const headers = init.headers as Headers; - expect(headers.get('Authorization')).toBe('Bearer valid-token'); - expect(headers.get('X-KiloCode-OrganizationId')).toBe(orgId); - }); - - it('does not include X-KiloCode-OrganizationId header when orgId not provided', async () => { - fetchMock.mockResolvedValue({ - ok: true, - json: async () => ({ balance: 5, isDepleted: false }), - } as Response); - - await validateBalanceOnly('valid-token', undefined, mockEnv); - - const [, init] = fetchMock.mock.calls[0] as [string, RequestInit]; - const headers = init.headers as Headers; - expect(headers.get('Authorization')).toBe('Bearer valid-token'); - expect(headers.get('X-KiloCode-OrganizationId')).toBeNull(); - }); - }); - }); - describe('extractProcedureName', () => { it('extracts procedure name from valid tRPC path', () => { expect(extractProcedureName('/trpc/initiateSessionStream')).toBe('initiateSessionStream'); @@ -341,84 +89,16 @@ describe('balance-validation', () => { }); }); - describe('fetchOrgIdForSession', () => { - const mockPersistenceEnv = { - CLOUD_AGENT_SESSION: { - idFromName: vi.fn(), - get: vi.fn(), - }, - SESSION_INGEST: { - fetch: vi.fn(), - }, - } as unknown as PersistenceEnv; - - beforeEach(() => { - vi.mocked(fetchSessionMetadata).mockReset(); - }); - - it('returns orgId when session metadata exists', async () => { - vi.mocked(fetchSessionMetadata).mockResolvedValue( - parseSessionMetadata({ - version: 1, - sessionId: 'agent_123', - orgId: 'org-456', - userId: 'user-789', - timestamp: Date.now(), - }) - ); - - const result = await fetchOrgIdForSession(mockPersistenceEnv, 'user-789', 'agent_123'); - - expect(result).toBe('org-456'); - expect(fetchSessionMetadata).toHaveBeenCalledWith( - mockPersistenceEnv, - 'user-789', - 'agent_123' - ); - }); - - it('returns undefined when session metadata does not exist', async () => { - vi.mocked(fetchSessionMetadata).mockResolvedValue(null); - - const result = await fetchOrgIdForSession(mockPersistenceEnv, 'user-789', 'agent_123'); - - expect(result).toBeUndefined(); - }); - - it('returns undefined when session metadata has no orgId (personal account)', async () => { - vi.mocked(fetchSessionMetadata).mockResolvedValue( - parseSessionMetadata({ - version: 1, - sessionId: 'agent_123', - userId: 'user-789', - timestamp: Date.now(), - }) - ); - - const result = await fetchOrgIdForSession(mockPersistenceEnv, 'user-789', 'agent_123'); - - expect(result).toBeUndefined(); - }); - - it('returns undefined and logs warning when fetchSessionMetadata throws', async () => { - vi.mocked(fetchSessionMetadata).mockRejectedValue(new Error('DO unavailable')); - - const result = await fetchOrgIdForSession(mockPersistenceEnv, 'user-789', 'agent_123'); - - expect(result).toBeUndefined(); - }); - }); - describe('BALANCE_REQUIRED_MUTATIONS', () => { it('contains expected V2 mutation procedures', () => { expect(BALANCE_REQUIRED_MUTATIONS.has('initiateFromKilocodeSessionV2')).toBe(true); expect(BALANCE_REQUIRED_MUTATIONS.has('sendMessageV2')).toBe(true); + expect(BALANCE_REQUIRED_MUTATIONS.has('prepareSession')).toBe(true); }); it('does not contain non-balance-required procedures', () => { expect(BALANCE_REQUIRED_MUTATIONS.has('deleteSession')).toBe(false); expect(BALANCE_REQUIRED_MUTATIONS.has('getSessionLogs')).toBe(false); - expect(BALANCE_REQUIRED_MUTATIONS.has('prepareSession')).toBe(false); }); }); }); diff --git a/services/cloud-agent-next/src/balance-validation.ts b/services/cloud-agent-next/src/balance-validation.ts index 26f913de94..aac39dd87c 100644 --- a/services/cloud-agent-next/src/balance-validation.ts +++ b/services/cloud-agent-next/src/balance-validation.ts @@ -1,75 +1,12 @@ -import { DEFAULT_BACKEND_URL } from './constants.js'; -import { logger } from './logger.js'; -import type { Env } from './types.js'; -import type { PersistenceEnv } from './persistence/types.js'; -import { fetchSessionMetadata } from './session-service.js'; - -const MIN_BALANCE_DOLLARS = 1; - /** - * Result of balance validation - either success or failure with HTTP status. - * Used when auth has already been validated by middleware. + * Result of a balance check — either success or failure with HTTP status. + * Used when auth has already been validated by middleware and balance is + * resolved as part of the cloud agent admission check. */ export type BalanceOnlyResult = | { success: true } | { success: false; status: 402 | 500; message: string }; -/** - * Validates balance only, skipping JWT validation. - * Use this when auth has already been validated by middleware. - * - * @param token - The already-validated JWT token - * @param orgId - Optional organization ID for org-specific balance check - * @param env - Worker environment with secrets and bindings - */ -export async function validateBalanceOnly( - token: string, - orgId: string | undefined, - env: Env -): Promise { - const backendUrl = env.KILOCODE_BACKEND_BASE_URL || DEFAULT_BACKEND_URL; - - const headers = new Headers({ - Authorization: `Bearer ${token}`, - }); - if (orgId) { - headers.set('X-KiloCode-OrganizationId', orgId); - } - - let response: Response; - try { - response = await fetch(`${backendUrl}/api/profile/balance`, { - method: 'GET', - headers, - }); - } catch (error) { - logger - .withFields({ error: error instanceof Error ? error.message : String(error) }) - .error('Failed to fetch balance'); - return { success: false, status: 500, message: 'Failed to verify balance' }; - } - - if (!response.ok) { - logger - .withFields({ status: response.status, statusText: response.statusText }) - .error('Balance API returned error'); - return { success: false, status: 500, message: 'Failed to verify balance' }; - } - - let data: { balance: number; isDepleted: boolean }; - try { - data = await response.json(); - } catch { - return { success: false, status: 500, message: 'Invalid balance response' }; - } - - if (data.isDepleted || typeof data.balance !== 'number' || data.balance < MIN_BALANCE_DOLLARS) { - return { success: false, status: 402, message: 'Insufficient credits: $1 minimum required' }; - } - - return { success: true }; -} - /** * Extracts the tRPC procedure name from a URL pathname. * @example "/trpc/initiateSessionStream" -> "initiateSessionStream" @@ -104,31 +41,6 @@ export function extractOrgIdFromUrl(url: URL): string | undefined { return undefined; } -/** - * Fetches the orgId for a session from the Durable Object metadata. - * Used by sendMessageV2 to get the org context for balance validation. - * - * @param env - Worker environment with DO bindings - * @param userId - User ID from auth - * @param sessionId - Session ID from URL input - * @returns The orgId if found, undefined otherwise - */ -export async function fetchOrgIdForSession( - env: PersistenceEnv, - userId: string, - sessionId: string -): Promise { - try { - const metadata = await fetchSessionMetadata(env, userId, sessionId); - return metadata?.identity.orgId; - } catch (error) { - logger - .withFields({ error: error instanceof Error ? error.message : String(error), sessionId }) - .warn('Failed to fetch session metadata for balance validation'); - return undefined; - } -} - /** * Set of V2 mutation procedure names that require balance validation. * @@ -137,6 +49,7 @@ export async function fetchOrgIdForSession( * once the queued message is flushed to the wrapper. */ export const BALANCE_REQUIRED_MUTATIONS = new Set([ + 'prepareSession', 'initiateFromKilocodeSessionV2', 'sendMessageV2', 'start', diff --git a/services/cloud-agent-next/src/cloud-agent-admission.test.ts b/services/cloud-agent-next/src/cloud-agent-admission.test.ts new file mode 100644 index 0000000000..d40b0cabd5 --- /dev/null +++ b/services/cloud-agent-next/src/cloud-agent-admission.test.ts @@ -0,0 +1,146 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { checkCloudAgentAdmission } from './cloud-agent-admission.js'; +import type { Env } from './types.js'; + +vi.mock('./logger.js', () => ({ + logger: { + withFields: () => ({ error: vi.fn(), warn: vi.fn() }), + error: vi.fn(), + warn: vi.fn(), + }, +})); + +describe('checkCloudAgentAdmission', () => { + const originalFetch = global.fetch; + let fetchMock: ReturnType; + + const mockEnv = { + KILOCODE_BACKEND_BASE_URL: 'https://app.kilo.ai', + } as unknown as Env; + + beforeEach(() => { + vi.clearAllMocks(); + fetchMock = vi.fn(); + global.fetch = fetchMock as unknown as typeof fetch; + }); + + afterEach(() => { + global.fetch = originalFetch; + }); + + function okWith(body: unknown) { + return { ok: true, status: 200, json: async () => body } as Response; + } + + it('returns classification without balance for a non-balance-required model', async () => { + fetchMock.mockResolvedValue( + okWith({ classification: 'byok', balance: null, isDepleted: null }) + ); + + const result = await checkCloudAgentAdmission({ + env: mockEnv, + token: 'jwt-token', + modelId: 'anthropic/claude-sonnet-4', + owner: { userId: 'user-1' }, + }); + + expect(result).toEqual({ classification: 'byok', balance: null, isDepleted: null }); + }); + + it('returns classification with balance for a balance-required model', async () => { + fetchMock.mockResolvedValue( + okWith({ classification: 'balance-required', balance: 5, isDepleted: false }) + ); + + const result = await checkCloudAgentAdmission({ + env: mockEnv, + token: 'jwt-token', + modelId: 'anthropic/claude-sonnet-4', + owner: { userId: 'user-1' }, + }); + + expect(result).toEqual({ classification: 'balance-required', balance: 5, isDepleted: false }); + }); + + it('posts the model with the bearer token and no org header for a user owner', async () => { + fetchMock.mockResolvedValue( + okWith({ classification: 'free', balance: null, isDepleted: null }) + ); + + await checkCloudAgentAdmission({ + env: mockEnv, + token: 'jwt-token', + modelId: 'kilo/free-model', + owner: { userId: 'user-1' }, + }); + + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe('https://app.kilo.ai/api/profile/cloud-agent-admission'); + expect(init.method).toBe('POST'); + expect(init.body).toBe(JSON.stringify({ modelId: 'kilo/free-model' })); + const headers = init.headers as Headers; + expect(headers.get('Authorization')).toBe('Bearer jwt-token'); + expect(headers.get('X-KiloCode-OrganizationId')).toBeNull(); + }); + + it('sends the organization header for an organization owner', async () => { + fetchMock.mockResolvedValue( + okWith({ classification: 'byok', balance: null, isDepleted: null }) + ); + + await checkCloudAgentAdmission({ + env: mockEnv, + token: 'jwt-token', + modelId: 'anthropic/claude-sonnet-4', + owner: { organizationId: 'org-1' }, + }); + + const [, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + const headers = init.headers as Headers; + expect(headers.get('X-KiloCode-OrganizationId')).toBe('org-1'); + }); + + it('falls back to the default backend URL when none is configured', async () => { + fetchMock.mockResolvedValue( + okWith({ classification: 'free', balance: null, isDepleted: null }) + ); + + await checkCloudAgentAdmission({ + env: { ...mockEnv, KILOCODE_BACKEND_BASE_URL: undefined } as Env, + token: 'jwt-token', + modelId: 'kilo/free-model', + owner: { userId: 'user-1' }, + }); + + const [url] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe('https://api.kilo.ai/api/profile/cloud-agent-admission'); + }); + + it.each([ + ['fetch rejects', () => fetchMock.mockRejectedValue(new Error('Network error'))], + ['non-ok response', () => fetchMock.mockResolvedValue({ ok: false, status: 500 } as Response)], + [ + 'non-JSON body', + () => + fetchMock.mockResolvedValue({ + ok: true, + status: 200, + json: async () => { + throw new Error('Invalid JSON'); + }, + } as unknown as Response), + ], + ['unexpected shape', () => fetchMock.mockResolvedValue(okWith({ classification: 'unknown' }))], + ])('throws a retryable SERVICE_UNAVAILABLE on %s', async (_label, arrange) => { + arrange(); + + await expect( + checkCloudAgentAdmission({ + env: mockEnv, + token: 'jwt-token', + modelId: 'anthropic/claude-sonnet-4', + owner: { userId: 'user-1' }, + }) + ).rejects.toMatchObject({ code: 'SERVICE_UNAVAILABLE' }); + }); +}); diff --git a/services/cloud-agent-next/src/cloud-agent-admission.ts b/services/cloud-agent-next/src/cloud-agent-admission.ts new file mode 100644 index 0000000000..ea779a0256 --- /dev/null +++ b/services/cloud-agent-next/src/cloud-agent-admission.ts @@ -0,0 +1,129 @@ +import { TRPCError } from '@trpc/server'; +import { z } from 'zod'; +import { DEFAULT_BACKEND_URL } from './constants.js'; +import { logger } from './logger.js'; +import type { Env } from './types.js'; + +/** + * How a Cloud Agent session for a given model should be billed. Mirrors the + * source-of-truth classifier in `apps/web` + * (`lib/cloud-agent-next/classify-model-billing`), which this worker asks over + * HTTP because it cannot import the app's model catalog. + */ +export type CloudAgentModelBilling = 'free' | 'byok' | 'balance-required'; + +export type CloudAgentModelBillingOwner = + | { userId: string; organizationId?: never } + | { organizationId: string; userId?: never }; + +/** + * Result of the single admission round-trip. `balance`/`isDepleted` are only + * populated when `classification === 'balance-required'` (the only case where + * balance gates the request); otherwise they are `null`. + */ +export type CloudAgentAdmission = { + classification: CloudAgentModelBilling; + balance: number | null; + isDepleted: boolean | null; +}; + +const admissionResponseSchema = z.object({ + classification: z.enum(['free', 'byok', 'balance-required']), + balance: z.number().nullable(), + isDepleted: z.boolean().nullable(), +}); + +export const INSUFFICIENT_CREDITS_MESSAGE = + 'Insufficient credits: a positive credit balance is required'; + +/** + * Whether a `balance-required` admission must be rejected for lack of funds. + * Callers gate on `classification === 'balance-required'` first; this mirrors + * the old worker-side balance check (`isDepleted || balance <= 0`). + */ +export function hasInsufficientBalance(admission: { + balance: number | null; + isDepleted: boolean | null; +}): boolean { + return admission.isDepleted === true || (admission.balance ?? 0) <= 0; +} + +const ADMISSION_UNAVAILABLE_MESSAGE = 'Cloud agent admission could not be verified'; + +/** + * Fail loud-but-retryable when the backend cannot answer, matching how + * `assertKiloModelAvailable` treats an unreachable model catalog. Rejecting + * (rather than silently skipping the balance gate) keeps platform-billed models + * from being served for free during an outage. + */ +function admissionUnavailable(): TRPCError { + return new TRPCError({ + code: 'SERVICE_UNAVAILABLE', + message: ADMISSION_UNAVAILABLE_MESSAGE, + cause: { + error: 'CLOUD_AGENT_ADMISSION_UNAVAILABLE', + message: ADMISSION_UNAVAILABLE_MESSAGE, + retryable: true, + }, + }); +} + +/** + * Ask `apps/web` whether a model is admissible for the given owner: how it is + * billed and, when platform-billed, whether the owner has balance. The owner is + * conveyed the same way the balance check conveyed it — the JWT identifies the + * user, and `X-KiloCode-OrganizationId` selects the organization owner (whose + * membership the backend validates). + */ +export async function checkCloudAgentAdmission(params: { + env: Env; + token: string; + modelId: string; + owner: CloudAgentModelBillingOwner; +}): Promise { + const backendUrl = params.env.KILOCODE_BACKEND_BASE_URL || DEFAULT_BACKEND_URL; + + const headers = new Headers({ + 'Content-Type': 'application/json', + Authorization: `Bearer ${params.token}`, + }); + if (params.owner.organizationId) { + headers.set('X-KiloCode-OrganizationId', params.owner.organizationId); + } + + let response: Response; + try { + response = await fetch(`${backendUrl}/api/profile/cloud-agent-admission`, { + method: 'POST', + headers, + body: JSON.stringify({ modelId: params.modelId }), + }); + } catch (error) { + logger + .withFields({ error: error instanceof Error ? error.message : String(error) }) + .error('Failed to fetch cloud agent admission'); + throw admissionUnavailable(); + } + + if (!response.ok) { + logger + .withFields({ status: response.status, statusText: response.statusText }) + .error('Cloud agent admission API returned an error'); + throw admissionUnavailable(); + } + + let data: unknown; + try { + data = await response.json(); + } catch { + logger.error('Cloud agent admission API returned a non-JSON body'); + throw admissionUnavailable(); + } + + const parsed = admissionResponseSchema.safeParse(data); + if (!parsed.success) { + logger.error('Cloud agent admission API returned an unexpected shape'); + throw admissionUnavailable(); + } + return parsed.data; +} diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts index 39ae833945..6b0a2bd4ed 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.test.ts @@ -2208,9 +2208,13 @@ describe('handleKiloFacadeRequest', () => { expect(admitPrompt).not.toHaveBeenCalled(); }); - it('rejects prompt_async attempts to bypass public balance validation after owner resolution', async () => { + it('does not let a balance bypass header change public prompt admission', async () => { const env = envStub(); - const validatePromptBalance = vi.fn(); + const validatePromptBalance = vi.fn().mockResolvedValue({ + success: false, + status: 402, + message: 'Insufficient credits: a positive credit balance is required', + }); const admitPrompt = vi.fn(); const resolveRootSessionForKiloSession = vi.fn(async () => ({ cloudAgentSessionId: 'agent_owned', @@ -2235,17 +2239,17 @@ describe('handleKiloFacadeRequest', () => { }, }); - expect(response.status).toBe(400); + expect(response.status).toBe(402); await expect(response.json()).resolves.toEqual({ - error: 'KILO_BALANCE_BYPASS_UNSUPPORTED', - message: 'Balance bypass is not supported for public Kilo prompt mutations', + error: 'KILO_BALANCE_VALIDATION_FAILED', + message: 'Insufficient credits: a positive credit balance is required', }); expect(resolveRootSessionForKiloSession).toHaveBeenCalledWith({ env, userId: 'usr_1', kiloSessionId, }); - expect(validatePromptBalance).not.toHaveBeenCalled(); + expect(validatePromptBalance).toHaveBeenCalled(); expect(admitPrompt).not.toHaveBeenCalled(); }); @@ -2286,7 +2290,7 @@ describe('handleKiloFacadeRequest', () => { const validatePromptBalance = vi.fn().mockResolvedValue({ success: false, status: 402, - message: 'Insufficient credits: $1 minimum required', + message: 'Insufficient credits: a positive credit balance is required', }); const response = await handleKiloFacadeRequest({ diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index 3ff7d2ba41..50b30bf275 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -10,11 +10,9 @@ import type { KiloSdkStoredMessage, ListCloudAgentRootSessionsParams, } from '../session-ingest-binding.js'; -import { - fetchOrgIdForSession, - validateBalanceOnly, - type BalanceOnlyResult, -} from '../balance-validation.js'; +import { type BalanceOnlyResult } from '../balance-validation.js'; +import { hasInsufficientBalance, INSUFFICIENT_CREDITS_MESSAGE } from '../cloud-agent-admission.js'; +import { preflightCloudAgentModelBilling } from '../model-billing-preflight.js'; import type { QueueExecutionTurnCommand, SubmittedSessionMessageRequest, @@ -127,6 +125,7 @@ export type KiloFacadeRequestDeps = { authToken: string; userId: string; cloudAgentSessionId: string; + requestedModel?: string; }) => Promise; interruptPrompt?: (params: { env: Env; @@ -710,9 +709,25 @@ async function defaultValidatePromptBalance(params: { authToken: string; userId: string; cloudAgentSessionId: string; + requestedModel?: string; }): Promise { - const orgId = await fetchOrgIdForSession(params.env, params.userId, params.cloudAgentSessionId); - return validateBalanceOnly(params.authToken, orgId, params.env); + const billing = await preflightCloudAgentModelBilling({ + env: params.env, + userId: params.userId, + authToken: params.authToken, + procedure: 'kilo.prompt_async', + body: { + cloudAgentSessionId: params.cloudAgentSessionId, + ...(params.requestedModel ? { agent: { model: params.requestedModel } } : {}), + }, + }); + if (billing.classification !== 'balance-required') { + return { success: true }; + } + if (hasInsufficientBalance(billing)) { + return { success: false, status: 402, message: INSUFFICIENT_CREDITS_MESSAGE }; + } + return { success: true }; } async function defaultAdmitPrompt(params: { @@ -754,18 +769,12 @@ async function admitBasicPrompt(params: { if (!params.authToken) { return facadeError(500, 'KILO_FACADE_UNAVAILABLE', 'Durable prompt admission is unavailable'); } - if (params.request.headers.get('x-skip-balance-check') !== null) { - return facadeError( - 400, - 'KILO_BALANCE_BYPASS_UNSUPPORTED', - 'Balance bypass is not supported for public Kilo prompt mutations' - ); - } const balance = await (params.deps?.validatePromptBalance ?? defaultValidatePromptBalance)({ env: params.env, authToken: params.authToken, userId: params.userId, cloudAgentSessionId: params.cloudAgentSessionId, + requestedModel: parsed.prompt.agent?.model, }); if (!balance.success) { return facadeError(balance.status, 'KILO_BALANCE_VALIDATION_FAILED', balance.message); diff --git a/services/cloud-agent-next/src/middleware/balance.test.ts b/services/cloud-agent-next/src/middleware/balance.test.ts index 474fb89af9..f742b5b0ea 100644 --- a/services/cloud-agent-next/src/middleware/balance.test.ts +++ b/services/cloud-agent-next/src/middleware/balance.test.ts @@ -1,14 +1,16 @@ import { Hono } from 'hono'; +import { TRPCError } from '@trpc/server'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { HonoContext } from '../hono-context.js'; import type { Env } from '../types.js'; -const { requireCurrentSessionAccessMock } = vi.hoisted(() => ({ - requireCurrentSessionAccessMock: vi.fn(), +const { preflightCloudAgentModelBillingMock } = vi.hoisted(() => ({ + preflightCloudAgentModelBillingMock: vi.fn(), })); vi.mock('../balance-validation.js', () => ({ BALANCE_REQUIRED_MUTATIONS: new Set([ + 'prepareSession', 'initiateFromKilocodeSessionV2', 'sendMessageV2', 'start', @@ -18,8 +20,6 @@ vi.mock('../balance-validation.js', () => ({ const match = pathname.match(/^\/trpc\/([^?/]+)/); return match ? match[1] : null; }, - fetchOrgIdForSession: vi.fn(), - validateBalanceOnly: vi.fn(), })); vi.mock('../logger.js', () => ({ @@ -28,27 +28,28 @@ vi.mock('../logger.js', () => ({ }, })); -vi.mock('../session-access.js', () => ({ - requireCurrentSessionAccess: requireCurrentSessionAccessMock, - projectSessionAccessHttpError: () => new Response('Session access denied', { status: 403 }), +vi.mock('../model-billing-preflight.js', () => ({ + preflightCloudAgentModelBilling: preflightCloudAgentModelBillingMock, })); const { balanceMiddleware } = await import('./balance.js'); -const { fetchOrgIdForSession, validateBalanceOnly } = await import('../balance-validation.js'); +const { INSUFFICIENT_CREDITS_MESSAGE } = await import('../cloud-agent-admission.js'); describe('balanceMiddleware', () => { const env = {} as Env; beforeEach(() => { vi.clearAllMocks(); - vi.mocked(validateBalanceOnly).mockResolvedValue({ success: true }); - requireCurrentSessionAccessMock.mockResolvedValue({ - kiloSessionId: 'ses_12345678901234567890123456', - organizationId: 'org-current', + // Default: balance-required with sufficient funds → admits. + preflightCloudAgentModelBillingMock.mockResolvedValue({ + classification: 'balance-required', + balance: 5, + isDepleted: false, + owner: { userId: 'user-123' }, }); }); - function createApp() { + function createApp(downstream?: () => void) { const app = new Hono(); app.use('/trpc/*', async (c, next) => { c.set('userId', 'user-123'); @@ -56,86 +57,158 @@ describe('balanceMiddleware', () => { await next(); }); app.use('/trpc/*', balanceMiddleware); - app.post('/trpc/:procedure', c => - c.json({ ok: true, validatedSessionAccess: c.get('validatedSessionAccess') }) - ); + app.post('/trpc/:procedure', c => { + downstream?.(); + return c.json({ ok: true, validatedSessionAccess: c.get('validatedSessionAccess') }); + }); return app; } - async function postTrpc(procedureName: string, body: unknown) { - return createApp().fetch( + async function postTrpc( + procedureName: string, + body: unknown, + options?: { headers?: HeadersInit; downstream?: () => void } + ) { + return createApp(options?.downstream).fetch( new Request(`https://worker.test/trpc/${procedureName}`, { method: 'POST', body: JSON.stringify(body), - headers: { 'Content-Type': 'application/json' }, + headers: { 'Content-Type': 'application/json', ...options?.headers }, }), env ); } it('returns non-retryable clientError for insufficient credits', async () => { - vi.mocked(validateBalanceOnly).mockResolvedValue({ - success: false, - status: 402, - message: 'Insufficient credits', + preflightCloudAgentModelBillingMock.mockResolvedValue({ + classification: 'balance-required', + balance: 0, + isDepleted: true, + owner: { userId: 'user-123' }, }); - const response = await postTrpc('start', {}); + const downstream = vi.fn(); + const response = await postTrpc('start', { agent: { model: 'paid/model' } }, { downstream }); const body: any = await response.json(); + expect(response.status).toBe(402); expect(body.error.data.clientError).toEqual({ code: 'PAYMENT_REQUIRED', - message: 'Insufficient credits', + message: INSUFFICIENT_CREDITS_MESSAGE, retryable: false, }); + expect(downstream).not.toHaveBeenCalled(); }); - it('returns retryable clientError for balance infrastructure failures', async () => { - vi.mocked(validateBalanceOnly).mockResolvedValue({ - success: false, - status: 500, - message: 'Failed to verify balance', - }); + it('returns retryable clientError when admission is unavailable', async () => { + preflightCloudAgentModelBillingMock.mockRejectedValue( + new TRPCError({ + code: 'SERVICE_UNAVAILABLE', + message: 'Cloud agent admission could not be verified', + }) + ); - const response = await postTrpc('start', {}); + const downstream = vi.fn(); + const response = await postTrpc('start', { agent: { model: 'paid/model' } }, { downstream }); const body: any = await response.json(); + expect(response.status).toBe(503); expect(body.error.data.clientError).toEqual({ - code: 'INTERNAL_SERVER_ERROR', - message: 'Failed to verify balance', + code: 'SERVICE_UNAVAILABLE', + message: 'Cloud agent admission could not be verified', retryable: true, }); + expect(downstream).not.toHaveBeenCalled(); }); - it('uses nested start organization context for balance validation', async () => { - const orgId = '11111111-2222-3333-4444-555555555555'; - - const response = await postTrpc('start', { - prompt: 'Start a cloud agent session', - options: { - kilocodeOrganizationId: orgId, - }, - }); + it.each(['prepareSession', 'start', 'initiateFromKilocodeSessionV2', 'sendMessageV2', 'send'])( + 'rejects balance-required %s before reaching its handler at zero balance', + async procedureName => { + const downstream = vi.fn(); + preflightCloudAgentModelBillingMock.mockResolvedValue({ + classification: 'balance-required', + balance: 0, + isDepleted: true, + owner: { userId: 'user-123' }, + }); + + const response = await postTrpc( + procedureName, + { cloudAgentSessionId: 'agent-existing' }, + { downstream } + ); + + expect(response.status).toBe(402); + expect(downstream).not.toHaveBeenCalled(); + } + ); + + it.each(['free', 'byok'] as const)( + 'admits a %s model without gating on balance', + async classification => { + preflightCloudAgentModelBillingMock.mockResolvedValue({ + classification, + balance: null, + isDepleted: null, + owner: { userId: 'user-123' }, + }); + const downstream = vi.fn(); + + const response = await postTrpc( + 'start', + { agent: { model: 'available/model' } }, + { downstream } + ); + + expect(response.status).toBe(200); + expect(downstream).toHaveBeenCalledOnce(); + } + ); + + it('admits a balance-required model with sufficient balance', async () => { + const downstream = vi.fn(); + + const response = await postTrpc('start', { agent: { model: 'paid/model' } }, { downstream }); expect(response.status).toBe(200); - expect(validateBalanceOnly).toHaveBeenCalledWith('token-123', orgId, env); - expect(fetchOrgIdForSession).not.toHaveBeenCalled(); + expect(downstream).toHaveBeenCalledOnce(); }); - it('validates caller-supplied organization context against the stored session scope', async () => { - const response = await postTrpc('send', { - cloudAgentSessionId: 'agent-existing', - kilocodeOrganizationId: 'org-requested', + it('does not let the removed header bypass paid zero-balance admission', async () => { + preflightCloudAgentModelBillingMock.mockResolvedValue({ + classification: 'balance-required', + balance: 0, + isDepleted: true, + owner: { userId: 'user-123' }, }); + const downstream = vi.fn(); - expect(response.status).toBe(200); - expect(requireCurrentSessionAccessMock).toHaveBeenCalledWith({ - env, - kiloUserId: 'user-123', - cloudAgentSessionId: 'agent-existing', - expectedOrganizationId: 'org-requested', + const response = await postTrpc( + 'start', + { agent: { model: 'paid/model' } }, + { headers: { 'x-skip-balance-check': 'true' }, downstream } + ); + + expect(response.status).toBe(402); + expect(downstream).not.toHaveBeenCalled(); + }); + + it('retains trusted existing-session access for downstream handlers', async () => { + preflightCloudAgentModelBillingMock.mockResolvedValue({ + classification: 'free', + balance: null, + isDepleted: null, + owner: { organizationId: 'org-current' }, + validatedSessionAccess: { + kiloUserId: 'user-123', + cloudAgentSessionId: 'agent-existing', + kiloSessionId: 'ses_12345678901234567890123456', + organizationId: 'org-current', + }, }); - expect(validateBalanceOnly).toHaveBeenCalledWith('token-123', 'org-current', env); + + const response = await postTrpc('send', { cloudAgentSessionId: 'agent-existing' }); + await expect(response.json()).resolves.toMatchObject({ validatedSessionAccess: { kiloUserId: 'user-123', diff --git a/services/cloud-agent-next/src/middleware/balance.ts b/services/cloud-agent-next/src/middleware/balance.ts index 69aacf9761..92bd86943e 100644 --- a/services/cloud-agent-next/src/middleware/balance.ts +++ b/services/cloud-agent-next/src/middleware/balance.ts @@ -1,14 +1,31 @@ import { createMiddleware } from 'hono/factory'; import type { Context, Next } from 'hono'; +import { TRPCError } from '@trpc/server'; import type { HonoContext } from '../hono-context.js'; import { logger } from '../logger.js'; import { buildTrpcErrorResponse } from '../trpc-error.js'; -import { - validateBalanceOnly, - extractProcedureName, - BALANCE_REQUIRED_MUTATIONS, -} from '../balance-validation.js'; -import { projectSessionAccessHttpError, requireCurrentSessionAccess } from '../session-access.js'; +import { extractProcedureName, BALANCE_REQUIRED_MUTATIONS } from '../balance-validation.js'; +import { preflightCloudAgentModelBilling } from '../model-billing-preflight.js'; +import { hasInsufficientBalance, INSUFFICIENT_CREDITS_MESSAGE } from '../cloud-agent-admission.js'; + +function projectBillingPreflightError(error: unknown): { status: number; message: string } { + if (!(error instanceof TRPCError)) { + return { status: 500, message: 'Billing admission is temporarily unavailable' }; + } + + switch (error.code) { + case 'BAD_REQUEST': + return { status: 400, message: error.message }; + case 'FORBIDDEN': + return { status: 403, message: error.message }; + case 'NOT_FOUND': + return { status: 404, message: error.message }; + case 'SERVICE_UNAVAILABLE': + return { status: 503, message: error.message }; + default: + return { status: 500, message: 'Billing admission is temporarily unavailable' }; + } +} /** * Middleware that validates user balance for mutations that require it. @@ -23,33 +40,10 @@ export const balanceMiddleware = createMiddleware( return; } - const skipBalanceCheck = c.req.header('x-skip-balance-check') === 'true'; - - if (skipBalanceCheck) { - logger.withFields({ procedure: procedureName }).info('Skipping balance check per header'); - await next(); - return; - } - - let orgId: string | undefined; - let sessionId: string | undefined; + let body: unknown; try { const clonedRequest = c.req.raw.clone(); - const body = await clonedRequest.json(); - if (body && typeof body === 'object') { - if ('kilocodeOrganizationId' in body && typeof body.kilocodeOrganizationId === 'string') { - orgId = body.kilocodeOrganizationId; - } - if (!orgId && 'options' in body && body.options && typeof body.options === 'object') { - const options = body.options as Record; - if (typeof options.kilocodeOrganizationId === 'string') { - orgId = options.kilocodeOrganizationId; - } - } - if ('cloudAgentSessionId' in body && typeof body.cloudAgentSessionId === 'string') { - sessionId = body.cloudAgentSessionId; - } - } + body = await clonedRequest.json(); } catch { return buildTrpcErrorResponse(400, 'Invalid request body', procedureName); } @@ -59,48 +53,39 @@ export const balanceMiddleware = createMiddleware( const authToken = c.get('authToken'); // authMiddleware runs before this, so authToken should always be set for /trpc/* routes - if (!authToken) { + if (!authToken || !userId) { return buildTrpcErrorResponse(401, 'Missing auth token', procedureName); } - if ( - (procedureName === 'sendMessageV2' || - procedureName === 'initiateFromKilocodeSessionV2' || - procedureName === 'send') && - sessionId && - userId - ) { - try { - const access = await requireCurrentSessionAccess({ - env: c.env, - kiloUserId: userId, - cloudAgentSessionId: sessionId, - expectedOrganizationId: orgId, - }); - c.set('validatedSessionAccess', { - kiloUserId: userId, - cloudAgentSessionId: sessionId, - ...access, - }); - orgId = access.organizationId ?? undefined; - } catch (error) { - const response = projectSessionAccessHttpError(error); - return buildTrpcErrorResponse(response.status, await response.text(), procedureName); - } + let billing; + try { + billing = await preflightCloudAgentModelBilling({ + env: c.env, + userId, + authToken, + procedure: procedureName, + body, + }); + } catch (error) { + const response = projectBillingPreflightError(error); + return buildTrpcErrorResponse(response.status, response.message, procedureName); + } + + if (billing.validatedSessionAccess) { + c.set('validatedSessionAccess', billing.validatedSessionAccess); + } + + if (billing.classification !== 'balance-required') { + await next(); + return; } - // Use balance-only validation since auth was already done by authMiddleware - const validationResult = await validateBalanceOnly(authToken, orgId, c.env); - if (!validationResult.success) { + if (hasInsufficientBalance(billing)) { logger - .withFields({ status: validationResult.status, procedure: procedureName }) + .withFields({ status: 402, procedure: procedureName }) .warn('Pre-flight balance validation failed for V2 mutation'); - return buildTrpcErrorResponse( - validationResult.status, - validationResult.message, - procedureName - ); + return buildTrpcErrorResponse(402, INSUFFICIENT_CREDITS_MESSAGE, procedureName); } await next(); diff --git a/services/cloud-agent-next/src/model-billing-preflight.test.ts b/services/cloud-agent-next/src/model-billing-preflight.test.ts new file mode 100644 index 0000000000..2c5e6aba21 --- /dev/null +++ b/services/cloud-agent-next/src/model-billing-preflight.test.ts @@ -0,0 +1,201 @@ +import { TRPCError } from '@trpc/server'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { CloudAgentSessionState } from './persistence/types.js'; +import type { Env } from './types.js'; + +const { + checkCloudAgentAdmissionMock, + getPgDbMock, + assertOrganizationMembershipMock, + requireCurrentSessionAccessMock, + requireSessionMetadataMock, +} = vi.hoisted(() => ({ + checkCloudAgentAdmissionMock: vi.fn(), + getPgDbMock: vi.fn(), + assertOrganizationMembershipMock: vi.fn(), + requireCurrentSessionAccessMock: vi.fn(), + requireSessionMetadataMock: vi.fn(), +})); + +vi.mock('./cloud-agent-admission.js', () => ({ + checkCloudAgentAdmission: checkCloudAgentAdmissionMock, +})); +vi.mock('./db/pg.js', () => ({ getPgDb: getPgDbMock })); +vi.mock('./session-access.js', () => ({ + assertOrganizationMembership: assertOrganizationMembershipMock, + requireCurrentSessionAccess: requireCurrentSessionAccessMock, +})); +vi.mock('./session/model-preflight.js', () => ({ + requireSessionMetadata: requireSessionMetadataMock, +})); + +const { preflightCloudAgentModelBilling } = await import('./model-billing-preflight.js'); + +const env = {} as Env; +const authToken = 'jwt-token'; +const db = {}; +const metadata: CloudAgentSessionState = { + metadataSchemaVersion: 2, + identity: { sessionId: 'agent-existing', userId: 'user-1' }, + auth: { kilocodeToken: 'stored-token' }, + agent: { model: 'stored/model' }, + lifecycle: { version: 1, timestamp: 1 }, +}; + +describe('Cloud Agent model billing preflight', () => { + beforeEach(() => { + vi.clearAllMocks(); + getPgDbMock.mockReturnValue(db); + checkCloudAgentAdmissionMock.mockResolvedValue({ + classification: 'balance-required', + balance: 5, + isDepleted: false, + }); + requireCurrentSessionAccessMock.mockResolvedValue({ + kiloSessionId: 'ses_12345678901234567890123456', + organizationId: 'org-current', + }); + requireSessionMetadataMock.mockResolvedValue(metadata); + }); + + it.each([ + ['prepareSession', { model: 'prepared/model' }, 'prepared/model'], + ['start', { agent: { model: 'started/model' } }, 'started/model'], + ['initiateFromKilocodeSessionV2', { cloudAgentSessionId: 'agent-existing' }, 'stored/model'], + ['sendMessageV2', { cloudAgentSessionId: 'agent-existing', model: 'v2/model' }, 'v2/model'], + [ + 'sendMessageV2', + { + cloudAgentSessionId: 'agent-existing', + payload: { type: 'command', command: 'compact' }, + model: 'caller/free-model', + }, + 'stored/model', + ], + [ + 'send', + { cloudAgentSessionId: 'agent-existing', agent: { model: 'send/model' } }, + 'send/model', + ], + [ + 'kilo.prompt_async', + { cloudAgentSessionId: 'agent-existing', agent: { model: 'sdk/model' } }, + 'sdk/model', + ], + ] as const)('resolves the effective model for %s', async (procedure, body, modelId) => { + await preflightCloudAgentModelBilling({ env, userId: 'user-1', authToken, procedure, body }); + + expect(checkCloudAgentAdmissionMock).toHaveBeenCalledWith({ + env, + token: authToken, + modelId, + owner: + procedure === 'prepareSession' || procedure === 'start' + ? { userId: 'user-1' } + : { + organizationId: 'org-current', + }, + }); + }); + + it('surfaces the balance returned by the admission check', async () => { + checkCloudAgentAdmissionMock.mockResolvedValue({ + classification: 'balance-required', + balance: 0, + isDepleted: true, + }); + + const result = await preflightCloudAgentModelBilling({ + env, + userId: 'user-1', + authToken, + procedure: 'prepareSession', + body: { model: 'prepared/model' }, + }); + + expect(result).toMatchObject({ + classification: 'balance-required', + balance: 0, + isDepleted: true, + owner: { userId: 'user-1' }, + }); + }); + + it('uses current session access rather than a caller-supplied organization', async () => { + const result = await preflightCloudAgentModelBilling({ + env, + userId: 'user-1', + authToken, + procedure: 'send', + body: { cloudAgentSessionId: 'agent-existing', kilocodeOrganizationId: 'org-requested' }, + }); + + expect(requireCurrentSessionAccessMock).toHaveBeenCalledWith({ + env, + kiloUserId: 'user-1', + cloudAgentSessionId: 'agent-existing', + expectedOrganizationId: 'org-requested', + validatedSessionAccess: undefined, + }); + expect(result.owner).toEqual({ organizationId: 'org-current' }); + expect(result.validatedSessionAccess).toMatchObject({ organizationId: 'org-current' }); + }); + + it('checks membership before the admission call for an organization model', async () => { + await preflightCloudAgentModelBilling({ + env, + userId: 'user-1', + authToken, + procedure: 'start', + body: { + agent: { model: 'org/model' }, + options: { kilocodeOrganizationId: 'org-requested' }, + }, + }); + + expect(assertOrganizationMembershipMock).toHaveBeenCalledWith(db, 'user-1', 'org-requested'); + expect(checkCloudAgentAdmissionMock).toHaveBeenCalledWith({ + env, + token: authToken, + modelId: 'org/model', + owner: { organizationId: 'org-requested' }, + }); + expect(assertOrganizationMembershipMock.mock.invocationCallOrder[0]).toBeLessThan( + checkCloudAgentAdmissionMock.mock.invocationCallOrder[0] + ); + }); + + it('does not run the admission call after a membership denial', async () => { + assertOrganizationMembershipMock.mockRejectedValue( + new TRPCError({ code: 'FORBIDDEN', message: 'You do not have access to this organization' }) + ); + + await expect( + preflightCloudAgentModelBilling({ + env, + userId: 'user-1', + authToken, + procedure: 'prepareSession', + body: { model: 'org/model', kilocodeOrganizationId: 'org-requested' }, + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(checkCloudAgentAdmissionMock).not.toHaveBeenCalled(); + }); + + it('fails before the admission call when existing session metadata is missing', async () => { + requireSessionMetadataMock.mockRejectedValue( + new TRPCError({ code: 'NOT_FOUND', message: 'Session not found' }) + ); + + await expect( + preflightCloudAgentModelBilling({ + env, + userId: 'user-1', + authToken, + procedure: 'send', + body: { cloudAgentSessionId: 'agent-existing' }, + }) + ).rejects.toMatchObject({ code: 'NOT_FOUND' }); + expect(checkCloudAgentAdmissionMock).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/model-billing-preflight.ts b/services/cloud-agent-next/src/model-billing-preflight.ts new file mode 100644 index 0000000000..fad967ea76 --- /dev/null +++ b/services/cloud-agent-next/src/model-billing-preflight.ts @@ -0,0 +1,188 @@ +import { + checkCloudAgentAdmission, + type CloudAgentModelBilling, + type CloudAgentModelBillingOwner, +} from './cloud-agent-admission.js'; +import { TRPCError } from '@trpc/server'; +import { getPgDb } from './db/pg.js'; +import { assertOrganizationMembership, requireCurrentSessionAccess } from './session-access.js'; +import { requireSessionMetadata } from './session/model-preflight.js'; +import type { Env, ValidatedSessionAccess } from './types.js'; + +type ModelBillingProcedure = + | 'prepareSession' + | 'start' + | 'initiateFromKilocodeSessionV2' + | 'sendMessageV2' + | 'send' + | 'kilo.prompt_async'; + +type ModelBillingPreflightInput = { + env: Env; + userId: string; + authToken: string; + procedure: string; + body: unknown; + validatedSessionAccess?: ValidatedSessionAccess; +}; + +export type ModelBillingPreflightResult = { + classification: CloudAgentModelBilling; + /** Owner balance; only populated when `classification === 'balance-required'`. */ + balance: number | null; + isDepleted: boolean | null; + owner: CloudAgentModelBillingOwner; + validatedSessionAccess?: ValidatedSessionAccess; +}; + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function isModelBillingProcedure(procedure: string): procedure is ModelBillingProcedure { + return ( + procedure === 'prepareSession' || + procedure === 'start' || + procedure === 'initiateFromKilocodeSessionV2' || + procedure === 'sendMessageV2' || + procedure === 'send' || + procedure === 'kilo.prompt_async' + ); +} + +function stringProperty(value: Record, property: string): string | undefined { + const candidate = value[property]; + return typeof candidate === 'string' ? candidate : undefined; +} + +function nestedRecord( + value: Record, + property: string +): Record | undefined { + const candidate = value[property]; + return isRecord(candidate) ? candidate : undefined; +} + +function requiredModel(model: string | undefined): string { + if (!model) { + throw new TRPCError({ + code: 'BAD_REQUEST', + message: 'Model is required for billing admission', + }); + } + return model; +} + +function requestedOrganizationId( + procedure: ModelBillingProcedure, + body: Record +): string | undefined { + if (procedure === 'start') { + return stringProperty(nestedRecord(body, 'options') ?? {}, 'kilocodeOrganizationId'); + } + return stringProperty(body, 'kilocodeOrganizationId'); +} + +function submittedModel( + procedure: 'prepareSession' | 'start', + body: Record +): string { + if (procedure === 'prepareSession') return requiredModel(stringProperty(body, 'model')); + return requiredModel(stringProperty(nestedRecord(body, 'agent') ?? {}, 'model')); +} + +function sessionId(body: Record): string { + const value = stringProperty(body, 'cloudAgentSessionId'); + if (!value) { + throw new TRPCError({ code: 'BAD_REQUEST', message: 'Cloud Agent session ID is required' }); + } + return value; +} + +function requestedExistingSessionModel( + procedure: Exclude< + ModelBillingProcedure, + 'prepareSession' | 'start' | 'initiateFromKilocodeSessionV2' + >, + body: Record +): string | undefined { + switch (procedure) { + case 'sendMessageV2': { + const payload = nestedRecord(body, 'payload'); + if (payload?.['type'] === 'prompt') { + return stringProperty(payload, 'model'); + } + return payload ? undefined : stringProperty(body, 'model'); + } + case 'send': + case 'kilo.prompt_async': + return stringProperty(nestedRecord(body, 'agent') ?? {}, 'model'); + } +} + +export async function preflightCloudAgentModelBilling( + input: ModelBillingPreflightInput +): Promise { + if (!isModelBillingProcedure(input.procedure)) { + throw new TRPCError({ code: 'BAD_REQUEST', message: 'Unsupported billing procedure' }); + } + const procedure = input.procedure; + if (!isRecord(input.body)) { + throw new TRPCError({ code: 'BAD_REQUEST', message: 'Invalid request body' }); + } + + if (procedure === 'prepareSession' || procedure === 'start') { + const organizationId = requestedOrganizationId(procedure, input.body); + if (organizationId) { + await assertOrganizationMembership(getPgDb(input.env), input.userId, organizationId); + } + const owner: CloudAgentModelBillingOwner = organizationId + ? { organizationId } + : { userId: input.userId }; + const admission = await checkCloudAgentAdmission({ + env: input.env, + token: input.authToken, + modelId: submittedModel(procedure, input.body), + owner, + }); + return { ...admission, owner }; + } + + const cloudAgentSessionId = sessionId(input.body); + const access = await requireCurrentSessionAccess({ + env: input.env, + kiloUserId: input.userId, + cloudAgentSessionId, + expectedOrganizationId: requestedOrganizationId(procedure, input.body), + validatedSessionAccess: input.validatedSessionAccess, + }); + const metadata = await requireSessionMetadata({ + env: input.env, + userId: input.userId, + cloudAgentSessionId, + procedure, + }); + const modelId = requiredModel( + procedure === 'initiateFromKilocodeSessionV2' + ? metadata.agent?.model + : (requestedExistingSessionModel(procedure, input.body) ?? metadata.agent?.model) + ); + const owner: CloudAgentModelBillingOwner = access.organizationId + ? { organizationId: access.organizationId } + : { userId: input.userId }; + const admission = await checkCloudAgentAdmission({ + env: input.env, + token: input.authToken, + modelId, + owner, + }); + return { + ...admission, + owner, + validatedSessionAccess: { + kiloUserId: input.userId, + cloudAgentSessionId, + ...access, + }, + }; +} diff --git a/services/cloud-agent-next/src/router/handlers/session-start.ts b/services/cloud-agent-next/src/router/handlers/session-start.ts index 56cfae0f34..21e617fd00 100644 --- a/services/cloud-agent-next/src/router/handlers/session-start.ts +++ b/services/cloud-agent-next/src/router/handlers/session-start.ts @@ -11,10 +11,7 @@ * session ownership state is created. */ import { protectedProcedure } from '../auth.js'; -import { TRPCError } from '@trpc/server'; -import { organization_memberships } from '@kilocode/db/schema'; import type { WorkerDb } from '@kilocode/db/client'; -import { and, eq } from 'drizzle-orm'; import { logger, withLogTags } from '../../logger.js'; import { getPgDb } from '../../db/pg.js'; import type * as z from 'zod'; @@ -28,6 +25,7 @@ import { import type { SessionCreateRequest } from '../../session/session-requests.js'; import { assertKiloModelAvailable } from '../../model-validation.js'; import { assertBitbucketRepositoryAccessBeforeSessionCreation } from '../../session/validate-repository-access.js'; +import { assertOrganizationMembership } from '../../session-access.js'; type SessionStartHandlers = { start: typeof startSessionHandler; @@ -91,30 +89,6 @@ function startInputToSessionCreateRequest( }; } -async function assertOrganizationMembership( - db: WorkerDb, - userId: string, - organizationId: string -): Promise { - const [membership] = await db - .select({ id: organization_memberships.id }) - .from(organization_memberships) - .where( - and( - eq(organization_memberships.organization_id, organizationId), - eq(organization_memberships.kilo_user_id, userId) - ) - ) - .limit(1); - - if (!membership) { - throw new TRPCError({ - code: 'FORBIDDEN', - message: 'You do not have access to this organization', - }); - } -} - const startSessionHandler = protectedProcedure .input(StartSessionInput) .output(StartSessionOutput) diff --git a/services/cloud-agent-next/src/session-access.ts b/services/cloud-agent-next/src/session-access.ts index 6a67f07f10..30ea61e040 100644 --- a/services/cloud-agent-next/src/session-access.ts +++ b/services/cloud-agent-next/src/session-access.ts @@ -1,9 +1,12 @@ import { getWorkerDb } from '@kilocode/db/client'; +import { organization_memberships } from '@kilocode/db/schema'; import { TRPCError } from '@trpc/server'; +import type { WorkerDb } from '@kilocode/db/client'; import { queryAccessibleCloudAgentSession, type AccessibleCloudAgentSession, } from '@kilocode/worker-utils/cloud-agent-session-access'; +import { and, eq } from 'drizzle-orm'; import type { Env, ValidatedSessionAccess } from './types.js'; type CurrentSessionAccessRequest = { @@ -15,6 +18,30 @@ type CurrentSessionAccessRequest = { validatedSessionAccess?: ValidatedSessionAccess; }; +export async function assertOrganizationMembership( + db: WorkerDb, + userId: string, + organizationId: string +): Promise { + const [membership] = await db + .select({ id: organization_memberships.id }) + .from(organization_memberships) + .where( + and( + eq(organization_memberships.organization_id, organizationId), + eq(organization_memberships.kilo_user_id, userId) + ) + ) + .limit(1); + + if (!membership) { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'You do not have access to this organization', + }); + } +} + function matchesExpectedSession( session: AccessibleCloudAgentSession, request: CurrentSessionAccessRequest diff --git a/services/cloud-agent-next/src/session-prepare.test.ts b/services/cloud-agent-next/src/session-prepare.test.ts index 339ec6f6c2..83a501c140 100644 --- a/services/cloud-agent-next/src/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session-prepare.test.ts @@ -158,7 +158,6 @@ function createInternalApiContext(options: { authToken?: string | null; internalApiSecret?: string | null; requestInternalApiKey?: string | null; - skipBalanceCheck?: boolean; doStub?: ReturnType; getBitbucketToken?: ReturnType; managedScmContainmentOrgIds?: string; @@ -191,9 +190,6 @@ function createInternalApiContext(options: { if (effectiveRequestInternalApiKey !== null) { headers.set('x-internal-api-key', effectiveRequestInternalApiKey); } - if (options.skipBalanceCheck) { - headers.set('x-skip-balance-check', 'true'); - } return { userId: effectiveUserId, @@ -1190,13 +1186,12 @@ describe('start endpoint', () => { assertKiloModelAvailableMock.mockResolvedValue(undefined); }); - it('rejects non-member organization profile resolution when balance validation is skipped', async () => { + it('rejects non-member organization profile resolution before initial admission', async () => { organizationMembershipLimitMock.mockResolvedValueOnce([]); const doStub = createMockDOStub(); - const context = createInternalApiContext({ doStub, skipBalanceCheck: true }); + const context = createInternalApiContext({ doStub }); const caller = appRouter.createCaller(context); - expect(context.request.headers.get('x-skip-balance-check')).toBe('true'); await expect( caller.start({ message: { prompt: 'Attempt organization profile access' }, diff --git a/services/cloud-agent-next/src/session/model-preflight.ts b/services/cloud-agent-next/src/session/model-preflight.ts index f4074e7443..8aa9531bdc 100644 --- a/services/cloud-agent-next/src/session/model-preflight.ts +++ b/services/cloud-agent-next/src/session/model-preflight.ts @@ -14,7 +14,7 @@ type ExistingPromptModelPreflightInput = StoredSessionPreflightInput & { requestedModel?: string; }; -async function requireSessionMetadata( +export async function requireSessionMetadata( input: StoredSessionPreflightInput ): Promise { const metadata = await fetchSessionMetadata(input.env, input.userId, input.cloudAgentSessionId); diff --git a/services/cloud-agent-next/test/e2e/client.ts b/services/cloud-agent-next/test/e2e/client.ts index 582d40b5e2..8823d55961 100644 --- a/services/cloud-agent-next/test/e2e/client.ts +++ b/services/cloud-agent-next/test/e2e/client.ts @@ -90,9 +90,6 @@ export async function trpcCall( const headers: Record = { 'Content-Type': 'application/json', Authorization: `Bearer ${mintApiToken(config.user, config.nextAuthSecret)}`, - // cloud-agent-client.ts sends this for App Builder callers; it cleanly - // skips dev billing checks and is safe to always send from the driver. - 'x-skip-balance-check': 'true', }; if (opts?.internalApiSecret) { headers['x-internal-api-key'] = opts.internalApiSecret; diff --git a/services/code-review-infra/src/code-review-orchestrator.ts b/services/code-review-infra/src/code-review-orchestrator.ts index c65ad7fcca..257ea446fc 100644 --- a/services/code-review-infra/src/code-review-orchestrator.ts +++ b/services/code-review-infra/src/code-review-orchestrator.ts @@ -888,7 +888,6 @@ export class CodeReviewOrchestrator extends DurableObject { id: string; userId: string; }; - skipBalanceCheck?: boolean; previousCloudAgentSessionId?: string; repositorySize?: string | null; runReviewDelayMs?: number; @@ -914,7 +913,6 @@ export class CodeReviewOrchestrator extends DurableObject { owner: params.owner, status: 'queued', updatedAt: new Date().toISOString(), - skipBalanceCheck: params.skipBalanceCheck, previousCloudAgentSessionId: params.sessionInput.platform === 'bitbucket' ? undefined @@ -1022,7 +1020,6 @@ export class CodeReviewOrchestrator extends DurableObject { authToken: this.state.authToken, sessionInput: this.state.sessionInput, owner: this.state.owner, - skipBalanceCheck: this.state.skipBalanceCheck, previousCloudAgentSessionId: undefined, repositorySize: this.state.repositorySize, runReviewDelayMs: retryDelayMs, @@ -1153,9 +1150,6 @@ export class CodeReviewOrchestrator extends DurableObject { Authorization: `Bearer ${this.state.authToken}`, 'x-internal-api-key': this.env.INTERNAL_API_SECRET, }; - if (this.state.skipBalanceCheck) { - internalHeaders['x-skip-balance-check'] = 'true'; - } // Step 1: Prepare session with callback target const callbackTarget = await callbackTargetForAttempt( @@ -1197,7 +1191,6 @@ export class CodeReviewOrchestrator extends DurableObject { reviewId: this.state.reviewId, callbackUrl: callbackTarget.url, createdOnPlatform: prepareInput.createdOnPlatform, - skipBalanceCheck: this.state.skipBalanceCheck, githubCloudReviewSkill: { attached: githubCloudReviewSkillAttached, name: GITHUB_CLOUD_REVIEW_SKILL_NAME, @@ -1237,9 +1230,6 @@ export class CodeReviewOrchestrator extends DurableObject { const userHeaders: Record = { Authorization: `Bearer ${this.state.authToken}`, }; - if (this.state.skipBalanceCheck) { - userHeaders['x-skip-balance-check'] = 'true'; - } console.log('[CodeReviewOrchestrator] Calling initiateFromKilocodeSessionV2', { reviewId: this.state.reviewId, @@ -1357,9 +1347,6 @@ export class CodeReviewOrchestrator extends DurableObject { const userHeaders: Record = { Authorization: `Bearer ${this.state.authToken}`, }; - if (this.state.skipBalanceCheck) { - userHeaders['x-skip-balance-check'] = 'true'; - } let health: CloudAgentSessionHealthOutput; try { @@ -1398,9 +1385,6 @@ export class CodeReviewOrchestrator extends DurableObject { Authorization: `Bearer ${this.state.authToken}`, 'x-internal-api-key': this.env.INTERNAL_API_SECRET, }; - if (this.state.skipBalanceCheck) { - internalHeaders['x-skip-balance-check'] = 'true'; - } // Step 1: Update callback target via updateSession (internal-only endpoint). // callbackTarget must be set through an internal procedure, not the diff --git a/services/code-review-infra/src/index.ts b/services/code-review-infra/src/index.ts index 65fdc9e4aa..d415653900 100644 --- a/services/code-review-infra/src/index.ts +++ b/services/code-review-infra/src/index.ts @@ -100,7 +100,6 @@ app.post('/review', async (c: Context) => { authToken: body.authToken, sessionInput: body.sessionInput, owner: body.owner, - skipBalanceCheck: body.skipBalanceCheck, previousCloudAgentSessionId: body.previousCloudAgentSessionId, repositorySize: body.repositorySize, }), diff --git a/services/code-review-infra/src/types.ts b/services/code-review-infra/src/types.ts index bef3c83666..e700968651 100644 --- a/services/code-review-infra/src/types.ts +++ b/services/code-review-infra/src/types.ts @@ -71,7 +71,6 @@ export interface CodeReview { totalTokensOut?: number; /** Accumulated cost in dollars across all LLM calls */ totalCost?: number; - skipBalanceCheck?: boolean; // Skip balance validation in cloud agent (for OSS sponsorship) /** Cloud-agent session ID from a previous completed review, for session continuation */ previousCloudAgentSessionId?: string; sandboxRetryAttempted?: boolean; @@ -139,7 +138,6 @@ export interface CodeReviewRequest { authToken: string; sessionInput: SessionInput; owner: Owner; - skipBalanceCheck?: boolean; /** Cloud-agent session ID from a previous completed review, for session continuation */ previousCloudAgentSessionId?: string; /** Provider-reported repository storage size, formatted for log correlation. */ diff --git a/services/webhook-agent-ingest/src/queue-consumer.test.ts b/services/webhook-agent-ingest/src/queue-consumer.test.ts index 1ccf684f29..478a985625 100644 --- a/services/webhook-agent-ingest/src/queue-consumer.test.ts +++ b/services/webhook-agent-ingest/src/queue-consumer.test.ts @@ -304,6 +304,7 @@ describe('handleWebhookDeliveryBatch Cloud Agent callback target', () => { await handleWebhookDeliveryBatch(batch, env); + expect(prepareRequests[0]?.headers.get('x-skip-balance-check')).toBeNull(); const prepareBody = await prepareRequests[0]?.json(); const expectedToken = await deriveCallbackToken({ secret: callbackTokenSecret, diff --git a/services/webhook-agent-ingest/src/queue-consumer.ts b/services/webhook-agent-ingest/src/queue-consumer.ts index 04b191762f..9168916d19 100644 --- a/services/webhook-agent-ingest/src/queue-consumer.ts +++ b/services/webhook-agent-ingest/src/queue-consumer.ts @@ -457,7 +457,6 @@ async function processWebhookMessage( 'Content-Type': 'application/json', Authorization: `Bearer ${token}`, 'x-internal-api-key': internalApiSecret, - 'x-skip-balance-check': 'true', }, body: JSON.stringify(prepareSessionBody), }) @@ -545,7 +544,6 @@ async function processWebhookMessage( 'Content-Type': 'application/json', Authorization: `Bearer ${token}`, 'x-internal-api-key': internalApiSecret, - 'x-skip-balance-check': 'true', }, body: JSON.stringify({ cloudAgentSessionId }), })