Skip to content

fix(ui): allow web search tests to use env keys #280

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
snvtac wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(ui): allow web search tests to use env keys #280

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 68 additions & 12 deletions ui/server/routes/config.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { prepareBackgroundSpawnOptions } from '../utils/processSpawn.js';
import { parse as parseYaml } from 'yaml';
import {
buildDefaultPilotDeckConfig,
buildRuntimeEnv,
configToYaml,
getPilotDeckConfigPath,
maskSecrets,
Expand All @@ -29,6 +30,60 @@ async function notifyGatewayConfigReload() {
}

const router = express.Router();
const MASK = '********';
const WEB_SEARCH_ENV_KEYS = {
glm: ['GLM_WEB_SEARCH_API_KEY', 'ZAI_API_KEY'],
tavily: ['TAVILY_API_KEY'],
custom: ['CUSTOM_WEB_SEARCH_API_KEY'],
};

function isRecord(value) {
return value && typeof value === 'object' && !Array.isArray(value);
}

function normalizeCredential(value) {
if (typeof value !== 'string') return '';
const trimmed = value.trim();
if (!trimmed || trimmed === MASK || trimmed === 'PLACEHOLDER_RUN_ONBOARDING_TO_REPLACE' || trimmed.startsWith('PLACEHOLDER_')) {
return '';
}
return trimmed;
}

function buildWebSearchCredentialEnv(requestCustomEnv) {
let diskConfig = buildDefaultPilotDeckConfig();
try {
diskConfig = readPilotDeckConfigFile().config;
} catch {
// Keep the test route usable even when the main config is temporarily invalid.
}

const diskCustomEnv = isRecord(diskConfig.customEnv) ? diskConfig.customEnv : {};
const restoredRequestCustomEnv = isRecord(requestCustomEnv)
? preserveMaskedSecrets(requestCustomEnv, diskCustomEnv)
: {};
const customEnv = isRecord(requestCustomEnv) && isRecord(restoredRequestCustomEnv)
? restoredRequestCustomEnv
: diskCustomEnv;

return {
...process.env,
...buildRuntimeEnv({
...diskConfig,
customEnv,
}),
};
}

function resolveWebSearchTestApiKey(apiKey, provider, env) {
const inlineKey = normalizeCredential(apiKey);
if (inlineKey) return inlineKey;
for (const key of WEB_SEARCH_ENV_KEYS[provider] ?? WEB_SEARCH_ENV_KEYS.glm) {
const envKey = normalizeCredential(env[key]);
if (envKey) return envKey;
}
return '';
}

function serializeConfigResponse(record, reloadResult = null) {
const validation = validatePilotDeckConfig(record.config);
Expand Down Expand Up @@ -312,16 +367,17 @@ router.post('/test-connection', async (req, res) => {
* established by `/test-connection`.
*/
router.post('/test-web-search', async (req, res) => {
const { provider, apiKey, endpoint, customProvider } = req.body || {};
const { provider, apiKey, endpoint, customProvider, customEnv } = req.body || {};
const selectedProvider = provider === 'tavily' || provider === 'custom' ? provider : 'glm';
const custom = customProvider && typeof customProvider === 'object' ? customProvider : {};
const customAuth = typeof custom.auth === 'string' ? custom.auth : 'bearer';
const customMethod = custom.method === 'GET' ? 'GET' : 'POST';
const queryParam = typeof custom.queryParam === 'string' && custom.queryParam.trim() ? custom.queryParam.trim() : 'query';
const apiKeyParam = typeof custom.apiKeyParam === 'string' && custom.apiKeyParam.trim() ? custom.apiKeyParam.trim() : 'api_key';
const resultsPath = typeof custom.resultsPath === 'string' ? custom.resultsPath.trim() : '';
const trimmedKey = typeof apiKey === 'string' ? apiKey.trim() : '';
if (!trimmedKey && !(selectedProvider === 'custom' && customAuth === 'none')) {
const credentialEnv = buildWebSearchCredentialEnv(customEnv);
const effectiveKey = resolveWebSearchTestApiKey(apiKey, selectedProvider, credentialEnv);
if (!effectiveKey && !(selectedProvider === 'custom' && customAuth === 'none')) {
return res.status(400).json({ ok: false, error: 'API key is required.' });
}
const trimmedEndpoint = typeof endpoint === 'string' ? endpoint.trim() : '';
Expand All @@ -347,7 +403,7 @@ router.post('/test-web-search', async (req, res) => {
Accept: 'application/json',
},
body: JSON.stringify({
api_key: trimmedKey,
api_key: effectiveKey,
query: 'hello',
max_results: 3,
include_answer: true,
Expand All @@ -363,13 +419,13 @@ router.post('/test-web-search', async (req, res) => {
headers['Content-Type'] = 'application/json';
body[queryParam] = 'hello';
}
if (customAuth === 'bearer' && trimmedKey) {
headers.Authorization = `Bearer ${trimmedKey}`;
} else if (customAuth === 'queryApiKey' && trimmedKey) {
url.searchParams.set(apiKeyParam, trimmedKey);
} else if (customAuth === 'bodyApiKey' && trimmedKey) {
if (customMethod === 'GET') url.searchParams.set(apiKeyParam, trimmedKey);
else body[apiKeyParam] = trimmedKey;
if (customAuth === 'bearer' && effectiveKey) {
headers.Authorization = `Bearer ${effectiveKey}`;
} else if (customAuth === 'queryApiKey' && effectiveKey) {
url.searchParams.set(apiKeyParam, effectiveKey);
} else if (customAuth === 'bodyApiKey' && effectiveKey) {
if (customMethod === 'GET') url.searchParams.set(apiKeyParam, effectiveKey);
else body[apiKeyParam] = effectiveKey;
}
requestUrl = url.toString();
requestInit = {
Expand All @@ -382,7 +438,7 @@ router.post('/test-web-search', async (req, res) => {
requestInit = {
method: 'POST',
headers: {
Authorization: `Bearer ${trimmedKey}`,
Authorization: `Bearer ${effectiveKey}`,
'Content-Type': 'application/json',
Accept: 'application/json',
},
Expand Down
250 changes: 250 additions & 0 deletions ui/server/routes/config.test.js
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,250 @@
import http from 'node:http';
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import express from 'express';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

const TEST_ENV_KEYS = [
'PILOT_HOME',
'PILOTDECK_CONFIG_PATH',
'TAVILY_API_KEY',
'CUSTOM_WEB_SEARCH_API_KEY',
'GLM_WEB_SEARCH_API_KEY',
'ZAI_API_KEY',
];

let previousEnv;
let tempRoot;

async function createConfigApp() {
vi.resetModules();
vi.doMock('../pilotdeck-bridge.js', () => ({
getPilotDeckGateway: vi.fn(async () => null),
}));
vi.doMock('../services/pilotdeckConfigReloader.js', () => ({
reloadPilotDeckConfig: vi.fn(async () => null),
}));
vi.doMock('../services/pilotdeckConfigWatcher.js', () => ({
suppressNextWatchEvent: vi.fn(),
}));
const { default: configRouter } = await import('./config.js');
const app = express();
app.use(express.json());
app.use('/api/config', configRouter);
return app;
}

async function postJson(app, body) {
const server = http.createServer(app);
await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve));
const address = server.address();
const port = typeof address === 'object' && address ? address.port : 0;

try {
return await new Promise((resolve, reject) => {
const req = http.request(
{
hostname: '127.0.0.1',
port,
path: '/api/config/test-web-search',
method: 'POST',
headers: { 'Content-Type': 'application/json' },
},
(res) => {
let raw = '';
res.setEncoding('utf8');
res.on('data', (chunk) => { raw += chunk; });
res.on('end', () => {
try {
resolve({ status: res.statusCode, body: JSON.parse(raw) });
} catch (error) {
reject(error);
}
});
},
);
req.on('error', reject);
req.end(JSON.stringify(body));
});
} finally {
await new Promise((resolve) => server.close(resolve));
}
}

function mockProviderFetch(responseBody = { results: [{ title: 'ok' }] }) {
const fetchMock = vi.fn(async () => new Response(JSON.stringify(responseBody), {
status: 200,
headers: { 'Content-Type': 'application/json' },
}));
vi.stubGlobal('fetch', fetchMock);
return fetchMock;
}

describe('config web search test route', () => {
beforeEach(() => {
previousEnv = Object.fromEntries(TEST_ENV_KEYS.map((key) => [key, process.env[key]]));
for (const key of TEST_ENV_KEYS) delete process.env[key];
tempRoot = mkdtempSync(join(tmpdir(), 'pilotdeck-web-search-test-'));
process.env.PILOT_HOME = join(tempRoot, 'pilot-home');
});

afterEach(() => {
for (const key of TEST_ENV_KEYS) {
if (previousEnv[key] === undefined) delete process.env[key];
else process.env[key] = previousEnv[key];
}
vi.unstubAllGlobals();
vi.restoreAllMocks();
if (tempRoot) rmSync(tempRoot, { recursive: true, force: true });
});

it('uses TAVILY_API_KEY when no inline Tavily key is provided', async () => {
process.env.TAVILY_API_KEY = 'env-tavily';
const fetchMock = mockProviderFetch();
const app = await createConfigApp();

const result = await postJson(app, { provider: 'tavily', apiKey: '' });

expect(result.status).toBe(200);
expect(result.body.ok).toBe(true);
const requestBody = JSON.parse(fetchMock.mock.calls[0][1].body);
expect(requestBody.api_key).toBe('env-tavily');
});

it('uses GLM_WEB_SEARCH_API_KEY before ZAI_API_KEY for GLM tests', async () => {
process.env.GLM_WEB_SEARCH_API_KEY = 'env-glm';
process.env.ZAI_API_KEY = 'env-zai';
const fetchMock = mockProviderFetch({ search_result: [{ title: 'ok' }] });
const app = await createConfigApp();

const result = await postJson(app, { provider: 'glm', apiKey: '' });

expect(result.status).toBe(200);
expect(result.body.ok).toBe(true);
expect(fetchMock.mock.calls[0][1].headers.Authorization).toBe('Bearer env-glm');
});

it('falls back to ZAI_API_KEY when GLM_WEB_SEARCH_API_KEY is missing', async () => {
process.env.ZAI_API_KEY = 'env-zai';
const fetchMock = mockProviderFetch({ search_result: [{ title: 'ok' }] });
const app = await createConfigApp();

const result = await postJson(app, { provider: 'glm', apiKey: '' });

expect(result.status).toBe(200);
expect(result.body.ok).toBe(true);
expect(fetchMock.mock.calls[0][1].headers.Authorization).toBe('Bearer env-zai');
});

it('uses request customEnv for custom bearer tests', async () => {
const fetchMock = mockProviderFetch();
const app = await createConfigApp();

const result = await postJson(app, {
provider: 'custom',
apiKey: '',
endpoint: 'https://example.com/search',
customProvider: { auth: 'bearer' },
customEnv: { CUSTOM_WEB_SEARCH_API_KEY: 'custom-env' },
});

expect(result.status).toBe(200);
expect(result.body.ok).toBe(true);
expect(fetchMock.mock.calls[0][1].headers.Authorization).toBe('Bearer custom-env');
});

it('preserves saved customEnv secrets when the request contains a masked value', async () => {
mkdirSync(process.env.PILOT_HOME, { recursive: true });
writeFileSync(
join(process.env.PILOT_HOME, 'pilotdeck.yaml'),
[
'schemaVersion: 1',
'customEnv:',
' TAVILY_API_KEY: saved-tavily',
'',
].join('\n'),
'utf8',
);
const fetchMock = mockProviderFetch();
const app = await createConfigApp();

const result = await postJson(app, {
provider: 'tavily',
apiKey: '',
customEnv: { TAVILY_API_KEY: '********' },
});

expect(result.status).toBe(200);
expect(result.body.ok).toBe(true);
const requestBody = JSON.parse(fetchMock.mock.calls[0][1].body);
expect(requestBody.api_key).toBe('saved-tavily');
});

it('does not revive saved customEnv secrets that were removed from the current request', async () => {
mkdirSync(process.env.PILOT_HOME, { recursive: true });
writeFileSync(
join(process.env.PILOT_HOME, 'pilotdeck.yaml'),
[
'schemaVersion: 1',
'customEnv:',
' TAVILY_API_KEY: saved-tavily',
'',
].join('\n'),
'utf8',
);
const fetchMock = mockProviderFetch();
const app = await createConfigApp();

const result = await postJson(app, {
provider: 'tavily',
apiKey: '',
customEnv: {},
});

expect(result.status).toBe(400);
expect(result.body).toEqual({ ok: false, error: 'API key is required.' });
expect(fetchMock).not.toHaveBeenCalled();
});

it('prefers inline API keys over environment keys', async () => {
process.env.TAVILY_API_KEY = 'env-tavily';
const fetchMock = mockProviderFetch();
const app = await createConfigApp();

const result = await postJson(app, { provider: 'tavily', apiKey: ' inline-tavily ' });

expect(result.status).toBe(200);
expect(result.body.ok).toBe(true);
const requestBody = JSON.parse(fetchMock.mock.calls[0][1].body);
expect(requestBody.api_key).toBe('inline-tavily');
});

it('rejects missing credentials when no inline or environment key is available', async () => {
const fetchMock = mockProviderFetch();
const app = await createConfigApp();

const result = await postJson(app, { provider: 'glm', apiKey: '' });

expect(result.status).toBe(400);
expect(result.body).toEqual({ ok: false, error: 'API key is required.' });
expect(fetchMock).not.toHaveBeenCalled();
});

it('allows custom providers with auth none and sends no API key', async () => {
const fetchMock = mockProviderFetch();
const app = await createConfigApp();

const result = await postJson(app, {
provider: 'custom',
apiKey: '',
endpoint: 'https://example.com/search',
customProvider: { auth: 'none', method: 'GET' },
});

expect(result.status).toBe(200);
expect(result.body.ok).toBe(true);
expect(fetchMock.mock.calls[0][0]).toBe('https://example.com/search?query=hello');
expect(fetchMock.mock.calls[0][1].headers.Authorization).toBeUndefined();
});
});
Loading