From 52c2e9eca5d1a5153e385a7084045872ff58b9fd Mon Sep 17 00:00:00 2001
From: Christian Gick <service@agiliton.eu>
Date: Fri, 10 Apr 2026 16:18:02 +0300
Subject: [PATCH] feat(agiliton-account): per-customer LiteLLM MCP server
 provisioning

At first login, provisionMcpServer() creates a sitebridge-{customer_id}
entry in LiteLLM DB via POST /v1/mcp/server, pointing to the customer's
demux URL. The virtual key is then scoped to ["sitebridge-{customer_id}"]
so LiteLLM routes tool calls only to that customer's WebSocket.

Also adds AGILITON_ACCOUNT_URL config for self-referencing in MCP URLs.

CF-3032

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
---
 .env.example        |  4 ++++
 src/auth/litellm.ts | 55 ++++++++++++++++++++++++++++++++++++++++++++-
 src/config.ts       |  3 +++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/.env.example b/.env.example
index 6986cb0..c1a64ba 100644
--- a/.env.example
+++ b/.env.example
@@ -19,6 +19,10 @@ JWT_EXPIRES_IN=7d
 # AES-256-GCM encryption for LiteLLM virtual keys (64 hex chars = 32 bytes)
 ENCRYPTION_KEY=changeme-generate-with-openssl-rand-hex-32
 
+# Self-referencing URL used when registering per-customer MCP servers in LiteLLM
+# LiteLLM calls this URL to forward tool calls to the right customer's WebSocket
+AGILITON_ACCOUNT_URL=http://agiliton-account:4100
+
 # LiteLLM
 LITELLM_URL=http://litellm:4000
 LITELLM_MASTER_KEY=sk-litellm-master-key
diff --git a/src/auth/litellm.ts b/src/auth/litellm.ts
index 3357086..00568f3 100644
--- a/src/auth/litellm.ts
+++ b/src/auth/litellm.ts
@@ -5,14 +5,67 @@ interface LiteLLMKeyResponse {
   key_alias: string;
 }
 
+interface LiteLLMMCPServerResponse {
+  server_id: string;
+  server_name: string;
+  alias: string;
+  url: string;
+}
+
+/**
+ * Register a per-customer sitebridge MCP server entry in LiteLLM DB.
+ * Returns the server_id (alias) used when provisioning the virtual key.
+ *
+ * LiteLLM does not support URL templates per virtual key (verified in 1.83.3),
+ * so we create one MCP server entry per customer with the customer-specific
+ * demux URL. The virtual key is then scoped to only that server.
+ */
+export async function provisionMcpServer(customerId: string): Promise<string> {
+  const serverAlias = `sitebridge-${customerId}`;
+  const demuxUrl = `${config.agilitonAccountUrl}/mcp/demux/${customerId}/mcp`;
+
+  const body = {
+    server_name: serverAlias,
+    alias: serverAlias,
+    description: `SiteBridge MCP bridge for customer ${customerId}`,
+    url: demuxUrl,
+    transport: "http",
+    auth_type: "bearer_token",
+    auth_value: config.mcpBridgeSecret,
+    allow_all_keys: false,
+  };
+
+  const res = await fetch(`${config.litellmUrl}/v1/mcp/server`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${config.litellmMasterKey}`,
+    },
+    body: JSON.stringify(body),
+  });
+
+  if (!res.ok) {
+    const text = await res.text();
+    throw new Error(`LiteLLM /v1/mcp/server failed ${res.status}: ${text}`);
+  }
+
+  return serverAlias;
+}
+
 /**
  * Provision a new virtual key for a customer. Called once at first login.
+ * Also creates a per-customer MCP server entry so LiteLLM routes sitebridge
+ * calls to the correct agiliton-account demux endpoint.
  */
 export async function provisionLiteLLMKey(customerId: string, email: string): Promise<string> {
+  // 1. Register per-customer MCP server
+  const mcpServerAlias = await provisionMcpServer(customerId);
+
+  // 2. Create virtual key scoped to that MCP server
   const body = {
     key_alias: `sb-${customerId}`,
     models: config.defaultModels,
-    mcp_servers: ["sitebridge"],
+    mcp_servers: [mcpServerAlias],
     max_budget: config.defaultBudgetUsd,
     budget_duration: config.defaultBudgetDuration,
     rpm_limit: config.defaultRpmLimit,
diff --git a/src/config.ts b/src/config.ts
index e8ae7c0..a131aff 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -31,6 +31,9 @@ export const config = {
   litellmUrl: envOptional("LITELLM_URL", "http://litellm:4000"),
   litellmMasterKey: env("LITELLM_MASTER_KEY"),
 
+  // Self-referencing URL for MCP demux (used in LiteLLM MCP server registration)
+  agilitonAccountUrl: envOptional("AGILITON_ACCOUNT_URL", "http://agiliton-account:4100"),
+
   // Shared secret used by LiteLLM when calling /mcp/demux/:id/mcp
   mcpBridgeSecret: env("MCP_BRIDGE_SECRET"),