fix(auth): handle device approval state errors

apps/web/src/app/api/device-auth/tokens/route.test.ts

@@ -46,6 +46,53 @@ describe('POST /api/device-auth/tokens', () => {
     expect(mockApprove).toHaveBeenCalledWith('ABCD-EFGH', 'user-1');
   });
 
+  it.each([
+    {
+      message: 'Device authorization request not found',
+      status: 404,
+      error: 'Not found',
+    },
+    {
+      message: 'Device authorization request is not pending',
+      status: 409,
+      error: 'Device authorization request can no longer be approved',
+    },
+    {
+      message: 'Device authorization request has expired',
+      status: 410,
+      error: 'Device authorization request has expired',
+    },
+  ])('returns $status when approval fails with "$message"', async ({ message, status, error }) => {
+    mockGetUserFromSession.mockResolvedValue({
+      user: defineTestUser({ id: 'user-1' }),
+      authFailedResponse: null,
+    });
+    mockApprove.mockRejectedValueOnce(new Error(message));
+
+    const response = await POST(
+      createRequest(JSON.stringify({ code: 'ABCD-EFGH' }), sessionHeaders)
+    );
+
+    expect(response.status).toBe(status);
+    expect(await response.json()).toEqual({ error });
+    expect(mockApprove).toHaveBeenCalledWith('ABCD-EFGH', 'user-1');
+  });
+
+  it.each([new Error('Database unavailable'), 'Unexpected rejection'])(
+    'rethrows unexpected approval failures: %s',
+    async error => {
+      mockGetUserFromSession.mockResolvedValue({
+        user: defineTestUser({ id: 'user-1' }),
+        authFailedResponse: null,
+      });
+      mockApprove.mockRejectedValueOnce(error);
+
+      await expect(
+        POST(createRequest(JSON.stringify({ code: 'ABCD-EFGH' }), sessionHeaders))
+      ).rejects.toBe(error);
+    }
+  );
+
   it.each([undefined, 'https://evil.example'])(
     'rejects a %s origin before authenticating',
     async origin => {

apps/web/src/app/api/device-auth/tokens/route.ts

@@ -31,7 +31,25 @@ export async function POST(request: Request) {
 
   const { code } = validation.data;
 
-  await approveDeviceAuthRequest(code, user.id);
+  try {
+    await approveDeviceAuthRequest(code, user.id);
+  } catch (error) {
+    if (error instanceof Error) {
+      if (error.message === 'Device authorization request not found') {
+        return NextResponse.json({ error: 'Not found' }, { status: 404 });
+      }
+      if (error.message === 'Device authorization request is not pending') {
+        return NextResponse.json(
+          { error: 'Device authorization request can no longer be approved' },
+          { status: 409 }
+        );
+      }
+      if (error.message === 'Device authorization request has expired') {
+        return NextResponse.json({ error: error.message }, { status: 410 });
+      }
+    }
+    throw error;
+  }
 
   return NextResponse.json({ success: true });
 }

apps/web/src/lib/device-auth/device-auth.test.ts

@@ -234,6 +234,127 @@ describe('Device Auth', () => {
       );
     });
 
+    test('concurrent approvals preserve the single winning user', async () => {
+      const otherUserId = `${testUserId}-other`;
+      await db.insert(kilocode_users).values({
+        id: otherUserId,
+        google_user_email: `other-${testUserEmail}`,
+        google_user_name: 'Other User',
+        google_user_image_url: 'https://example.com/avatar.jpg',
+        stripe_customer_id: 'cus_other',
+      });
+
+      try {
+        for (let i = 0; i < 20; i++) {
+          const { code, deviceCode } = await createDeviceAuthRequest({});
+          const results = await Promise.allSettled(
+            [testUserId, otherUserId].map(async userId => {
+              await approveDeviceAuthRequest(code, userId);
+              return userId;
+            })
+          );
+          const approved = results.filter(result => result.status === 'fulfilled');
+          const rejected = results.filter(result => result.status === 'rejected');
+
+          expect(approved).toHaveLength(1);
+          expect(rejected).toHaveLength(1);
+          expect(rejected[0]?.reason).toEqual(
+            new Error('Device authorization request is not pending')
+          );
+          expect((await getDeviceAuthRequest(code))?.kilo_user_id).toBe(approved[0]?.value);
+          expect((await consumeDeviceAuthByDeviceCode(deviceCode)).userId).toBe(approved[0]?.value);
+        }
+      } finally {
+        await db
+          .delete(device_auth_requests)
+          .where(eq(device_auth_requests.kilo_user_id, otherUserId));
+        await db.delete(kilocode_users).where(eq(kilocode_users.id, otherUserId));
+      }
+    });
+
+    test.each(['denied', 'consumed', 'expired'] as const)(
+      'does not change a %s request',
+      async status => {
+        const { code } = await createDeviceAuthRequest({});
+        await db
+          .update(device_auth_requests)
+          .set({ status, kilo_user_id: testUserId })
+          .where(eq(device_auth_requests.code, code));
+        const before = await getDeviceAuthRequest(code);
+
+        await expect(approveDeviceAuthRequest(code, testUserId)).rejects.toThrow(
+          'Device authorization request is not pending'
+        );
+
+        expect(await getDeviceAuthRequest(code)).toEqual(before);
+      }
+    );
+
+    test('reports a conflict when denial wins before the expiry fallback', async () => {
+      const { code } = await createDeviceAuthRequest({});
+      await db
+        .update(device_auth_requests)
+        .set({
+          expires_at: new Date(Date.now() - 1000).toISOString(),
+          kilo_user_id: testUserId,
+        })
+        .where(eq(device_auth_requests.code, code));
+      const pendingRequest = await getDeviceAuthRequest(code);
+      const staleSelect = Object.assign(db.select({}), {
+        from: jest.fn().mockReturnValue({
+          where: jest.fn().mockReturnValue({
+            limit: jest.fn().mockImplementation(async () => {
+              await denyDeviceAuthRequest(code);
+              return [pendingRequest];
+            }),
+          }),
+        }),
+      });
+      const selectSpy = jest.spyOn(db, 'select').mockReturnValueOnce(staleSelect);
+
+      try {
+        await expect(approveDeviceAuthRequest(code, testUserId)).rejects.toThrow(
+          'Device authorization request is not pending'
+        );
+        expect((await getDeviceAuthRequest(code))?.status).toBe('denied');
+      } finally {
+        selectSpy.mockRestore();
+      }
+    });
+
+    test('reports expiration when the atomic update reaches the expiry deadline', async () => {
+      const { code } = await createDeviceAuthRequest({});
+      const deadline = new Date();
+      await db
+        .update(device_auth_requests)
+        .set({ expires_at: deadline.toISOString(), kilo_user_id: testUserId })
+        .where(eq(device_auth_requests.code, code));
+      jest.useFakeTimers({
+        now: deadline,
+        doNotFake: [
+          'hrtime',
+          'nextTick',
+          'performance',
+          'queueMicrotask',
+          'setImmediate',
+          'clearImmediate',
+          'setInterval',
+          'clearInterval',
+          'setTimeout',
+          'clearTimeout',
+        ],
+      });
+
+      try {
+        await expect(approveDeviceAuthRequest(code, testUserId)).rejects.toThrow(
+          'Device authorization request has expired'
+        );
+        expect((await getDeviceAuthRequest(code))?.status).toBe('expired');
+      } finally {
+        jest.useRealTimers();
+      }
+    });
+
     test('throws error for expired request', async () => {
       const { code } = await createDeviceAuthRequest({});
 
@@ -250,6 +371,29 @@ describe('Device Auth', () => {
   });
 
   describe('denyDeviceAuthRequest', () => {
+    test('concurrent approval and denial preserve the single winning transition', async () => {
+      for (let i = 0; i < 20; i++) {
+        const { code } = await createDeviceAuthRequest({});
+        try {
+          const results = await Promise.allSettled([
+            approveDeviceAuthRequest(code, testUserId).then(() => 'approved' as const),
+            denyDeviceAuthRequest(code).then(() => 'denied' as const),
+          ]);
+          const succeeded = results.filter(result => result.status === 'fulfilled');
+          const rejected = results.filter(result => result.status === 'rejected');
+
+          expect(succeeded).toHaveLength(1);
+          expect(rejected).toHaveLength(1);
+          expect(rejected[0]?.reason).toEqual(
+            new Error('Device authorization request is not pending')
+          );
+          expect((await getDeviceAuthRequest(code))?.status).toBe(succeeded[0]?.value);
+        } finally {
+          await db.delete(device_auth_requests).where(eq(device_auth_requests.code, code));
+        }
+      }
+    });
+
     test('denies a pending request', async () => {
       const { code } = await createDeviceAuthRequest({});
 

apps/web/src/lib/device-auth/device-auth.ts

@@ -1,7 +1,7 @@
 import 'server-only';
 import { db } from '@/lib/drizzle';
 import { device_auth_requests, device_sessions, kilocode_users } from '@kilocode/db/schema';
-import { eq, and, lt, gt, isNull, isNotNull, sql } from 'drizzle-orm';
+import { eq, and, lt, lte, gt, isNull, isNotNull, sql } from 'drizzle-orm';
 import { generateApiToken } from '@/lib/tokens';
 import { randomInt, createHash, randomBytes } from 'node:crypto';
 import { createDeviceSession, issueSessionCredentials } from '@/lib/auth/device-sessions';
@@ -132,6 +132,27 @@ export function isDeviceAuthRequestExpired(request: {
  * Approve a device authorization request.
  */
 export async function approveDeviceAuthRequest(code: string, userId: string): Promise<void> {
+  const now = new Date().toISOString();
+  const [approved] = await db
+    .update(device_auth_requests)
+    .set({
+      status: 'approved',
+      kilo_user_id: userId,
+      approved_at: now,
+    })
+    .where(
+      and(
+        eq(device_auth_requests.code, code),
+        eq(device_auth_requests.status, 'pending'),
+        gt(device_auth_requests.expires_at, now)
+      )
+    )
+    .returning({ id: device_auth_requests.id });
+
+  if (approved) {
+    return;
+  }
+
   const request = await getDeviceAuthRequest(code);
 
   if (!request) {
@@ -142,22 +163,22 @@ export async function approveDeviceAuthRequest(code: string, userId: string): Pr
     throw new Error('Device authorization request is not pending');
   }
 
-  if (isDeviceAuthRequestExpired(request)) {
-    await db
-      .update(device_auth_requests)
-      .set({ status: 'expired' })
-      .where(eq(device_auth_requests.code, code));
-    throw new Error('Device authorization request has expired');
-  }
-
-  await db
+  const [expired] = await db
     .update(device_auth_requests)
-    .set({
-      status: 'approved',
-      kilo_user_id: userId,
-      approved_at: new Date().toISOString(),
-    })
-    .where(eq(device_auth_requests.code, code));
+    .set({ status: 'expired' })
+    .where(
+      and(
+        eq(device_auth_requests.code, code),
+        eq(device_auth_requests.status, 'pending'),
+        lte(device_auth_requests.expires_at, now)
+      )
+    )
+    .returning({ id: device_auth_requests.id });
+
+  if (!expired) {
+    throw new Error('Device authorization request is not pending');
+  }
+  throw new Error('Device authorization request has expired');
 }
 
 /**
@@ -174,10 +195,15 @@ export async function denyDeviceAuthRequest(code: string): Promise<void> {
     throw new Error('Device authorization request is not pending');
   }
 
-  await db
+  const [denied] = await db
     .update(device_auth_requests)
     .set({ status: 'denied' })
-    .where(eq(device_auth_requests.code, code));
+    .where(and(eq(device_auth_requests.code, code), eq(device_auth_requests.status, 'pending')))
+    .returning({ id: device_auth_requests.id });
+
+  if (!denied) {
+    throw new Error('Device authorization request is not pending');
+  }
 }
 
 /**