refactor(cloud-agent-next): gate billing by model classification

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',
+    });
+  });
+});

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<NextResponse<AdmissionResponse>> {
+  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 });
+  }
+}

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
       </SetPageTitle>
       <MobileSidebarToggle />
       <div className="w-full max-w-2xl space-y-4">
-        {/* Insufficient balance banner */}
-        {hasInsufficientBalance && eligibilityData && !hasLimitedAccess && (
+        {balanceNotice === 'zero-credit' && eligibilityData && (
           <InsufficientBalanceBanner
             balance={eligibilityData.balance}
             organizationId={organizationId}
-            content={{ type: 'productName', productName: 'Cloud Agent' }}
+            colorScheme="info"
+            content={{
+              type: 'custom',
+              title: 'Free and BYOK models remain available',
+              description:
+                'Use free or eligible BYOK models in Cloud Agent. Add Credits to use other models.',
+              compactActionText: 'Add Credits',
+            }}
           />
         )}
 
-        {/* Free-models-available banner when balance is low but free models are usable */}
-        {hasLimitedAccess && eligibilityData && (
+        {balanceNotice === 'low-funds' && eligibilityData && (
           <InsufficientBalanceBanner
             balance={eligibilityData.balance}
             organizationId={organizationId}
             colorScheme="info"
             content={{
               type: 'custom',
-              title: 'Free Models Available',
-              description:
-                'You can use free models in Cloud Agent. Add credits to unlock all models.',
-              compactActionText: 'Add credits to unlock all models',
+              title: 'Credits are below $1',
+              description: 'All models remain available while your balance is positive.',
+              compactActionText: 'Add Credits',
             }}
           />
         )}

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');
+  });
+});

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<T extends BalanceRestrictedModelOption>(
+  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';
+}