Merge pull request #126 from rbuchmayer-pplx/feat-api-key-provider

feat: per-call API key provider for multi-tenant embedders (v1.2.0)

Author: rbuchmayer-pplx

.claude-plugin/marketplace.json

@@ -6,14 +6,14 @@
   },
   "metadata": {
     "description": "Official Perplexity AI plugin providing real-time web search, reasoning, and research capabilities",
-    "version": "1.1.0"
+    "version": "1.2.0"
   },
   "plugins": [
     {
       "name": "perplexity",
       "source": "./",
       "description": "Real-time web search, reasoning, and research through Perplexity's API",
-      "version": "1.1.0",
+      "version": "1.2.0",
       "author": {
         "name": "Perplexity AI",
         "email": "api@perplexity.ai"

README.md

@@ -147,6 +147,26 @@ npm install && npm run build && npm run start:http
 
 The server will be accessible at `http://localhost:8080/mcp`
 
+## Use as a Library
+
+The package also exports the server factory for embedding in your own Node process:
+
+```ts
+import { createPerplexityServer } from "@perplexity-ai/mcp-server";
+
+// Single-tenant: reads PERPLEXITY_API_KEY from the environment.
+const server = createPerplexityServer("my-service");
+
+// Multi-tenant hosts resolve the key per call instead. When a provider is
+// set, the environment variable is never consulted, and a provider that
+// returns no key fails the call rather than falling back.
+const tenantServer = createPerplexityServer("my-service", {
+  apiKey: () => currentRequestApiKey,
+});
+```
+
+Mount the returned server on any MCP transport (stdio, streamable HTTP, in-memory).
+
 ## Troubleshooting
 
 - **API Key Issues**: Ensure `PERPLEXITY_API_KEY` is set correctly

package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "@perplexity-ai/mcp-server",
-  "version": "1.1.0",
+  "version": "1.2.0",
   "mcpName": "ai.perplexity/mcp-server",
   "description": "Real-time web search, reasoning, and research through Perplexity's API",
   "keywords": [

package.json

@@ -3,12 +3,12 @@
   "name": "ai.perplexity/mcp-server",
   "title": "Perplexity API Platform",
   "description": "Real-time web search, reasoning, and research through Perplexity's API",
-  "version": "1.1.0",
+  "version": "1.2.0",
   "packages": [
     {
       "registryType": "npm",
       "identifier": "@perplexity-ai/mcp-server",
-      "version": "1.1.0",
+      "version": "1.2.0",
       "transport": {
         "type": "stdio"
       }

server.json

@@ -386,6 +386,41 @@ describe("Perplexity MCP Server", () => {
       });
     });
 
+    it("should use the provider key for the server-side cancel on timeout", async () => {
+      process.env.PERPLEXITY_TIMEOUT_MS = "100";
+      const authHeaders: string[] = [];
+
+      global.fetch = vi.fn().mockImplementation((url, options) => {
+        authHeaders.push(
+          (options?.headers as Record<string, string>)["Authorization"]
+        );
+        if (String(url).includes("/cancel")) {
+          return Promise.resolve({ ok: true, json: async () => ({}) } as unknown as Response);
+        }
+        const signal = options?.signal as AbortSignal | undefined;
+        const stream = new ReadableStream<Uint8Array>({
+          start(controller) {
+            controller.enqueue(
+              encodeSse([{ type: "response.created", response: { id: "resp_stall" } }])
+            );
+            signal?.addEventListener("abort", () => {
+              controller.error(
+                new DOMException("The operation was aborted.", "AbortError")
+              );
+            });
+          },
+        });
+        return Promise.resolve({ ok: true, body: stream } as unknown as Response);
+      });
+
+      await expect(
+        performAgentResponse(TEST_MESSAGES, "medium", undefined, undefined, undefined, () => "pplx-tenant-cancel")
+      ).rejects.toThrow("Request timeout");
+      await vi.waitFor(() => expect(authHeaders).toHaveLength(2));
+      // Both the original call and the fire-and-forget cancel carry the provider key.
+      expect(authHeaders).toEqual(["Bearer pplx-tenant-cancel", "Bearer pplx-tenant-cancel"]);
+    });
+
     it("should return the answer when the stream stays open after completion", async () => {
       process.env.PERPLEXITY_TIMEOUT_MS = "100";
       const cancelCalls: string[] = [];

src/index.test.ts

@@ -7,14 +7,18 @@ import type {
   AgentSearchResult,
   AgentToolOptions,
   AgentCallHooks,
+  ApiKeyProvider,
+  PerplexityServerOptions,
   SearchResponse,
   UndiciRequestOptions
 } from "./types.js";
 import { AgentResponseSchema, SearchResponseSchema } from "./validation.js";
 
+export type { ApiKeyProvider, PerplexityServerOptions } from "./types.js";
+
 const PERPLEXITY_API_KEY = process.env.PERPLEXITY_API_KEY;
 const PERPLEXITY_BASE_URL = process.env.PERPLEXITY_BASE_URL || "https://api.perplexity.ai";
-const VERSION = "1.1.0";
+const VERSION = "1.2.0";
 
 // Agent API presets backing each tool: https://docs.perplexity.ai/docs/agent-api/presets
 export const ASK_PRESET = "fast";
@@ -68,9 +72,21 @@ async function makeApiRequest(
   body: Record<string, unknown>,
   serviceOrigin: string | undefined,
   signal?: AbortSignal,
+  apiKey?: ApiKeyProvider,
 ): Promise<Response> {
-  if (!PERPLEXITY_API_KEY) {
-    throw new Error("PERPLEXITY_API_KEY environment variable is required");
+  // A configured provider fully replaces the env var: falling back would let
+  // a multi-tenant misconfiguration silently bill the process-wide key.
+  let resolvedApiKey: string | undefined;
+  if (apiKey) {
+    resolvedApiKey = apiKey();
+    if (!resolvedApiKey) {
+      throw new Error("API key provider returned no key");
+    }
+  } else {
+    resolvedApiKey = PERPLEXITY_API_KEY;
+    if (!resolvedApiKey) {
+      throw new Error("PERPLEXITY_API_KEY environment variable is required");
+    }
   }
 
   // Read timeout fresh each time to respect env var changes
@@ -92,7 +108,7 @@ async function makeApiRequest(
   try {
     const headers: Record<string, string> = {
       "Content-Type": "application/json",
-      "Authorization": `Bearer ${PERPLEXITY_API_KEY}`,
+      "Authorization": `Bearer ${resolvedApiKey}`,
       "User-Agent": `perplexity-mcp/${VERSION}`,
       "X-Source": "pplx-mcp-server",
     };
@@ -134,9 +150,9 @@ async function makeApiRequest(
 }
 
 /** Best-effort cancellation of an agent run so an abandoned request stops billing. */
-export async function cancelAgentResponse(responseId: string, serviceOrigin?: string): Promise<void> {
+export async function cancelAgentResponse(responseId: string, serviceOrigin?: string, apiKey?: ApiKeyProvider): Promise<void> {
   try {
-    await makeApiRequest(`v1/agent/${encodeURIComponent(responseId)}/cancel`, {}, serviceOrigin);
+    await makeApiRequest(`v1/agent/${encodeURIComponent(responseId)}/cancel`, {}, serviceOrigin, undefined, apiKey);
   } catch {
     // The run may already be terminal; nothing actionable either way.
   }
@@ -151,6 +167,7 @@ export async function consumeAgentStream(
   hooks?: AgentCallHooks,
   serviceOrigin?: string,
   deadlineSignal?: AbortSignal,
+  apiKey?: ApiKeyProvider,
 ): Promise<AgentResponse> {
   const body = response.body;
   if (!body) {
@@ -265,7 +282,7 @@ export async function consumeAgentStream(
     if (hooks?.signal?.aborted || deadlineSignal?.aborted) {
       if (responseId) {
         // Stop the server-side run so an abandoned request stops billing.
-        void cancelAgentResponse(responseId, serviceOrigin);
+        void cancelAgentResponse(responseId, serviceOrigin, apiKey);
       }
       if (hooks?.signal?.aborted) {
         throw new Error("Request cancelled");
@@ -277,7 +294,7 @@ export async function consumeAgentStream(
 
   if (hooks?.signal?.aborted) {
     if (responseId) {
-      void cancelAgentResponse(responseId, serviceOrigin);
+      void cancelAgentResponse(responseId, serviceOrigin, apiKey);
     }
     throw new Error("Request cancelled");
   }
@@ -385,7 +402,8 @@ export async function performAgentResponse(
   preset: string,
   serviceOrigin?: string,
   options?: AgentToolOptions,
-  hooks?: AgentCallHooks
+  hooks?: AgentCallHooks,
+  apiKey?: ApiKeyProvider,
 ): Promise<string> {
   const webSearchTool = buildWebSearchTool(options);
 
@@ -418,8 +436,8 @@ export async function performAgentResponse(
   }
 
   try {
-    const response = await makeApiRequest("v1/agent", body, serviceOrigin, deadline.signal);
-    const agentResponse = await consumeAgentStream(response, hooks, serviceOrigin, deadline.signal);
+    const response = await makeApiRequest("v1/agent", body, serviceOrigin, deadline.signal, apiKey);
+    const agentResponse = await consumeAgentStream(response, hooks, serviceOrigin, deadline.signal, apiKey);
     return formatAgentResponseText(agentResponse);
   } catch (error) {
     if (hooks?.signal?.aborted) {
@@ -463,7 +481,8 @@ export async function performSearch(
   maxTokensPerPage: number = 1024,
   country?: string,
   filters?: Pick<AgentToolOptions, "search_recency_filter" | "search_domain_filter">,
-  serviceOrigin?: string
+  serviceOrigin?: string,
+  apiKey?: ApiKeyProvider,
 ): Promise<string> {
   const body: Record<string, unknown> = {
     query: query,
@@ -474,7 +493,7 @@ export async function performSearch(
     ...(filters?.search_domain_filter && { search_domain_filter: filters.search_domain_filter }),
   };
 
-  const response = await makeApiRequest("search", body, serviceOrigin);
+  const response = await makeApiRequest("search", body, serviceOrigin, undefined, apiKey);
 
   let data: SearchResponse;
   try {
@@ -518,7 +537,7 @@ function buildHooks(extra: ToolExtra | undefined): AgentCallHooks {
   };
 }
 
-export function createPerplexityServer(serviceOrigin?: string) {
+export function createPerplexityServer(serviceOrigin?: string, serverOptions?: PerplexityServerOptions) {
   const server = new McpServer(
     {
       name: "ai.perplexity/mcp-server",
@@ -604,6 +623,7 @@ export function createPerplexityServer(serviceOrigin?: string) {
         serviceOrigin,
         Object.keys(options).length > 0 ? options : undefined,
         buildHooks(extra),
+        serverOptions?.apiKey,
       );
       return {
         content: [{ type: "text" as const, text: result }],
@@ -640,6 +660,7 @@ export function createPerplexityServer(serviceOrigin?: string) {
         serviceOrigin,
         undefined,
         buildHooks(extra),
+        serverOptions?.apiKey,
       );
       return {
         content: [{ type: "text" as const, text: result }],
@@ -686,6 +707,7 @@ export function createPerplexityServer(serviceOrigin?: string) {
         serviceOrigin,
         Object.keys(options).length > 0 ? options : undefined,
         buildHooks(extra),
+        serverOptions?.apiKey,
       );
       return {
         content: [{ type: "text" as const, text: result }],
@@ -745,7 +767,7 @@ export function createPerplexityServer(serviceOrigin?: string) {
         ...(search_domain_filter && { search_domain_filter }),
       };
 
-      const result = await performSearch(query, maxResults, maxTokensPerPage, countryCode, filters, serviceOrigin);
+      const result = await performSearch(query, maxResults, maxTokensPerPage, countryCode, filters, serviceOrigin, serverOptions?.apiKey);
       return {
         content: [{ type: "text" as const, text: result }],
         structuredContent: { results: result },

src/server.ts

@@ -1,5 +1,6 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import { createPerplexityServer, ASK_PRESET, REASON_PRESET, RESEARCH_PRESET } from "./server.js";
+import type { PerplexityServerOptions } from "./types.js";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
@@ -39,8 +40,8 @@ function agentSseResponse(text: string): Response {
   return { ok: true, body: stream } as unknown as Response;
 }
 
-async function connectInMemoryClient() {
-  const server = createPerplexityServer();
+async function connectInMemoryClient(serverOptions?: PerplexityServerOptions) {
+  const server = createPerplexityServer(undefined, serverOptions);
   const client = new Client({ name: "test-client", version: "1.0.0" });
   const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
   await Promise.all([
@@ -362,6 +363,51 @@ describe("Transport Integration Tests", () => {
       }
     });
 
+    it("should use the configured API key provider instead of the env key", async () => {
+      process.env.PERPLEXITY_API_KEY = "pplx-env-key-must-not-be-used";
+      global.fetch = vi.fn().mockResolvedValue(agentSseResponse("tenant answer"));
+
+      const { client, server } = await connectInMemoryClient({
+        apiKey: () => "pplx-tenant-a",
+      });
+      try {
+        const result: any = await client.callTool({
+          name: "perplexity_ask",
+          arguments: { messages: [{ role: "user", content: "test" }] },
+        });
+
+        expect(result.isError).toBeFalsy();
+        const headers = (global.fetch as ReturnType<typeof vi.fn>).mock.calls[0][1]
+          .headers as Record<string, string>;
+        expect(headers["Authorization"]).toBe("Bearer pplx-tenant-a");
+      } finally {
+        await client.close();
+        await server.close();
+      }
+    });
+
+    it("should fail the call when the API key provider returns no key", async () => {
+      process.env.PERPLEXITY_API_KEY = "pplx-env-key-must-not-be-used";
+      global.fetch = vi.fn();
+
+      const { client, server } = await connectInMemoryClient({
+        apiKey: () => undefined,
+      });
+      try {
+        const result: any = await client.callTool({
+          name: "perplexity_ask",
+          arguments: { messages: [{ role: "user", content: "test" }] },
+        });
+
+        expect(result.isError).toBe(true);
+        expect(result.content[0].text).toContain("API key provider returned no key");
+        expect(global.fetch).not.toHaveBeenCalled();
+      } finally {
+        await client.close();
+        await server.close();
+      }
+    });
+
     it("should forward search filters to the search API request body", async () => {
       global.fetch = vi.fn().mockResolvedValue({
         ok: true,

src/transport.test.ts

@@ -96,6 +96,23 @@ export interface AgentProgressUpdate {
   message: string;
 }
 
+/**
+ * Resolves the API key for an upstream Perplexity API call. Invoked once per
+ * request, so a closure over per-request state (e.g. an inbound Authorization
+ * header) gives each embedded server instance its own key.
+ */
+export type ApiKeyProvider = () => string | undefined;
+
+export interface PerplexityServerOptions {
+  /**
+   * Per-call API key resolution for embedders hosting the server for more
+   * than one key (multi-tenant). When set, the PERPLEXITY_API_KEY environment
+   * variable is never consulted: a provider that returns no key fails the
+   * call rather than silently falling back to the process-wide key.
+   */
+  apiKey?: ApiKeyProvider;
+}
+
 export interface AgentCallHooks {
   /** Abort signal from the MCP request; triggers server-side cancellation. */
   signal?: AbortSignal;