From 906c5c3f18fae81903ed055855bf94973b3c7f66 Mon Sep 17 00:00:00 2001
From: Elie Habib <elie.habib@gmail.com>
Date: Fri, 27 Mar 2026 11:49:01 +0400
Subject: [PATCH] test(auth): cover widget-agent fallback cleanup (#2354)

---
 api/widget-agent.ts              | 10 +++++++---
 src/components/McpDataPanel.ts   | 10 +++++-----
 tests/widget-agent-auth.test.mts | 28 ++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/api/widget-agent.ts b/api/widget-agent.ts
index 094a66cdb..13291f32b 100644
--- a/api/widget-agent.ts
+++ b/api/widget-agent.ts
@@ -24,11 +24,15 @@ import { validateBearerToken } from '../server/auth-session';
 const RELAY_BASE = 'https://proxy.worldmonitor.app';
 const WIDGET_AGENT_KEY = process.env.WIDGET_AGENT_KEY ?? '';
 const PRO_WIDGET_KEY = process.env.PRO_WIDGET_KEY ?? '';
+const WORLDMONITOR_VALID_KEY_SET = new Set(
+  (process.env.WORLDMONITOR_VALID_KEYS ?? '')
+    .split(',')
+    .map((v) => v.trim())
+    .filter(Boolean),
+);
 
 function hasValidWorldMonitorKey(key: string): boolean {
-  if (!key) return false;
-  const validKeys = (process.env.WORLDMONITOR_VALID_KEYS ?? '').split(',').map((v) => v.trim()).filter(Boolean);
-  return validKeys.includes(key);
+  return Boolean(key) && WORLDMONITOR_VALID_KEY_SET.has(key);
 }
 
 function json(body: unknown, status: number, cors: Record<string, string>): Response {
diff --git a/src/components/McpDataPanel.ts b/src/components/McpDataPanel.ts
index d04a9aadd..a8edd61cc 100644
--- a/src/components/McpDataPanel.ts
+++ b/src/components/McpDataPanel.ts
@@ -176,11 +176,11 @@ export class McpDataPanel extends Panel {
 
     try {
       const testerKey = getBrowserTesterKey();
-      const headers: Record<string, string> = {
-        'Content-Type': 'application/json',
-        'X-Widget-Key': getWidgetAgentKey(),
-        'X-Pro-Key': getProWidgetKey(),
-      };
+      const widgetKey = getWidgetAgentKey();
+      const proKey = getProWidgetKey();
+      const headers: Record<string, string> = { 'Content-Type': 'application/json' };
+      if (widgetKey) headers['X-Widget-Key'] = widgetKey;
+      if (proKey) headers['X-Pro-Key'] = proKey;
       if (testerKey) headers['X-WorldMonitor-Key'] = testerKey;
       const res = await fetch(widgetAgentUrl(), {
         method: 'POST',
diff --git a/tests/widget-agent-auth.test.mts b/tests/widget-agent-auth.test.mts
index 8e6a59150..ec3498433 100644
--- a/tests/widget-agent-auth.test.mts
+++ b/tests/widget-agent-auth.test.mts
@@ -77,6 +77,34 @@ describe('widget-agent unified tester key auth', () => {
     });
   });
 
+  it('falls back to legacy tester keys when X-WorldMonitor-Key is invalid', async () => {
+    const res = await handler(new Request('https://www.worldmonitor.app/api/widget-agent', {
+      method: 'POST',
+      headers: {
+        Origin: 'https://www.worldmonitor.app',
+        'Content-Type': 'application/json',
+        'X-WorldMonitor-Key': 'wrong-key',
+        'X-Pro-Key': 'server-pro-key',
+      },
+      body: JSON.stringify({ prompt: 'Build a widget', mode: 'create', tier: 'basic' }),
+    }));
+
+    assert.equal(res.status, 200);
+    assert.equal(fetchMock.mock.calls.length, 1);
+
+    const call = fetchMock.mock.calls[0];
+    const init = call.arguments[1] as RequestInit;
+    const headers = new Headers(init.headers);
+    assert.equal(headers.get('X-Widget-Key'), 'server-widget-key');
+    assert.equal(headers.get('X-Pro-Key'), 'server-pro-key');
+
+    assert.deepEqual(JSON.parse(String(init.body)), {
+      prompt: 'Build a widget',
+      mode: 'create',
+      tier: 'pro',
+    });
+  });
+
   it('rejects invalid X-WorldMonitor-Key before relay fetch', async () => {
     const res = await handler(new Request('https://www.worldmonitor.app/api/widget-agent', {
       method: 'POST',