fix: fallback to local extract on ai allowlist miss

2026-09-17 13:24:13 +08:00 · 2026-09-17 13:06:52 +08:00
parent 1e054578d8
commit 67c6c1f82c
9 changed files with 233 additions and 66 deletions
@@ -1,5 +1,7 @@
 import { test, expect } from '@playwright/test';
-import { WORKER_URL, createTestAddress, deleteAddress } from '../../fixtures/test-helpers';
+import { WORKER_URL, WORKER_URL_ENV_OFF, createTestAddress, deleteAddress } from '../../fixtures/test-helpers';
+
+const ADMIN_HEADERS = { 'x-admin-auth': 'e2e-admin-pass' };

 test.describe('Telegram AI extraction rendering', () => {
  test('realtime mail stores AI extraction metadata for Telegram rendering', async ({ request }) => {
@@ -102,4 +104,159 @@ test.describe('Telegram AI extraction rendering', () => {
      await deleteAddress(request, jwt);
    }
  });
+
+  test('ai allowlist only gates Workers AI and falls back to local rules on miss', async ({ request }) => {
+    const { jwt, address } = await createTestAddress(request, 'tg-ai-local-fallback');
+
+    try {
+      const offSettingsRes = await request.post(`${WORKER_URL}/admin/ai_extract/settings`, {
+        headers: ADMIN_HEADERS,
+        data: {
+          enableAllowList: false,
+          allowList: [],
+        },
+      });
+      expect(offSettingsRes.ok()).toBe(true);
+
+      const aiRaw = [
+        'From: sender@test.example.com',
+        `To: ${address}`,
+        'Subject: AI allowlist off',
+        `Message-ID: <ai-allowlist-off-${Date.now()}@test>`,
+        'MIME-Version: 1.0',
+        'Content-Type: text/plain; charset=utf-8',
+        '',
+        'Your verification code is: 593817',
+      ].join('\r\n');
+
+      const aiReceiveRes = await request.post(`${WORKER_URL}/__test/receive_mail`, {
+        data: {
+          from: 'sender@test.example.com',
+          to: address,
+          raw: aiRaw,
+          extract_mode: 'ai',
+          ai_extract_result: {
+            type: 'auth_link',
+            result: 'https://example.com/ai-used',
+            result_text: '',
+          },
+        },
+      });
+      expect(aiReceiveRes.ok()).toBe(true);
+      expect((await aiReceiveRes.json()).success).toBe(true);
+
+      const aiMailsRes = await request.get(`${WORKER_URL}/api/mails?limit=10&offset=0`, {
+        headers: { Authorization: `Bearer ${jwt}` },
+      });
+      expect(aiMailsRes.ok()).toBe(true);
+      const aiMailsBody = await aiMailsRes.json();
+      expect(aiMailsBody.results[0].metadata).toBeTruthy();
+      expect(JSON.parse(aiMailsBody.results[0].metadata).ai_extract).toEqual({
+        type: 'auth_link',
+        result: 'https://example.com/ai-used',
+        result_text: '',
+      });
+
+      const onSettingsRes = await request.post(`${WORKER_URL}/admin/ai_extract/settings`, {
+        headers: ADMIN_HEADERS,
+        data: {
+          enableAllowList: true,
+          allowList: ['allowed@example.com'],
+        },
+      });
+      expect(onSettingsRes.ok()).toBe(true);
+
+      const raw = [
+        'From: sender@test.example.com',
+        `To: ${address}`,
+        'Subject: AI allowlist fallback',
+        `Message-ID: <ai-allowlist-fallback-${Date.now()}@test>`,
+        'MIME-Version: 1.0',
+        'Content-Type: text/plain; charset=utf-8',
+        '',
+        'Your verification code is: 593817',
+      ].join('\r\n');
+
+      const receiveRes = await request.post(`${WORKER_URL}/__test/receive_mail`, {
+        data: {
+          from: 'sender@test.example.com',
+          to: address,
+          raw,
+          extract_mode: 'ai',
+          // If Workers AI were called, this would win. An allowlist miss must
+          // skip AI and use local extraction instead.
+          ai_extract_result: {
+            type: 'auth_link',
+            result: 'https://example.com/should-not-be-used',
+            result_text: '',
+          },
+        },
+      });
+      expect(receiveRes.ok()).toBe(true);
+      expect((await receiveRes.json()).success).toBe(true);
+
+      const mailsRes = await request.get(`${WORKER_URL}/api/mails?limit=10&offset=0`, {
+        headers: { Authorization: `Bearer ${jwt}` },
+      });
+      expect(mailsRes.ok()).toBe(true);
+      const { results } = await mailsRes.json();
+      expect(results.length).toBeGreaterThanOrEqual(2);
+
+      const metadata = JSON.parse(results[0].metadata);
+      expect(metadata.ai_extract).toEqual({
+        type: 'auth_code',
+        result: '593817',
+        result_text: '',
+      });
+    } finally {
+      await request.post(`${WORKER_URL}/admin/ai_extract/settings`, {
+        headers: ADMIN_HEADERS,
+        data: {
+          enableAllowList: false,
+          allowList: [],
+        },
+      });
+      await deleteAddress(request, jwt);
+    }
+  });
+
+  test('env-off worker keeps extraction disabled without test overrides', async ({ request }) => {
+    test.skip(!WORKER_URL_ENV_OFF, 'WORKER_URL_ENV_OFF is not configured');
+
+    const { jwt, address, address_id } = await createTestAddress(request, 'tg-extract-off', 'test.example.com', WORKER_URL_ENV_OFF);
+
+    try {
+      const raw = [
+        'From: sender@test.example.com',
+        `To: ${address}`,
+        'Subject: Extraction disabled',
+        `Message-ID: <extract-off-${Date.now()}@test>`,
+        'MIME-Version: 1.0',
+        'Content-Type: text/plain; charset=utf-8',
+        '',
+        'Your verification code is: 374829',
+      ].join('\r\n');
+
+      const receiveRes = await request.post(`${WORKER_URL_ENV_OFF}/__test/receive_mail`, {
+        data: {
+          from: 'sender@test.example.com',
+          to: address,
+          raw,
+        },
+      });
+      expect(receiveRes.ok()).toBe(true);
+      expect((await receiveRes.json()).success).toBe(true);
+
+      const mailsRes = await request.get(`${WORKER_URL_ENV_OFF}/api/mails?limit=10&offset=0`, {
+        headers: { Authorization: `Bearer ${jwt}` },
+      });
+      expect(mailsRes.ok()).toBe(true);
+      const { results } = await mailsRes.json();
+      expect(results).toHaveLength(1);
+      expect(results[0].metadata).toBeFalsy();
+    } finally {
+      const deleteRes = await request.delete(`${WORKER_URL_ENV_OFF}/admin/delete_address/${address_id}`);
+      expect(deleteRes.ok()).toBe(true);
+    }
+  });
 });