fix(cloud-agent): restore sends and workflow issuance with scoped tokens

apps/web/src/app/api/cloud-agent-next/balance/route.auth.test.ts

@@ -0,0 +1,178 @@
+import jwt from 'jsonwebtoken';
+import { buildModernKiloTokenPayload } from '@kilocode/worker-utils/kilo-token-policy';
+import { GET } from './route';
+import { createControlTokenForRequest } from '@/lib/auth/resource-delegation';
+import type { User } from '@kilocode/db/schema';
+import { GET as getGenericBalance } from '@/app/api/profile/balance/route';
+
+// Keep getUserFromAuth and validateAuthorizationHeader real; substitute only
+// persistence, request context, and unrelated sign-in integrations.
+const mockHeaders = jest.fn();
+const mockFindUser = jest.fn();
+const mockMembership = jest.fn();
+const mockBalance = jest.fn();
+jest.mock('next/headers', () => ({ headers: () => mockHeaders(), cookies: jest.fn() }));
+jest.mock('next-auth', () => ({
+  __esModule: true,
+  default: jest.fn(),
+  getServerSession: jest.fn(),
+}));
+jest.mock('@/lib/user', () => ({ findUserById: (...args: unknown[]) => mockFindUser(...args) }));
+jest.mock('@/lib/drizzle', () => ({
+  db: { query: { kilocode_users: { findFirst: (...args: unknown[]) => mockFindUser(...args) } } },
+  readDb: {},
+}));
+jest.mock('@/lib/organizations/organizations', () => ({
+  isOrganizationMember: (...args: unknown[]) => mockMembership(...args),
+}));
+jest.mock('@/lib/organizations/organization-usage', () => ({
+  getBalanceAndOrgSettings: (...args: unknown[]) => mockBalance(...args),
+}));
+jest.mock('@/lib/config.server', () => ({
+  NEXTAUTH_SECRET: 'balance-test-secret',
+  BLACKLIST_TLDS: [],
+  isResourceTokenIssuanceEnabled: () => true,
+}));
+jest.mock('@/lib/constants', () => ({ ORGANIZATION_ID_HEADER: 'X-KiloCode-OrganizationId' }));
+jest.mock('@/lib/dotenvx', () => ({ getEnvVariable: jest.fn() }));
+jest.mock('@/lib/blacklist-domains-config', () => ({ getBlacklistedDomains: async () => [] }));
+jest.mock('@/lib/utils.server', () => ({
+  warnExceptInTest: jest.fn(),
+  sentryLogger: () => jest.fn(),
+}));
+jest.mock('@/lib/posthog', () => ({ __esModule: true, default: jest.fn() }));
+jest.mock('@/lib/auth/magic-link-tokens', () => ({}));
+jest.mock('@/lib/impact/debug', () => ({}));
+jest.mock('@/lib/impact/referral', () => ({}));
+jest.mock('@/lib/organizations/trial-utils', () => ({}));
+jest.mock('@/lib/organizations/organization-seats', () => ({}));
+jest.mock('@/lib/organizations/sales-demo', () => ({}));
+jest.mock('@/lib/organizations/organization-sso-policy', () => ({}));
+jest.mock('@/lib/organizations/verified-domain-membership', () => ({}));
+jest.mock('@/lib/organizations/verified-domain-destination', () => ({}));
+jest.mock('@/lib/account-linking-session', () => ({}));
+jest.mock('@/lib/admin/admin-access-log', () => ({}));
+jest.mock('@/lib/user/sso', () => ({}));
+jest.mock('@/lib/web-session-revocation', () => ({}));
+
+const organizationId = '11111111-1111-4111-8111-111111111111';
+const user = {
+  id: 'user_123',
+  api_token_pepper: 'current-pepper',
+  google_user_email: 'test@example.com',
+};
+
+function controlToken(overrides: Record<string, unknown> = {}, secret = 'balance-test-secret') {
+  const now = Math.floor(Date.now() / 1000);
+  // The same builder and claims used by createModernControlToken.
+  const payload = buildModernKiloTokenPayload({
+    userId: user.id,
+    pepper: user.api_token_pepper,
+    env: process.env.NODE_ENV,
+    audience: 'cloud-agent-next',
+    issuedAt: now,
+    expiresAt: now + 60,
+    tokenPurpose: 'human-api',
+    credentialExchange: false,
+    extra: {
+      organizationId,
+      tokenSource: 'cloud-agent',
+      runtimeAdmission: {
+        source: 'user',
+        authorizationUserId: user.id,
+        authorizationPepper: user.api_token_pepper,
+      },
+    },
+  });
+  return jwt.sign({ ...payload, ...overrides }, secret, { algorithm: 'HS256' });
+}
+
+beforeEach(() => {
+  jest.clearAllMocks();
+  mockFindUser.mockResolvedValue(user);
+  mockMembership.mockResolvedValue(true);
+  mockBalance.mockResolvedValue({ balance: 12 });
+  jest.spyOn(console, 'warn').mockImplementation(() => {});
+});
+afterEach(() => jest.restoreAllMocks());
+
+function requestWith(token: string) {
+  mockHeaders.mockResolvedValue(new Headers({ Authorization: `Bearer ${token}` }));
+}
+
+it.each(['human-api', 'device-access'])(
+  'accepts issuer-shaped %s control only at the dedicated endpoint',
+  async tokenPurpose => {
+    requestWith(
+      controlToken({
+        tokenPurpose,
+        ...(tokenPurpose === 'device-access' ? { deviceSessionId: 'device_123' } : {}),
+      })
+    );
+    expect((await GET()).status).toBe(200);
+    expect(mockMembership).toHaveBeenCalledWith(organizationId, user.id, expect.anything());
+    expect(mockBalance).toHaveBeenCalledWith(organizationId, user);
+    mockBalance.mockClear();
+    expect((await getGenericBalance()).status).toBe(401);
+    expect(mockBalance).not.toHaveBeenCalled();
+  }
+);
+
+it.each([
+  { name: 'wrong audience', overrides: { aud: 'kilo-api' }, secret: 'balance-test-secret' },
+  { name: 'wrong signature', overrides: {}, secret: 'wrong-secret' },
+  {
+    name: 'wrong pepper',
+    overrides: { apiTokenPepper: 'stale-pepper' },
+    secret: 'balance-test-secret',
+  },
+  { name: 'expired', overrides: { exp: 1 }, secret: 'balance-test-secret' },
+])('rejects $name before balance access', async ({ overrides, secret }) => {
+  requestWith(controlToken(overrides, secret));
+  expect((await GET()).status).toBe(401);
+  expect(mockBalance).not.toHaveBeenCalled();
+});
+
+it('rejects an organization the user cannot access', async () => {
+  mockMembership.mockResolvedValue(false);
+  requestWith(controlToken());
+  expect((await GET()).status).toBe(403);
+  expect(mockBalance).not.toHaveBeenCalled();
+});
+
+it('rejects malformed bearer input', async () => {
+  requestWith('not-a-jwt');
+  expect((await GET()).status).toBe(401);
+  expect(mockBalance).not.toHaveBeenCalled();
+});
+
+it('accepts the actual control issuer output with its environment and pepper claims', async () => {
+  const now = Math.floor(Date.now() / 1000);
+  const principal = jwt.sign(
+    buildModernKiloTokenPayload({
+      userId: user.id,
+      pepper: user.api_token_pepper,
+      env: process.env.NODE_ENV,
+      audience: 'kilo-api',
+      issuedAt: now,
+      expiresAt: now + 60,
+      tokenPurpose: 'human-api',
+      credentialExchange: true,
+    }),
+    'balance-test-secret',
+    { algorithm: 'HS256' }
+  );
+  const issued = await createControlTokenForRequest(user as User, 'cloud-agent-next', {
+    headers: new Headers({ Authorization: `Bearer ${principal}` }),
+    tokenSource: 'cloud-agent',
+  });
+  expect(jwt.verify(issued.token, 'balance-test-secret')).toMatchObject({
+    aud: 'cloud-agent-next',
+    env: process.env.NODE_ENV,
+    apiTokenPepper: user.api_token_pepper,
+    credentialExchange: false,
+  });
+  requestWith(issued.token);
+  expect((await GET()).status).toBe(200);
+  expect((await getGenericBalance()).status).toBe(401);
+});

apps/web/src/app/api/cloud-agent-next/balance/route.test.ts

@@ -0,0 +1,41 @@
+import { GET } from './route';
+import { getUserFromAuth } from '@/lib/user/server';
+import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage';
+import { NextResponse } from 'next/server';
+
+jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() }));
+jest.mock('@/lib/organizations/organization-usage', () => ({
+  getBalanceAndOrgSettings: jest.fn(),
+}));
+
+const auth = jest.mocked(getUserFromAuth);
+const balance = jest.mocked(getBalanceAndOrgSettings);
+
+afterEach(() => jest.resetAllMocks());
+
+it('requires Cloud Agent audience authentication and uses the authorized organization', async () => {
+  const user = { id: 'user_123' };
+  auth.mockResolvedValue({ user, organizationId: 'org_123' } as Awaited<
+    ReturnType<typeof getUserFromAuth>
+  >);
+  balance.mockResolvedValue({ balance: 12 } as Awaited<
+    ReturnType<typeof getBalanceAndOrgSettings>
+  >);
+
+  const response = await GET();
+
+  expect(auth).toHaveBeenCalledWith({ adminOnly: false, expectedAudience: 'cloud-agent-next' });
+  expect(balance).toHaveBeenCalledWith('org_123', user);
+  expect(await response.json()).toEqual({ balance: 12, isDepleted: false });
+});
+
+it.each([401, 403])(
+  'does not read balances when authentication or organization access fails (%s)',
+  async status => {
+    const authFailedResponse = NextResponse.json({ error: 'Unauthorized' }, { status });
+    auth.mockResolvedValue({ authFailedResponse } as Awaited<ReturnType<typeof getUserFromAuth>>);
+
+    expect(await GET()).toBe(authFailedResponse);
+    expect(balance).not.toHaveBeenCalled();
+  }
+);

apps/web/src/app/api/cloud-agent-next/balance/route.ts

@@ -0,0 +1,19 @@
+import { CLOUD_AGENT_NEXT_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences';
+import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage';
+import { getUserFromAuth } from '@/lib/user/server';
+import { NextResponse } from 'next/server';
+
+export async function GET(): Promise<
+  NextResponse<{ error: string } | { balance: number; isDepleted: boolean }>
+> {
+  const { user, authFailedResponse, organizationId } = await getUserFromAuth({
+    adminOnly: false,
+    expectedAudience: CLOUD_AGENT_NEXT_AUDIENCE,
+  });
+
+  if (authFailedResponse) return authFailedResponse;
+
+  const { balance } = await getBalanceAndOrgSettings(organizationId, user);
+
+  return NextResponse.json({ balance, isDepleted: balance <= 0 });
+}

apps/web/src/lib/auth/cloud-agent-workflow-user.test.ts

@@ -0,0 +1,76 @@
+import jwt from 'jsonwebtoken';
+import { eq } from 'drizzle-orm';
+import { db } from '@/lib/drizzle';
+import { kilocode_users } from '@kilocode/db/schema';
+import { insertTestUser } from '@/tests/helpers/user.helper';
+import { generateCloudAgentWorkflowToken } from '@/lib/tokens';
+import { prepareCloudAgentWorkflowUser } from './cloud-agent-workflow-user';
+
+const issuance = { enabled: true };
+jest.mock('@/lib/config.server', () => ({
+  ...jest.requireActual('@/lib/config.server'),
+  isResourceTokenIssuanceEnabled: () => issuance.enabled,
+}));
+
+beforeEach(() => {
+  issuance.enabled = true;
+});
+
+test('initializes a null pepper and issues modern review admission with the persisted value', async () => {
+  const user = await insertTestUser({ api_token_pepper: null });
+  expect(() =>
+    generateCloudAgentWorkflowToken(user, { tokenSource: 'code-review', expiresIn: 3600 })
+  ).toThrow('current user pepper');
+  const prepared = await prepareCloudAgentWorkflowUser(user);
+  const [persisted] = await db.select().from(kilocode_users).where(eq(kilocode_users.id, user.id));
+  expect(prepared.api_token_pepper).toEqual(expect.any(String));
+  expect(prepared.api_token_pepper).toBe(persisted.api_token_pepper);
+  const claims = jwt.decode(
+    generateCloudAgentWorkflowToken(prepared, { tokenSource: 'code-review', expiresIn: 3600 })
+  );
+  expect(claims).toMatchObject({
+    aud: 'cloud-agent-next',
+    tokenPurpose: 'internal-service',
+    credentialExchange: false,
+    apiTokenPepper: persisted.api_token_pepper,
+    runtimeAdmission: {
+      source: 'automation',
+      authorizationUserId: user.id,
+      authorizationPepper: persisted.api_token_pepper,
+    },
+  });
+});
+
+test('concurrent initializers return the same persisted pepper', async () => {
+  const user = await insertTestUser({ api_token_pepper: null });
+  const results = await Promise.all(
+    Array.from({ length: 8 }, () => prepareCloudAgentWorkflowUser(user))
+  );
+  expect(new Set(results.map(result => result.api_token_pepper)).size).toBe(1);
+  expect(results[0].api_token_pepper).toEqual(expect.any(String));
+});
+
+test('preserves a pepper assigned after the user snapshot was loaded', async () => {
+  const user = await insertTestUser({ api_token_pepper: null });
+  await db
+    .update(kilocode_users)
+    .set({ api_token_pepper: 'concurrent-rotation' })
+    .where(eq(kilocode_users.id, user.id));
+  expect((await prepareCloudAgentWorkflowUser(user)).api_token_pepper).toBe('concurrent-rotation');
+});
+
+test('preserves existing peppers and leaves legacy issuance unchanged', async () => {
+  const existing = await insertTestUser({ api_token_pepper: 'existing-pepper' });
+  expect(await prepareCloudAgentWorkflowUser(existing)).toBe(existing);
+  issuance.enabled = false;
+  const user = await insertTestUser({ api_token_pepper: null });
+  expect(await prepareCloudAgentWorkflowUser(user)).toBe(user);
+  const [persisted] = await db.select().from(kilocode_users).where(eq(kilocode_users.id, user.id));
+  expect(persisted.api_token_pepper).toBeNull();
+});
+
+test('fails closed if the user was deleted', async () => {
+  const user = await insertTestUser({ api_token_pepper: null });
+  await db.delete(kilocode_users).where(eq(kilocode_users.id, user.id));
+  await expect(prepareCloudAgentWorkflowUser(user)).rejects.toThrow('not found');
+});

apps/web/src/lib/auth/cloud-agent-workflow-user.ts

@@ -0,0 +1,24 @@
+import { randomUUID } from 'node:crypto';
+import { eq, sql } from 'drizzle-orm';
+import { kilocode_users, type User } from '@kilocode/db/schema';
+import { db } from '@/lib/drizzle';
+import { isResourceTokenIssuanceEnabled, type ResourceTokenFamily } from '@/lib/config.server';
+
+export async function prepareCloudAgentWorkflowUser(
+  user: User,
+  requiredFamilies: readonly ResourceTokenFamily[] = ['cloud-agent-next']
+): Promise<User> {
+  if (!requiredFamilies.some(isResourceTokenIssuanceEnabled) || user.api_token_pepper !== null) {
+    return user;
+  }
+
+  // Use the primary's persisted value: another issuer or revocation may have
+  // assigned a pepper since this user was loaded. Never overwrite that value.
+  const [currentUser] = await db
+    .update(kilocode_users)
+    .set({ api_token_pepper: sql`COALESCE(${kilocode_users.api_token_pepper}, ${randomUUID()})` })
+    .where(eq(kilocode_users.id, user.id))
+    .returning();
+  if (!currentUser) throw new Error(`User ${user.id} not found`);
+  return currentUser;
+}

apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts

@@ -1,17 +1,15 @@
 const mockGetFixTicketById = jest.fn();
 
-jest.mock('@/lib/drizzle', () => ({
-  db: {
-    select: () => ({
-      from: () => ({
-        where: () => ({
-          limit: () => [{ id: 'user-1', api_token_pepper: 'pepper' }],
-        }),
-      }),
-    }),
-  },
+jest.mock('@/lib/config.server', () => ({
+  ...jest.requireActual('@/lib/config.server'),
+  isResourceTokenIssuanceEnabled: () => true,
 }));
 
+import { insertTestUser } from '@/tests/helpers/user.helper';
+import { db } from '@/lib/drizzle';
+import { kilocode_users } from '@kilocode/db/schema';
+import { eq } from 'drizzle-orm';
+
 jest.mock('@/lib/tokens', () => ({
   generateCloudAgentWorkflowToken: jest.fn(() => 'workflow-token'),
   TOKEN_EXPIRY: { default: 3600 },
@@ -30,7 +28,12 @@ const mockGenerateCloudAgentWorkflowToken = jest.mocked(generateCloudAgentWorkfl
 
 const organizationId = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa';
 
-beforeEach(() => {
+afterEach(async () => {
+  await db.delete(kilocode_users).where(eq(kilocode_users.id, 'user-1'));
+});
+
+beforeEach(async () => {
+  await insertTestUser({ id: 'user-1', api_token_pepper: null });
   jest.clearAllMocks();
   mockGetFixTicketById.mockResolvedValue({
     repo_full_name: 'kilo/repo',
@@ -56,8 +59,13 @@ describe('prepareFixPayload workflow token ownership', () => {
         },
       });
 
+      const [persisted] = await db
+        .select()
+        .from(kilocode_users)
+        .where(eq(kilocode_users.id, 'user-1'));
+      expect(persisted.api_token_pepper).toEqual(expect.any(String));
       expect(mockGenerateCloudAgentWorkflowToken).toHaveBeenCalledWith(
-        expect.objectContaining({ id: 'user-1' }),
+        expect.objectContaining({ id: 'user-1', api_token_pepper: persisted.api_token_pepper }),
         expect.objectContaining({ organizationId: expectedOrganizationId })
       );
     }