diff --git a/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx b/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx index 6a4f98f219..b9fa73dac0 100644 --- a/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx +++ b/apps/web/src/components/cloud-agent-next/NewSessionPanel.tsx @@ -162,6 +162,7 @@ export function NewSessionPanel({ organizationId, isDevcontainerAvailable }: New : personalEligibilityQuery.isPending; const hasInsufficientBalance = !isEligibilityLoading && eligibilityData && !eligibilityData.isEligible; + const hasLimitedAccess = !isEligibilityLoading && eligibilityData?.accessLevel === 'limited'; // --------------------------------------------------------------------------- // Models @@ -171,20 +172,19 @@ export function NewSessionPanel({ organizationId, isDevcontainerAvailable }: New const allModels = modelsData?.data || []; - const modelOptions = useMemo( - () => - appendCloudAgentNextLocalTestModel( - allModels.map(model => ({ - id: model.id, - name: model.name, - isFree: model.isFree, - mayTrainOnYourPrompts: model.mayTrainOnYourPrompts, - hasUserByokAvailable: model.hasUserByokAvailable, - variants: model.opencode?.variants ? Object.keys(model.opencode.variants) : undefined, - })) - ), - [allModels] - ); + const modelOptions = useMemo(() => { + const options = allModels.map(model => ({ + id: model.id, + name: model.name, + isFree: model.isFree, + mayTrainOnYourPrompts: model.mayTrainOnYourPrompts, + hasUserByokAvailable: model.hasUserByokAvailable, + 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]); // --------------------------------------------------------------------------- // Form state @@ -872,12 +872,21 @@ 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 && + (!hasInsufficientBalance || limitedAccessModelIsAllowed) && !attachmentUpload.hasUploadingAttachments; const handleStartSession = useCallback(async () => { @@ -1158,7 +1167,7 @@ export function NewSessionPanel({ organizationId, isDevcontainerAvailable }: New
{/* Insufficient balance banner */} - {hasInsufficientBalance && eligibilityData && ( + {hasInsufficientBalance && eligibilityData && !hasLimitedAccess && ( )} + {/* Free-models-available banner when balance is low but free models are usable */} + {hasLimitedAccess && eligibilityData && ( + + )} + {/* Textarea + model toolbar container */}
= minBalance ? 'full' : 'limited'; + return { + balance, + minBalance, + accessLevel, + isEligible: accessLevel === 'full', + }; +} 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 new file mode 100644 index 0000000000..97f6a5365d --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/balance-check-eligibility.test.ts @@ -0,0 +1,143 @@ +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.clearAllMocks(); + 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 new file mode 100644 index 0000000000..fd3d6c0ac5 --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/balance-check-eligibility.ts @@ -0,0 +1,66 @@ +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/cloud-agent-client.test.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts new file mode 100644 index 0000000000..6e046ac84f --- /dev/null +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, it, jest, beforeEach } from '@jest/globals'; + +jest.mock('@/lib/dotenvx', () => ({ + getEnvVariable: jest.fn(() => 'http://cloud-agent-next'), +})); + +jest.mock('@/lib/config.server', () => ({ + INTERNAL_API_SECRET: 'test-secret', +})); + +jest.mock('@trpc/client', () => ({ + createTRPCClient: jest.fn(() => ({})), + httpLink: jest.fn(), + 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'); + +const { + createCloudAgentNextClient: mockCreateCloudAgentNextClient, + createAppBuilderCloudAgentNextClient: mockCreateAppBuilderCloudAgentNextClient, + createCloudAgentNextClientForModel, +} = clientModule; + +beforeEach(() => { + mockCreateCloudAgentNextClient.mockClear(); + mockCreateAppBuilderCloudAgentNextClient.mockClear(); +}); + +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(); + }); + + it('returns the AppBuilder client when the model is free', () => { + const result = createCloudAgentNextClientForModel('token', { + isFree: true, + hasUserByokAvailable: false, + }); + expect(result).toEqual({ marker: 'appbuilder' }); + expect(mockCreateAppBuilderCloudAgentNextClient).toHaveBeenCalledWith('token'); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + }); + + 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(); + }); +}); 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 7c85a75f57..e04c84fba0 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 @@ -809,3 +809,28 @@ export function createAppBuilderCloudAgentNextClient(authToken: string): CloudAg 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); +} diff --git a/apps/web/src/routers/app-builder-router.ts b/apps/web/src/routers/app-builder-router.ts index ae8cd5eec4..52241b769c 100644 --- a/apps/web/src/routers/app-builder-router.ts +++ b/apps/web/src/routers/app-builder-router.ts @@ -12,6 +12,7 @@ import { } from '@/routers/app-builder/schemas'; import { getBalanceForUser } from '@/lib/user/balance'; import { MIN_BALANCE_FOR_APP_BUILDER } from '@/lib/app-builder/constants'; +import { buildAccessLevelEligibility } from '@/lib/access-level-eligibility'; import { generateImageUploadUrl } from '@/lib/r2/cloud-agent-attachments'; export const appBuilderRouter = createTRPCRouter({ @@ -109,19 +110,7 @@ export const appBuilderRouter = createTRPCRouter({ */ checkEligibility: baseProcedure.query(async ({ ctx }) => { const { balance } = await getBalanceForUser(ctx.user); - - // Determine access level based on balance - // Currently only 'full' or 'limited' - change to 'blocked' if we need to restrict entirely - const accessLevel: 'full' | 'limited' | 'blocked' = - balance >= MIN_BALANCE_FOR_APP_BUILDER ? 'full' : 'limited'; - - return { - balance, - minBalance: MIN_BALANCE_FOR_APP_BUILDER, - accessLevel, - // Keep isEligible for backwards compatibility (true if full access) - isEligible: accessLevel === 'full', - }; + return buildAccessLevelEligibility(balance, MIN_BALANCE_FOR_APP_BUILDER); }), /** diff --git a/apps/web/src/routers/cloud-agent-next-eligibility.ts b/apps/web/src/routers/cloud-agent-next-eligibility.ts index 376a7e269e..6d705bcf64 100644 --- a/apps/web/src/routers/cloud-agent-next-eligibility.ts +++ b/apps/web/src/routers/cloud-agent-next-eligibility.ts @@ -1,15 +1,15 @@ +import { + buildAccessLevelEligibility, + type AccessLevel, + type AccessLevelEligibility, +} from '@/lib/access-level-eligibility'; + export const CLOUD_AGENT_NEXT_MIN_BALANCE_DOLLARS = 1; -export type CloudAgentNextEligibility = { - balance: number; - minBalance: number; - isEligible: boolean; -}; +export type CloudAgentNextAccessLevel = AccessLevel; + +export type CloudAgentNextEligibility = AccessLevelEligibility; export function buildCloudAgentNextEligibility(balance: number): CloudAgentNextEligibility { - return { - balance, - minBalance: CLOUD_AGENT_NEXT_MIN_BALANCE_DOLLARS, - isEligible: balance >= CLOUD_AGENT_NEXT_MIN_BALANCE_DOLLARS, - }; + return buildAccessLevelEligibility(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 ef25f3d86b..97fcc6a793 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -47,6 +47,20 @@ const mockCreateCloudAgentNextClient = jest.fn(() => ({ sendMessage: mockSendMessage, })); +const mockCreateCloudAgentNextClientForModel = jest.fn( + (_authToken: string, _eligibility: unknown) => ({ + prepareSession: mockPrepareSession, + sendMessage: mockSendMessage, + }) +); + +const mockComputeCloudAgentNextBalanceCheckEligibility = jest.fn< + (...args: unknown[]) => Promise<{ + isFree: boolean; + hasUserByokAvailable: boolean; + }> +>(); + const mockIsFeatureFlagEnabledOrDevelopment = jest.fn<(flagName: string, distinctId: string) => Promise>(); const mockVerifyUserOwnsSessionV2ByCloudAgentId = @@ -79,9 +93,14 @@ 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, })); @@ -139,7 +158,12 @@ let createCaller: (ctx: { user: User }) => { contentType: 'application/pdf'; contentLength: number; }) => Promise; - checkEligibility: () => Promise<{ balance: number; minBalance: number; isEligible: boolean }>; + checkEligibility: () => Promise<{ + balance: number; + minBalance: number; + isEligible: boolean; + accessLevel: 'full' | 'limited' | 'blocked'; + }>; listGitHubRepositories: (input: { forceRefresh: boolean }) => Promise; listGitLabRepositories: (input: { forceRefresh: boolean }) => Promise; }; @@ -229,9 +253,9 @@ describe('cloudAgentNextRouter helper procedures', () => { }); it.each([ - { balance: 1, isEligible: true }, - { balance: 0.99, isEligible: false }, - ])('reports eligibility for a $balance balance', async ({ balance, isEligible }) => { + { 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 }); @@ -240,6 +264,7 @@ describe('cloudAgentNextRouter helper procedures', () => { balance, minBalance: 1, isEligible, + accessLevel, }); expect(mockGetBalanceForUser).toHaveBeenCalledWith(user); expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); @@ -273,6 +298,10 @@ 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 () => { @@ -376,4 +405,77 @@ 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, + }); + + await caller.prepareSession({ + prompt: 'Test prompt', + mode: 'code', + model: 'kilo/paid-model', + githubRepo: 'acme/repo', + autoInitiate: true, + devcontainer: false, + }); + + expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', { + isFree: false, + hasUserByokAvailable: false, + }); + }); }); diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index beb2e662d3..d02ee9579e 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -2,8 +2,10 @@ 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'; @@ -135,7 +137,12 @@ export const cloudAgentNextRouter = createTRPCRouter({ } const authToken = generateCloudAgentToken(ctx.user); - const client = createCloudAgentNextClient(authToken); + const eligibility = await computeCloudAgentNextBalanceCheckEligibility({ + fromDb: db, + user: ctx.user, + modelId: input.model, + }); + const client = createCloudAgentNextClientForModel(authToken, eligibility); const { gitlabProject, githubRepo, attachments, images, ...restInput } = input; diff --git a/apps/web/src/routers/organizations/organization-app-builder-router.ts b/apps/web/src/routers/organizations/organization-app-builder-router.ts index cd85f185cd..ea381ab1e9 100644 --- a/apps/web/src/routers/organizations/organization-app-builder-router.ts +++ b/apps/web/src/routers/organizations/organization-app-builder-router.ts @@ -17,6 +17,7 @@ import { } from '@/routers/app-builder/schemas'; import { getBalanceForOrganizationUser } from '@/lib/organizations/organization-usage'; import { MIN_BALANCE_FOR_APP_BUILDER } from '@/lib/app-builder/constants'; +import { buildAccessLevelEligibility } from '@/lib/access-level-eligibility'; import { generateImageUploadUrl } from '@/lib/r2/cloud-agent-attachments'; // Input schemas with required organizationId @@ -139,19 +140,7 @@ export const organizationAppBuilderRouter = createTRPCRouter({ */ checkEligibility: organizationMemberProcedure.query(async ({ ctx, input }) => { const { balance } = await getBalanceForOrganizationUser(input.organizationId, ctx.user.id); - - // Determine access level based on balance - // Currently only 'full' or 'limited' - change to 'blocked' if we need to restrict entirely - const accessLevel: 'full' | 'limited' | 'blocked' = - balance >= MIN_BALANCE_FOR_APP_BUILDER ? 'full' : 'limited'; - - return { - balance, - minBalance: MIN_BALANCE_FOR_APP_BUILDER, - accessLevel, - // Keep isEligible for backwards compatibility (true if full access) - isEligible: accessLevel === 'full', - }; + return buildAccessLevelEligibility(balance, MIN_BALANCE_FOR_APP_BUILDER); }), /** 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 7e32e2a896..554b016234 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 @@ -59,6 +59,20 @@ const mockCreateCloudAgentNextClient = jest.fn(() => ({ sendMessage: mockSendMessage, })); +const mockCreateCloudAgentNextClientForModel = jest.fn( + (_authToken: string, _eligibility: unknown) => ({ + prepareSession: mockPrepareSession, + sendMessage: mockSendMessage, + }) +); + +const mockComputeCloudAgentNextBalanceCheckEligibility = jest.fn< + (...args: unknown[]) => Promise<{ + isFree: boolean; + hasUserByokAvailable: boolean; + }> +>(); + const mockIsFeatureFlagEnabledOrDevelopment = jest.fn<(flagName: string, distinctId: string) => Promise>(); const mockVerifyOrgOwnsSessionV2ByCloudAgentId = @@ -101,9 +115,14 @@ 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, })); @@ -195,6 +214,7 @@ let createCaller: (ctx: { user: User }) => { balance: number; minBalance: number; isEligible: boolean; + accessLevel: 'full' | 'limited' | 'blocked'; }>; listGitHubRepositories: (input: { organizationId: string; @@ -324,21 +344,28 @@ describe('organizationCloudAgentNextRouter helper procedures', () => { }); it.each([ - { balance: 1, isEligible: true }, - { balance: 0.99, isEligible: false }, - ])('reports organization eligibility for a $balance balance', async ({ balance, isEligible }) => { - mockGetBalanceForOrganizationUser.mockResolvedValue({ balance }); - const caller = createCaller({ user: { id: 'member-user', is_admin: false } as User }); + { balance: 1, isEligible: true, accessLevel: 'full' as const }, + { balance: 0.99, isEligible: false, accessLevel: 'limited' as const }, + ])( + 'reports organization eligibility for a $balance balance', + async ({ balance, isEligible, accessLevel }) => { + mockGetBalanceForOrganizationUser.mockResolvedValue({ balance }); + const caller = createCaller({ user: { id: 'member-user', is_admin: false } as User }); - await expect(caller.checkEligibility({ organizationId: ORGANIZATION_ID })).resolves.toEqual({ - balance, - minBalance: 1, - isEligible, - }); - expect(mockEnsureOrganizationAccess).toHaveBeenCalledWith('member-user', ORGANIZATION_ID); - expect(mockGetBalanceForOrganizationUser).toHaveBeenCalledWith(ORGANIZATION_ID, 'member-user'); - expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); - }); + await expect(caller.checkEligibility({ organizationId: ORGANIZATION_ID })).resolves.toEqual({ + balance, + minBalance: 1, + isEligible, + accessLevel, + }); + expect(mockEnsureOrganizationAccess).toHaveBeenCalledWith('member-user', ORGANIZATION_ID); + expect(mockGetBalanceForOrganizationUser).toHaveBeenCalledWith( + ORGANIZATION_ID, + 'member-user' + ); + expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); + } + ); it('rejects eligibility checks before reading balance when membership is denied', async () => { mockEnsureOrganizationAccess.mockImplementation(() => { @@ -442,6 +469,10 @@ 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 () => { @@ -558,6 +589,79 @@ 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 () => { + const caller = createCaller({ user: { id: 'user-paid', is_admin: false } as User }); + + await caller.prepareSession({ + organizationId: ORGANIZATION_ID, + prompt: 'Test prompt', + mode: 'code', + model: 'kilo/paid-model', + githubRepo: 'acme/repo', + autoInitiate: true, + devcontainer: false, + }); + + expect(mockCreateCloudAgentNextClientForModel).toHaveBeenCalledWith('cloud-agent-token', { + isFree: false, + hasUserByokAvailable: false, + }); + }); }); describe('organizationCloudAgentNextRouter Bitbucket repository listing', () => { 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 2e4ab5fa90..ce672ff0e7 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,8 +2,10 @@ 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'; @@ -224,7 +226,13 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ } const authToken = generateCloudAgentToken(ctx.user); - const client = createCloudAgentNextClient(authToken); + const eligibility = await computeCloudAgentNextBalanceCheckEligibility({ + fromDb: db, + user: ctx.user, + modelId: input.model, + organizationId: input.organizationId, + }); + const client = createCloudAgentNextClientForModel(authToken, eligibility); const { gitlabProject,