Refactor AI request headers for greater flexibility

Replaced `AI_CLIENT_ID` and `AI_USERNAME` environment variables with a single `AI_REQUEST_HEADERS` variable. This new variable accepts a JSON string, allowing users to define arbitrary custom headers for AI provider requests. This change enhances the extensibility of AI provider integrations.

Additionally, improved error handling for AI API call failures by providing more concise error messages and removing excessive logging of truncated response bodies.

Author: vinzee

packages/api/src/config.ts

@@ -53,8 +53,7 @@ export const AI_PROVIDER = env.AI_PROVIDER as string; // 'anthropic' | 'openai'
 export const AI_API_KEY = env.AI_API_KEY as string;
 export const AI_BASE_URL = env.AI_BASE_URL as string;
 export const AI_MODEL_NAME = env.AI_MODEL_NAME as string;
-export const AI_CLIENT_ID = env.AI_CLIENT_ID as string;
-export const AI_USERNAME = env.AI_USERNAME as string;
+export const AI_REQUEST_HEADERS = env.AI_REQUEST_HEADERS as string;
 
 // Legacy Anthropic-specific configuration (backward compatibility)
 export const ANTHROPIC_API_KEY = env.ANTHROPIC_API_KEY as string;

packages/api/src/controllers/__tests__/ai.test.ts

@@ -8,20 +8,22 @@ const mockOpenAIModel = {
   modelId: 'gpt-4o',
 } as unknown as LanguageModel;
 
-const mockAnthropicFactory = jest.fn(() => mockAnthropicModel);
-const mockCreateAnthropic = jest.fn(() => mockAnthropicFactory);
+const mockAnthropicFactory = jest.fn((_model?: string) => mockAnthropicModel);
+const mockCreateAnthropic = jest.fn(
+  (_opts?: Record<string, unknown>) => mockAnthropicFactory,
+);
 
-const mockOpenAIChatFactory = jest.fn(() => mockOpenAIModel);
-const mockCreateOpenAI = jest.fn(() => ({
+const mockOpenAIChatFactory = jest.fn((_model?: string) => mockOpenAIModel);
+const mockCreateOpenAI = jest.fn((_opts?: Record<string, unknown>) => ({
   chat: mockOpenAIChatFactory,
 }));
 
 jest.mock('@ai-sdk/anthropic', () => ({
-  createAnthropic: (...args: unknown[]) => mockCreateAnthropic(...args),
+  createAnthropic: (opts: Record<string, unknown>) => mockCreateAnthropic(opts),
 }));
 
 jest.mock('@ai-sdk/openai', () => ({
-  createOpenAI: (...args: unknown[]) => mockCreateOpenAI(...args),
+  createOpenAI: (opts: Record<string, unknown>) => mockCreateOpenAI(opts),
 }));
 
 jest.mock('@/utils/logger', () => ({
@@ -208,69 +210,43 @@ describe('openai provider', () => {
     );
   });
 
-  describe('custom headers', () => {
-    it('adds X-Client-Id header when AI_CLIENT_ID is set', () => {
+  describe('AI_REQUEST_HEADERS', () => {
+    it('passes parsed headers to createOpenAI', () => {
       setConfig({
         AI_PROVIDER: 'openai',
         AI_API_KEY: 'sk-test',
         AI_MODEL_NAME: 'gpt-4o',
-        AI_CLIENT_ID: 'MyApp',
+        AI_REQUEST_HEADERS: '{"X-Custom":"val1","X-Other":"val2"}',
       });
       getAIModel();
       expect(mockCreateOpenAI).toHaveBeenCalledWith(
         expect.objectContaining({
-          headers: expect.objectContaining({
-            'X-Client-Id': 'MyApp',
-          }),
+          headers: { 'X-Custom': 'val1', 'X-Other': 'val2' },
         }),
       );
     });
 
-    it('adds X-Username header when AI_USERNAME is set', () => {
+    it('throws when AI_REQUEST_HEADERS is invalid JSON', () => {
       setConfig({
         AI_PROVIDER: 'openai',
         AI_API_KEY: 'sk-test',
         AI_MODEL_NAME: 'gpt-4o',
-        AI_USERNAME: 'testuser',
+        AI_REQUEST_HEADERS: '{bad',
       });
-      getAIModel();
-      expect(mockCreateOpenAI).toHaveBeenCalledWith(
-        expect.objectContaining({
-          headers: expect.objectContaining({
-            'X-Username': 'testuser',
-          }),
-        }),
-      );
-    });
-
-    it('adds both headers when both are set', () => {
-      setConfig({
-        AI_PROVIDER: 'openai',
-        AI_API_KEY: 'sk-test',
-        AI_MODEL_NAME: 'gpt-4o',
-        AI_CLIENT_ID: 'MyApp',
-        AI_USERNAME: 'testuser',
-      });
-      getAIModel();
-      expect(mockCreateOpenAI).toHaveBeenCalledWith(
-        expect.objectContaining({
-          headers: {
-            'X-Client-Id': 'MyApp',
-            'X-Username': 'testuser',
-          },
-        }),
+      expect(() => getAIModel()).toThrow(
+        'AI_REQUEST_HEADERS is not valid JSON',
       );
     });
 
-    it('omits headers object when neither is set', () => {
+    it('omits headers when AI_REQUEST_HEADERS is not set', () => {
       setConfig({
         AI_PROVIDER: 'openai',
         AI_API_KEY: 'sk-test',
         AI_MODEL_NAME: 'gpt-4o',
       });
       getAIModel();
-      const call = mockCreateOpenAI.mock.calls[0][0];
-      expect(call.headers).toBeUndefined();
+      const call = mockCreateOpenAI.mock.calls[0]?.[0];
+      expect(call?.headers).toBeUndefined();
     });
   });
 });

packages/api/src/controllers/ai.ts

@@ -17,6 +17,7 @@ import z from 'zod';
 
 import * as config from '@/config';
 import { ISource } from '@/models/source';
+import { parseJSON } from '@/utils/common';
 import { Api500Error } from '@/utils/errors';
 import logger from '@/utils/logger';
 
@@ -385,13 +386,12 @@ function getOpenAIModel(): LanguageModel {
     );
   }
 
-  const headers: Record<string, string> = {};
-  if (config.AI_CLIENT_ID) {
-    headers['X-Client-Id'] = config.AI_CLIENT_ID;
-  }
-  if (config.AI_USERNAME) {
-    headers['X-Username'] = config.AI_USERNAME;
-  }
+  const headers: Record<string, string> = config.AI_REQUEST_HEADERS
+    ? parseJSON<Record<string, string>>(
+        config.AI_REQUEST_HEADERS,
+        'AI_REQUEST_HEADERS',
+      )
+    : {};
 
   const openai = createOpenAI({
     apiKey,

packages/api/src/routers/api/ai.ts

@@ -110,22 +110,9 @@ ${JSON.stringify(allFieldsWithKeys.slice(0, 200).map(f => ({ field: f.key, type:
         return res.json(chartConfig);
       } catch (err) {
         if (err instanceof APICallError) {
-          const truncatedBody =
-            typeof err.responseBody === 'string' &&
-            err.responseBody.length > 500
-              ? err.responseBody.slice(0, 500) + '...[truncated]'
-              : err.responseBody;
-          logger.error(
-            {
-              statusCode: err.statusCode,
-              responseBody: truncatedBody,
-              url: err.url,
-              isRetryable: err.isRetryable,
-              cause: err.cause instanceof Error ? err.cause.message : undefined,
-            },
-            'AI API call failed',
+          throw new Api500Error(
+            `AI Provider Error. Status: ${err.statusCode}. Message: ${err.message}`,
           );
-          throw new Api500Error(`AI Provider Error: ${err.message}`);
         }
         throw err;
       }

packages/api/src/utils/common.ts

@@ -27,6 +27,14 @@ export const tryJSONStringify = (json: Json) => {
   return result;
 };
 
+export function parseJSON<T = unknown>(raw: string, label: string): T {
+  try {
+    return JSON.parse(raw) as T;
+  } catch (e) {
+    throw new Error(`${label} is not valid JSON: ${(e as Error).message}`);
+  }
+}
+
 export const truncateString = (str: string, length: number) => {
   if (str.length > length) {
     return str.substring(0, length) + '...';