feat(agiliton-account): Phase 1 service scaffold

TypeScript + Fastify service implementing:
- Google + Microsoft SSO (POST /v1/auth/sso/{google,microsoft})
- JWT issuance + LiteLLM virtual key provisioning on first login
- AES-256-GCM encrypted virtual key storage in Postgres
- Conversation CRUD (GET/POST/DELETE /v1/conversations, /messages)
- GDPR export + soft-delete (/v1/me/export, /v1/me/delete)
- WebSocket MCP bridge (/v1/mcp-bridge) with JWT auth
- MCP demux endpoint (/mcp/demux/:customer_id/mcp) for LiteLLM tool routing
- DB migration script creating sb_customers, sb_conversations, sb_messages
- 9 unit tests (bridge + crypto), all passing
- Dockerfile + docker-compose targeting port 4100

CF-3032

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Christian Gick
2026-04-10 16:06:16 +03:00
commit 7ab23554c0
26 changed files with 4855 additions and 0 deletions

42
.env.example Normal file
View File

@@ -0,0 +1,42 @@
# agiliton-account — copy to .env and fill in values
# Secrets are read from vault on the infra VM via entrypoint.sh
PORT=4100
HOST=0.0.0.0
LOG_LEVEL=info
NODE_ENV=production
# Database (shared agiliton Postgres)
DATABASE_URL=postgres://agiliton:PASSWORD@postgres:5432/agiliton
# Redis
REDIS_URL=redis://redis:6379
# JWT (32+ random bytes, hex)
JWT_SECRET=changeme-generate-with-openssl-rand-hex-32
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
# LiteLLM
LITELLM_URL=http://litellm:4000
LITELLM_MASTER_KEY=sk-litellm-master-key
# Shared secret for LiteLLM→agiliton-account MCP bridge calls
MCP_BRIDGE_SECRET=changeme-generate-with-openssl-rand-hex-32
# Google OAuth (create at console.cloud.google.com)
GOOGLE_CLIENT_ID=xxx.apps.googleusercontent.com
GOOGLE_CLIENT_SECRET=GOCSPX-xxx
# Microsoft OAuth (optional, create at portal.azure.com)
MS_CLIENT_ID=
MS_TENANT_ID=common
# Customer defaults
DEFAULT_BUDGET_USD=30.0
DEFAULT_BUDGET_DURATION=30d
DEFAULT_RPM_LIMIT=30
DEFAULT_MODELS=claude-sonnet-4-6,claude-opus-4-6,grok-heavy
MCP_BRIDGE_TIMEOUT_MS=15000

5
.gitignore vendored Normal file
View File

@@ -0,0 +1,5 @@
node_modules/
dist/
.env
*.env.local
*.log

23
Dockerfile Normal file
View File

@@ -0,0 +1,23 @@
FROM node:22-alpine AS builder
WORKDIR /app
COPY package*.json ./
RUN npm ci
COPY tsconfig.json ./
COPY src ./src
RUN npm run build
FROM node:22-alpine AS runtime
WORKDIR /app
COPY package*.json ./
RUN npm ci --omit=dev
COPY --from=builder /app/dist ./dist
ENV NODE_ENV=production
EXPOSE 4100
HEALTHCHECK --interval=30s --timeout=5s --start-period=10s \
CMD wget -qO- http://localhost:4100/health || exit 1
CMD ["node", "dist/index.js"]

20
docker-compose.yml Normal file
View File

@@ -0,0 +1,20 @@
services:
agiliton-account:
build: .
image: gitea.agiliton.internal/infrastructure/agiliton-account:latest
restart: unless-stopped
ports:
- "4100:4100"
env_file:
- .env
labels:
- "com.centurylinklabs.watchtower.enable=true"
networks:
- agiliton-internal
depends_on:
- postgres
- redis
networks:
agiliton-internal:
external: true

3366
package-lock.json generated Normal file

File diff suppressed because it is too large Load Diff

35
package.json Normal file
View File

@@ -0,0 +1,35 @@
{
"name": "agiliton-account",
"version": "1.0.0",
"description": "Agiliton customer account service — SSO auth, conversation history, MCP WebSocket bridge",
"type": "module",
"main": "dist/index.js",
"scripts": {
"build": "tsc",
"start": "node dist/index.js",
"dev": "tsx watch src/index.ts",
"test": "vitest run",
"test:watch": "vitest",
"migrate": "node dist/db/migrate.js",
"typecheck": "tsc --noEmit"
},
"dependencies": {
"@fastify/cors": "^9.0.0",
"@fastify/jwt": "^8.0.0",
"@fastify/websocket": "^10.0.0",
"fastify": "^4.28.0",
"google-auth-library": "^9.14.0",
"pg": "^8.12.0",
"redis": "^4.7.0",
"uuid": "^10.0.0",
"@mswjs/interceptors": "^0.37.0"
},
"devDependencies": {
"@types/node": "^22.0.0",
"@types/pg": "^8.11.0",
"@types/uuid": "^10.0.0",
"tsx": "^4.19.0",
"typescript": "^5.5.0",
"vitest": "^2.0.0"
}
}

View File

@@ -0,0 +1,117 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
// Mock config before importing bridge
vi.mock("../config.js", () => ({
config: {
mcpBridgeTimeoutMs: 1000,
jwtSecret: "test-secret",
encryptionKey: "0".repeat(64),
litellmUrl: "http://localhost:4000",
litellmMasterKey: "test-key",
mcpBridgeSecret: "bridge-secret",
defaultBudgetUsd: 30,
defaultBudgetDuration: "30d",
defaultRpmLimit: 30,
defaultModels: ["claude-sonnet-4-6"],
googleClientId: "test",
googleClientSecret: "test",
msClientId: "",
msTenantId: "common",
dbUrl: "postgres://localhost/test",
redisUrl: "redis://localhost:6379",
port: 4100,
host: "0.0.0.0",
jwtExpiresIn: "7d",
},
}));
import {
registerCustomerWs,
forwardMcpCall,
isCustomerOnline,
McpBridgeError,
} from "../mcp/bridge.js";
function makeMockWs() {
const listeners: Record<string, Array<(data: unknown) => void>> = {};
return {
send: vi.fn(),
close: vi.fn(),
on: (event: string, cb: (data: unknown) => void) => {
listeners[event] = listeners[event] ?? [];
listeners[event].push(cb);
},
emit: (event: string, data: unknown) => {
listeners[event]?.forEach((cb) => cb(data));
},
_listeners: listeners,
};
}
describe("MCP bridge", () => {
beforeEach(() => {
// Clear sessions between tests by re-importing is complex; use unique IDs
});
it("registers a customer WS and marks them online", () => {
const ws = makeMockWs();
registerCustomerWs("cust-1", ws as never);
expect(isCustomerOnline("cust-1")).toBe(true);
});
it("forwards a call and resolves on WS response", async () => {
const ws = makeMockWs();
registerCustomerWs("cust-2", ws as never);
const callPromise = forwardMcpCall(
"cust-2",
{ method: "tools/call", params: { name: "screenshot" }, id: "req-1" },
5000
);
// Simulate extension responding
ws.emit("message", Buffer.from(JSON.stringify({ id: "req-1", result: { image: "base64data" } })));
const result = await callPromise;
expect(result).toEqual({ image: "base64data" });
});
it("rejects with McpBridgeError when customer is offline", async () => {
await expect(
forwardMcpCall("offline-cust", { method: "tools/call", id: "req-2" }, 100)
).rejects.toThrow(McpBridgeError);
});
it("rejects with McpBridgeError on timeout", async () => {
const ws = makeMockWs();
registerCustomerWs("cust-3", ws as never);
await expect(
forwardMcpCall(
"cust-3",
{ method: "tools/call", id: "req-3" },
50 // 50ms timeout
)
).rejects.toThrow("timeout");
});
it("marks customer offline after WS close", () => {
const ws = makeMockWs();
registerCustomerWs("cust-4", ws as never);
expect(isCustomerOnline("cust-4")).toBe(true);
ws.emit("close", undefined);
expect(isCustomerOnline("cust-4")).toBe(false);
});
it("replaces stale WS on second registration", () => {
const ws1 = makeMockWs();
const ws2 = makeMockWs();
registerCustomerWs("cust-5", ws1 as never);
registerCustomerWs("cust-5", ws2 as never);
expect(ws1.close).toHaveBeenCalled();
expect(isCustomerOnline("cust-5")).toBe(true);
});
});

View File

@@ -0,0 +1,30 @@
import { describe, it, expect, vi } from "vitest";
vi.mock("../config.js", () => ({
config: {
encryptionKey: "a".repeat(64), // 32-byte hex key
},
}));
import { encrypt, decrypt } from "../utils/crypto.js";
describe("AES-256-GCM encrypt/decrypt", () => {
it("round-trips a plaintext string", () => {
const original = "sk-litellm-abc123supersecret";
const buf = encrypt(original);
expect(buf.length).toBeGreaterThan(28);
expect(decrypt(buf)).toBe(original);
});
it("produces different ciphertexts for the same input (random IV)", () => {
const key = "same-key-test";
const buf1 = encrypt(key);
const buf2 = encrypt(key);
expect(buf1.toString("hex")).not.toBe(buf2.toString("hex"));
});
it("round-trips a unicode string", () => {
const original = "sk-🔑-unicode-test-αβγ";
expect(decrypt(encrypt(original))).toBe(original);
});
});

157
src/auth/customers.ts Normal file
View File

@@ -0,0 +1,157 @@
import { query, queryOne } from "../db/pool.js";
import { encrypt, decrypt } from "../utils/crypto.js";
import { provisionLiteLLMKey } from "./litellm.js";
export interface Customer {
id: string;
email: string;
name: string | null;
sso_provider: string;
sso_sub: string;
tier: string;
system_prompt: string | null;
preferred_model: string | null;
litellm_key_alias: string;
litellm_key_encrypted: Buffer;
created_at: Date;
last_login_at: Date | null;
deleted_at: Date | null;
}
export interface CustomerPublic {
id: string;
email: string;
name: string | null;
tier: string;
system_prompt: string | null;
preferred_model: string | null;
}
export function toPublic(c: Customer): CustomerPublic {
return {
id: c.id,
email: c.email,
name: c.name,
tier: c.tier,
system_prompt: c.system_prompt,
preferred_model: c.preferred_model,
};
}
export function getVirtualKey(c: Customer): string {
return decrypt(c.litellm_key_encrypted);
}
/**
* Upsert a customer from SSO data. Provisions a LiteLLM key on first login.
* Returns the customer row and the plaintext virtual key.
*/
export async function upsertCustomer(params: {
sso_provider: string;
sso_sub: string;
email: string;
name: string;
}): Promise<{ customer: Customer; virtualKey: string; isNew: boolean }> {
// Check for existing customer (not deleted)
const existing = await queryOne<Customer>(
`SELECT * FROM sb_customers
WHERE sso_provider = $1 AND sso_sub = $2 AND deleted_at IS NULL`,
[params.sso_provider, params.sso_sub]
);
if (existing) {
// Update last_login_at and name
await query(
`UPDATE sb_customers SET last_login_at = NOW(), name = $1 WHERE id = $2`,
[params.name, existing.id]
);
existing.last_login_at = new Date();
existing.name = params.name;
return { customer: existing, virtualKey: getVirtualKey(existing), isNew: false };
}
// First login — provision LiteLLM key, then insert
// We need the customer ID before the key alias, so generate a UUID first
const { v4: uuidv4 } = await import("uuid");
const customerId = uuidv4();
const virtualKey = await provisionLiteLLMKey(customerId, params.email);
const keyAlias = `sb-${customerId}`;
const encryptedKey = encrypt(virtualKey);
const [customer] = await query<Customer>(
`INSERT INTO sb_customers
(id, email, name, sso_provider, sso_sub, litellm_key_alias, litellm_key_encrypted, last_login_at)
VALUES ($1, $2, $3, $4, $5, $6, $7, NOW())
RETURNING *`,
[customerId, params.email, params.name, params.sso_provider, params.sso_sub, keyAlias, encryptedKey]
);
return { customer, virtualKey, isNew: true };
}
export async function getCustomerById(id: string): Promise<Customer | null> {
return queryOne<Customer>(
`SELECT * FROM sb_customers WHERE id = $1 AND deleted_at IS NULL`,
[id]
);
}
export async function updateCustomerConfig(
id: string,
patch: { system_prompt?: string | null; preferred_model?: string | null }
): Promise<void> {
const fields: string[] = [];
const values: unknown[] = [];
let i = 1;
if ("system_prompt" in patch) {
fields.push(`system_prompt = $${i++}`);
values.push(patch.system_prompt);
}
if ("preferred_model" in patch) {
fields.push(`preferred_model = $${i++}`);
values.push(patch.preferred_model);
}
if (fields.length === 0) return;
values.push(id);
await query(
`UPDATE sb_customers SET ${fields.join(", ")} WHERE id = $${i}`,
values
);
}
export async function softDeleteCustomer(id: string): Promise<void> {
await query(
`UPDATE sb_customers SET deleted_at = NOW() WHERE id = $1`,
[id]
);
}
export async function exportCustomerData(id: string): Promise<object> {
const customer = await getCustomerById(id);
if (!customer) return {};
const conversations = await query(
`SELECT id, title, created_at, updated_at FROM sb_conversations
WHERE customer_id = $1 ORDER BY updated_at DESC`,
[id]
);
const messages = await query(
`SELECT m.id, m.conversation_id, m.role, m.content, m.created_at
FROM sb_messages m
JOIN sb_conversations c ON c.id = m.conversation_id
WHERE c.customer_id = $1
ORDER BY m.created_at`,
[id]
);
return {
customer: toPublic(customer),
conversations,
messages,
exported_at: new Date().toISOString(),
};
}

40
src/auth/google.ts Normal file
View File

@@ -0,0 +1,40 @@
import { OAuth2Client } from "google-auth-library";
import { config } from "../config.js";
const client = new OAuth2Client(config.googleClientId, config.googleClientSecret);
export interface GoogleProfile {
sub: string;
email: string;
name: string;
}
/**
* Exchange an authorization code from chrome.identity.launchWebAuthFlow for
* a Google ID token, then verify and return the user profile.
*
* redirectUri must match what was registered in the Google Cloud Console AND
* what chrome.identity uses (typically https://<ext-id>.chromiumapp.org/).
*/
export async function exchangeGoogleCode(
code: string,
redirectUri: string
): Promise<GoogleProfile> {
const { tokens } = await client.getToken({ code, redirect_uri: redirectUri });
if (!tokens.id_token) throw new Error("No id_token in Google response");
const ticket = await client.verifyIdToken({
idToken: tokens.id_token,
audience: config.googleClientId,
});
const payload = ticket.getPayload();
if (!payload?.sub || !payload.email) {
throw new Error("Invalid Google token payload");
}
return {
sub: payload.sub,
email: payload.email,
name: payload.name ?? payload.email,
};
}

52
src/auth/litellm.ts Normal file
View File

@@ -0,0 +1,52 @@
import { config } from "../config.js";
interface LiteLLMKeyResponse {
key: string;
key_alias: string;
}
/**
* Provision a new virtual key for a customer. Called once at first login.
*/
export async function provisionLiteLLMKey(customerId: string, email: string): Promise<string> {
const body = {
key_alias: `sb-${customerId}`,
models: config.defaultModels,
mcp_servers: ["sitebridge"],
max_budget: config.defaultBudgetUsd,
budget_duration: config.defaultBudgetDuration,
rpm_limit: config.defaultRpmLimit,
metadata: { customer_id: customerId, email },
};
const res = await fetch(`${config.litellmUrl}/key/generate`, {
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 key/generate failed ${res.status}: ${text}`);
}
const data = (await res.json()) as LiteLLMKeyResponse;
return data.key;
}
/**
* Look up a virtual key's metadata via LiteLLM.
*/
export async function getLiteLLMKeyInfo(keyAlias: string): Promise<unknown> {
const res = await fetch(
`${config.litellmUrl}/key/info?key_alias=${encodeURIComponent(keyAlias)}`,
{
headers: { Authorization: `Bearer ${config.litellmMasterKey}` },
}
);
if (!res.ok) throw new Error(`LiteLLM key/info failed ${res.status}`);
return res.json();
}

56
src/config.ts Normal file
View File

@@ -0,0 +1,56 @@
import { readFileSync } from "fs";
function env(key: string, fallback?: string): string {
const val = process.env[key] ?? fallback;
if (val === undefined) throw new Error(`Missing required env var: ${key}`);
return val;
}
function envOptional(key: string, fallback = ""): string {
return process.env[key] ?? fallback;
}
export const config = {
port: parseInt(envOptional("PORT", "4100")),
host: envOptional("HOST", "0.0.0.0"),
// JWT
jwtSecret: env("JWT_SECRET"),
jwtExpiresIn: envOptional("JWT_EXPIRES_IN", "7d"),
// Database
dbUrl: env("DATABASE_URL"),
// Redis
redisUrl: envOptional("REDIS_URL", "redis://localhost:6379"),
// Encryption key for LiteLLM virtual keys stored in DB (32-byte hex)
encryptionKey: env("ENCRYPTION_KEY"),
// LiteLLM
litellmUrl: envOptional("LITELLM_URL", "http://litellm:4000"),
litellmMasterKey: env("LITELLM_MASTER_KEY"),
// Shared secret used by LiteLLM when calling /mcp/demux/:id/mcp
mcpBridgeSecret: env("MCP_BRIDGE_SECRET"),
// Google OAuth
googleClientId: env("GOOGLE_CLIENT_ID"),
googleClientSecret: env("GOOGLE_CLIENT_SECRET"),
// Microsoft OAuth
msClientId: envOptional("MS_CLIENT_ID"),
msTenantId: envOptional("MS_TENANT_ID", "common"),
// Default budget per customer (USD, 30 days)
defaultBudgetUsd: parseFloat(envOptional("DEFAULT_BUDGET_USD", "30.0")),
defaultBudgetDuration: envOptional("DEFAULT_BUDGET_DURATION", "30d"),
defaultRpmLimit: parseInt(envOptional("DEFAULT_RPM_LIMIT", "30")),
defaultModels: envOptional(
"DEFAULT_MODELS",
"claude-sonnet-4-6,claude-opus-4-6,grok-heavy"
).split(","),
// WebSocket MCP bridge timeout (ms)
mcpBridgeTimeoutMs: parseInt(envOptional("MCP_BRIDGE_TIMEOUT_MS", "15000")),
} as const;

109
src/db/conversations.ts Normal file
View File

@@ -0,0 +1,109 @@
import { query, queryOne } from "./pool.js";
export interface Conversation {
id: string;
customer_id: string;
title: string | null;
created_at: Date;
updated_at: Date;
}
export interface Message {
id: string;
conversation_id: string;
role: string;
content: unknown;
response_id: string | null;
created_at: Date;
}
export async function listConversations(
customerId: string,
limit = 50,
offset = 0
): Promise<Conversation[]> {
return query<Conversation>(
`SELECT id, customer_id, title, created_at, updated_at
FROM sb_conversations
WHERE customer_id = $1
ORDER BY updated_at DESC
LIMIT $2 OFFSET $3`,
[customerId, limit, offset]
);
}
export async function createConversation(
customerId: string,
title?: string
): Promise<Conversation> {
const [row] = await query<Conversation>(
`INSERT INTO sb_conversations (customer_id, title)
VALUES ($1, $2)
RETURNING *`,
[customerId, title ?? null]
);
return row;
}
export async function getConversation(
id: string,
customerId: string
): Promise<Conversation | null> {
return queryOne<Conversation>(
`SELECT * FROM sb_conversations WHERE id = $1 AND customer_id = $2`,
[id, customerId]
);
}
export async function deleteConversation(
id: string,
customerId: string
): Promise<boolean> {
const { rowCount } = await (await import("./pool.js")).pool.query(
`DELETE FROM sb_conversations WHERE id = $1 AND customer_id = $2`,
[id, customerId]
);
return (rowCount ?? 0) > 0;
}
export async function getMessages(conversationId: string): Promise<Message[]> {
return query<Message>(
`SELECT id, conversation_id, role, content, response_id, created_at
FROM sb_messages
WHERE conversation_id = $1
ORDER BY created_at`,
[conversationId]
);
}
export async function appendMessage(params: {
conversationId: string;
role: string;
content: unknown;
responseId?: string;
}): Promise<Message> {
const [msg] = await query<Message>(
`INSERT INTO sb_messages (conversation_id, role, content, response_id)
VALUES ($1, $2, $3, $4)
RETURNING *`,
[params.conversationId, params.role, JSON.stringify(params.content), params.responseId ?? null]
);
// Touch updated_at on the conversation
await query(
`UPDATE sb_conversations SET updated_at = NOW() WHERE id = $1`,
[params.conversationId]
);
return msg;
}
export async function updateConversationTitle(
id: string,
title: string
): Promise<void> {
await query(
`UPDATE sb_conversations SET title = $1, updated_at = NOW() WHERE id = $2`,
[title, id]
);
}

69
src/db/migrate.ts Normal file
View File

@@ -0,0 +1,69 @@
/**
* Run this once to create the agiliton-account tables in the agiliton DB.
* Usage: npm run migrate
*/
import { pool } from "./pool.js";
const MIGRATION = `
-- Customers (one row per SSO identity)
CREATE TABLE IF NOT EXISTS sb_customers (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
email TEXT UNIQUE NOT NULL,
name TEXT,
sso_provider TEXT NOT NULL, -- 'google' | 'microsoft'
sso_sub TEXT NOT NULL, -- provider subject id
tier TEXT NOT NULL DEFAULT 'free',
system_prompt TEXT,
preferred_model TEXT,
litellm_key_alias TEXT NOT NULL, -- e.g. sb-<customer_id>
litellm_key_encrypted BYTEA NOT NULL,-- virtual key, AES-256-GCM encrypted
created_at TIMESTAMPTZ DEFAULT NOW(),
last_login_at TIMESTAMPTZ,
deleted_at TIMESTAMPTZ, -- soft delete for GDPR grace period
UNIQUE (sso_provider, sso_sub)
);
-- Conversations (per customer)
CREATE TABLE IF NOT EXISTS sb_conversations (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
customer_id UUID NOT NULL REFERENCES sb_customers(id) ON DELETE CASCADE,
title TEXT,
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_sb_conv_customer
ON sb_conversations (customer_id, updated_at DESC);
-- Messages (per conversation)
CREATE TABLE IF NOT EXISTS sb_messages (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
conversation_id UUID NOT NULL REFERENCES sb_conversations(id) ON DELETE CASCADE,
role TEXT NOT NULL, -- 'user' | 'assistant' | 'tool'
content JSONB NOT NULL, -- full OpenAI-style content blocks
response_id TEXT, -- LiteLLM /v1/responses id for threading
created_at TIMESTAMPTZ DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_sb_msg_conv
ON sb_messages (conversation_id, created_at);
`;
async function migrate() {
const client = await pool.connect();
try {
await client.query("BEGIN");
await client.query(MIGRATION);
await client.query("COMMIT");
console.log("[migrate] done");
} catch (err) {
await client.query("ROLLBACK");
console.error("[migrate] failed", err);
process.exit(1);
} finally {
client.release();
await pool.end();
}
}
migrate();

24
src/db/pool.ts Normal file
View File

@@ -0,0 +1,24 @@
import { Pool } from "pg";
import { config } from "../config.js";
export const pool = new Pool({ connectionString: config.dbUrl });
pool.on("error", (err) => {
console.error("[db] unexpected pool error", err);
});
export async function query<T extends object = Record<string, unknown>>(
sql: string,
params?: unknown[]
): Promise<T[]> {
const { rows } = await pool.query<T>(sql, params);
return rows;
}
export async function queryOne<T extends object = Record<string, unknown>>(
sql: string,
params?: unknown[]
): Promise<T | null> {
const rows = await query<T>(sql, params);
return rows[0] ?? null;
}

64
src/index.ts Normal file
View File

@@ -0,0 +1,64 @@
import Fastify from "fastify";
import cors from "@fastify/cors";
import jwt from "@fastify/jwt";
import websocket from "@fastify/websocket";
import { config } from "./config.js";
import { pool } from "./db/pool.js";
import { ssoRoutes } from "./routes/sso.js";
import { meRoutes } from "./routes/me.js";
import { conversationRoutes } from "./routes/conversations.js";
import { mcpBridgeRoute } from "./routes/mcpBridge.js";
import { mcpDemuxRoutes } from "./routes/mcpDemux.js";
import { getBridgeStats } from "./mcp/bridge.js";
const app = Fastify({
logger: {
level: process.env.LOG_LEVEL ?? "info",
transport:
process.env.NODE_ENV !== "production"
? { target: "pino-pretty", options: { colorize: true } }
: undefined,
},
});
await app.register(cors, {
origin: [
// Allow chrome-extension:// origins from the SiteBridge extension
/^chrome-extension:\/\//,
// Allow agiliton.cloud subdomains for future web UI
/^https:\/\/.*\.agiliton\.cloud$/,
"https://agiliton.eu",
],
credentials: true,
});
await app.register(jwt, { secret: config.jwtSecret });
await app.register(websocket);
// Routes
await app.register(ssoRoutes);
await app.register(meRoutes);
await app.register(conversationRoutes);
await app.register(mcpBridgeRoute);
await app.register(mcpDemuxRoutes);
// Health check
app.get("/health", async () => ({
status: "ok",
bridge: getBridgeStats(),
}));
// Graceful shutdown
const shutdown = async () => {
console.log("[agiliton-account] shutting down...");
await app.close();
await pool.end();
process.exit(0);
};
process.on("SIGTERM", shutdown);
process.on("SIGINT", shutdown);
await app.listen({ port: config.port, host: config.host });
console.log(`[agiliton-account] listening on ${config.host}:${config.port}`);

124
src/mcp/bridge.ts Normal file
View File

@@ -0,0 +1,124 @@
/**
* WebSocket MCP bridge.
*
* Each authenticated extension connects here. We maintain a map of
* customer_id → WebSocket. When LiteLLM fires an MCP tool call via
* POST /mcp/demux/:customer_id/mcp, we forward the JSON-RPC request
* over the customer's WebSocket and await the result within a timeout.
*
* Protocol:
* LiteLLM → agiliton-account → (WS) → SiteBridge extension
* SiteBridge extension → (WS) → agiliton-account → LiteLLM
*
* Messages on the WebSocket are JSON-RPC 2.0 objects.
*/
import type { WebSocket } from "@fastify/websocket";
interface PendingCall {
resolve: (result: unknown) => void;
reject: (err: Error) => void;
timer: NodeJS.Timeout;
}
// In-memory map: customer_id → ws + pending calls
const sessions = new Map<
string,
{ ws: WebSocket; pending: Map<string | number, PendingCall> }
>();
export function registerCustomerWs(customerId: string, ws: WebSocket): void {
// Clean up any stale session for this customer
const existing = sessions.get(customerId);
if (existing) {
console.log(`[mcp-bridge] replacing stale ws for customer=${customerId}`);
try { existing.ws.close(); } catch {}
existing.pending.forEach((p) => p.reject(new Error("WebSocket replaced")));
}
const pending = new Map<string | number, PendingCall>();
sessions.set(customerId, { ws, pending });
console.log(`[mcp-bridge] customer=${customerId} connected (total=${sessions.size})`);
ws.on("message", (raw: Buffer | string) => {
let msg: { id?: string | number; result?: unknown; error?: unknown };
try {
msg = JSON.parse(raw.toString());
} catch {
console.warn(`[mcp-bridge] customer=${customerId} sent unparseable message`);
return;
}
if (msg.id !== undefined) {
const call = pending.get(msg.id);
if (call) {
clearTimeout(call.timer);
pending.delete(msg.id);
if (msg.error) {
call.reject(new Error(JSON.stringify(msg.error)));
} else {
call.resolve(msg.result);
}
}
}
});
ws.on("close", () => {
const session = sessions.get(customerId);
if (session?.ws === ws) {
sessions.delete(customerId);
session!.pending.forEach((p) => p.reject(new Error("WebSocket closed")));
console.log(`[mcp-bridge] customer=${customerId} disconnected (total=${sessions.size})`);
}
});
}
export function isCustomerOnline(customerId: string): boolean {
return sessions.has(customerId);
}
/**
* Forward a JSON-RPC MCP call to the customer's WebSocket and return the result.
* Throws if customer is offline or times out.
*/
export async function forwardMcpCall(
customerId: string,
request: { method: string; params?: unknown; id: string | number },
timeoutMs: number
): Promise<unknown> {
const session = sessions.get(customerId);
if (!session) {
throw new McpBridgeError("tool_unavailable: customer not online", -32001);
}
return new Promise((resolve, reject) => {
const timer = setTimeout(() => {
session.pending.delete(request.id);
reject(new McpBridgeError(`tool_unavailable: timeout after ${timeoutMs}ms`, -32002));
}, timeoutMs);
session.pending.set(request.id, { resolve, reject, timer });
try {
session.ws.send(JSON.stringify({ jsonrpc: "2.0", ...request }));
} catch (err) {
clearTimeout(timer);
session.pending.delete(request.id);
reject(err);
}
});
}
export class McpBridgeError extends Error {
constructor(
message: string,
public readonly code: number = -32000
) {
super(message);
this.name = "McpBridgeError";
}
}
export function getBridgeStats() {
return { online: sessions.size };
}

27
src/middleware/auth.ts Normal file
View File

@@ -0,0 +1,27 @@
import type { FastifyRequest, FastifyReply } from "fastify";
import { getCustomerById } from "../auth/customers.js";
import type { Customer } from "../auth/customers.js";
declare module "fastify" {
interface FastifyRequest {
customer?: Customer;
}
}
export async function requireAuth(
req: FastifyRequest,
reply: FastifyReply
): Promise<void> {
try {
// @fastify/jwt decorates verify on the request
await req.jwtVerify();
const payload = req.user as { sub: string };
const customer = await getCustomerById(payload.sub);
if (!customer) {
return reply.code(401).send({ error: "customer_not_found" });
}
req.customer = customer;
} catch {
return reply.code(401).send({ error: "unauthorized" });
}
}

View File

@@ -0,0 +1,95 @@
import type { FastifyInstance } from "fastify";
import { requireAuth } from "../middleware/auth.js";
import {
listConversations,
createConversation,
getConversation,
getMessages,
appendMessage,
deleteConversation,
} from "../db/conversations.js";
export async function conversationRoutes(app: FastifyInstance) {
app.addHook("onRequest", requireAuth);
/** GET /v1/conversations?limit=50&offset=0 */
app.get<{ Querystring: { limit?: string; offset?: string } }>(
"/v1/conversations",
async (req) => {
const limit = Math.min(parseInt(req.query.limit ?? "50"), 200);
const offset = parseInt(req.query.offset ?? "0");
return listConversations(req.customer!.id, limit, offset);
}
);
/** POST /v1/conversations */
app.post<{ Body: { title?: string } }>(
"/v1/conversations",
{
schema: {
body: {
type: "object",
properties: { title: { type: "string" } },
},
},
},
async (req, reply) => {
const conv = await createConversation(req.customer!.id, req.body.title);
return reply.code(201).send(conv);
}
);
/** GET /v1/conversations/:id */
app.get<{ Params: { id: string } }>(
"/v1/conversations/:id",
async (req, reply) => {
const conv = await getConversation(req.params.id, req.customer!.id);
if (!conv) return reply.code(404).send({ error: "not_found" });
const messages = await getMessages(req.params.id);
return { ...conv, messages };
}
);
/** POST /v1/conversations/:id/messages */
app.post<{
Params: { id: string };
Body: { role: string; content: unknown; response_id?: string };
}>(
"/v1/conversations/:id/messages",
{
schema: {
body: {
type: "object",
required: ["role", "content"],
properties: {
role: { type: "string", enum: ["user", "assistant", "tool"] },
content: {},
response_id: { type: "string" },
},
},
},
},
async (req, reply) => {
const conv = await getConversation(req.params.id, req.customer!.id);
if (!conv) return reply.code(404).send({ error: "not_found" });
const msg = await appendMessage({
conversationId: req.params.id,
role: req.body.role,
content: req.body.content,
responseId: req.body.response_id,
});
return reply.code(201).send(msg);
}
);
/** DELETE /v1/conversations/:id */
app.delete<{ Params: { id: string } }>(
"/v1/conversations/:id",
async (req, reply) => {
const deleted = await deleteConversation(req.params.id, req.customer!.id);
if (!deleted) return reply.code(404).send({ error: "not_found" });
return reply.code(204).send();
}
);
}

61
src/routes/mcpBridge.ts Normal file
View File

@@ -0,0 +1,61 @@
/**
* WS /v1/mcp-bridge
*
* JWT-authenticated WebSocket endpoint. On connect, registers the customer's
* extension as the MCP server for that customer. MCP JSON-RPC calls arrive
* from LiteLLM via /mcp/demux/:id/mcp and are forwarded here.
*/
import type { FastifyInstance } from "fastify";
import type { WebSocket } from "@fastify/websocket";
import { registerCustomerWs } from "../mcp/bridge.js";
import { getCustomerById } from "../auth/customers.js";
export async function mcpBridgeRoute(app: FastifyInstance) {
app.get(
"/v1/mcp-bridge",
{ websocket: true },
async (socket: WebSocket, req) => {
// Verify JWT from query param or Authorization header
// chrome.identity WebSocket handshakes can't set headers, so we also
// accept ?token=<jwt> as a query param.
let customerId: string | null = null;
try {
const url = new URL(req.url, "http://localhost");
const tokenParam = url.searchParams.get("token");
let payload: { sub: string };
if (tokenParam) {
payload = app.jwt.verify<{ sub: string }>(tokenParam);
} else {
await req.jwtVerify();
payload = req.user as { sub: string };
}
const customer = await getCustomerById(payload.sub);
if (!customer) throw new Error("customer not found");
customerId = customer.id;
} catch (err) {
socket.send(
JSON.stringify({
jsonrpc: "2.0",
error: { code: -32001, message: "unauthorized" },
})
);
socket.close(1008, "unauthorized");
return;
}
registerCustomerWs(customerId, socket);
// Send a ready signal so the extension knows the bridge is active
socket.send(
JSON.stringify({
jsonrpc: "2.0",
method: "bridge/ready",
params: { customer_id: customerId },
})
);
}
);
}

74
src/routes/mcpDemux.ts Normal file
View File

@@ -0,0 +1,74 @@
/**
* POST /mcp/demux/:customer_id/mcp
*
* Called by LiteLLM when it needs to execute a sitebridge MCP tool call.
* Forwards the JSON-RPC request to the customer's WebSocket and returns the result.
*
* Authentication: shared-secret bearer token (MCP_BRIDGE_SECRET env var).
* This endpoint is NOT exposed to customers — it's internal LiteLLM→agiliton-account only.
*/
import type { FastifyInstance } from "fastify";
import { forwardMcpCall, isCustomerOnline, McpBridgeError } from "../mcp/bridge.js";
import { config } from "../config.js";
import { randomUUID } from "crypto";
export async function mcpDemuxRoutes(app: FastifyInstance) {
app.post<{ Params: { customer_id: string } }>(
"/mcp/demux/:customer_id/mcp",
async (req, reply) => {
// Verify shared secret
const authHeader = req.headers.authorization ?? "";
const token = authHeader.startsWith("Bearer ") ? authHeader.slice(7) : "";
if (token !== config.mcpBridgeSecret) {
return reply.code(401).send({ error: "unauthorized" });
}
const { customer_id } = req.params;
if (!isCustomerOnline(customer_id)) {
// Return a JSON-RPC error in MCP format
return reply.code(200).send({
jsonrpc: "2.0",
id: (req.body as { id?: unknown })?.id ?? null,
error: {
code: -32001,
message: "tool_unavailable: customer not online",
},
});
}
const body = req.body as {
method: string;
params?: unknown;
id?: string | number;
};
const callId = body.id ?? randomUUID();
try {
const result = await forwardMcpCall(
customer_id,
{ method: body.method, params: body.params, id: callId },
config.mcpBridgeTimeoutMs
);
return reply.send({
jsonrpc: "2.0",
id: callId,
result,
});
} catch (err) {
const code = err instanceof McpBridgeError ? err.code : -32000;
const message = err instanceof Error ? err.message : "internal_error";
req.log.warn({ err, customer_id }, "MCP demux call failed");
return reply.code(200).send({
jsonrpc: "2.0",
id: callId,
error: { code, message },
});
}
}
);
}

53
src/routes/me.ts Normal file
View File

@@ -0,0 +1,53 @@
import type { FastifyInstance } from "fastify";
import { requireAuth } from "../middleware/auth.js";
import {
toPublic,
updateCustomerConfig,
softDeleteCustomer,
exportCustomerData,
getVirtualKey,
} from "../auth/customers.js";
export async function meRoutes(app: FastifyInstance) {
app.addHook("onRequest", requireAuth);
/** GET /v1/me */
app.get("/v1/me", async (req, reply) => {
return toPublic(req.customer!);
});
/** PATCH /v1/me/config */
app.patch<{
Body: { system_prompt?: string | null; preferred_model?: string | null };
}>(
"/v1/me/config",
{
schema: {
body: {
type: "object",
properties: {
system_prompt: { type: ["string", "null"] },
preferred_model: { type: ["string", "null"] },
},
},
},
},
async (req, reply) => {
await updateCustomerConfig(req.customer!.id, req.body);
return reply.code(204).send();
}
);
/** GET /v1/me/export — GDPR data export */
app.get("/v1/me/export", async (req, reply) => {
const data = await exportCustomerData(req.customer!.id);
reply.header("Content-Disposition", `attachment; filename="agiliton-export.json"`);
return data;
});
/** POST /v1/me/delete — GDPR right to erasure */
app.post("/v1/me/delete", async (req, reply) => {
await softDeleteCustomer(req.customer!.id);
return reply.code(200).send({ deleted: true });
});
}

159
src/routes/sso.ts Normal file
View File

@@ -0,0 +1,159 @@
import type { FastifyInstance } from "fastify";
import { exchangeGoogleCode } from "../auth/google.js";
import { upsertCustomer, getVirtualKey } from "../auth/customers.js";
import { config } from "../config.js";
export async function ssoRoutes(app: FastifyInstance) {
/**
* POST /v1/auth/sso/google
* Body: { code: string, redirect_uri: string }
* Returns: { jwt, virtual_key, customer_id, is_new }
*/
app.post<{ Body: { code: string; redirect_uri: string } }>(
"/v1/auth/sso/google",
{
schema: {
body: {
type: "object",
required: ["code", "redirect_uri"],
properties: {
code: { type: "string" },
redirect_uri: { type: "string" },
},
},
},
},
async (req, reply) => {
let profile;
try {
profile = await exchangeGoogleCode(req.body.code, req.body.redirect_uri);
} catch (err) {
req.log.warn({ err }, "Google OAuth exchange failed");
return reply.code(400).send({ error: "oauth_exchange_failed" });
}
const { customer, virtualKey, isNew } = await upsertCustomer({
sso_provider: "google",
sso_sub: profile.sub,
email: profile.email,
name: profile.name,
});
const jwt = app.jwt.sign(
{ sub: customer.id, email: customer.email },
{ expiresIn: config.jwtExpiresIn }
);
return reply.code(isNew ? 201 : 200).send({
jwt,
virtual_key: virtualKey,
customer_id: customer.id,
is_new: isNew,
});
}
);
/**
* POST /v1/auth/sso/microsoft
* Body: { code: string, redirect_uri: string }
*
* Microsoft MSAL token exchange — requires MS_CLIENT_ID + MS_TENANT_ID env vars.
* Uses the MSAL token endpoint directly (no MSAL SDK dependency to keep bundle small).
*/
app.post<{ Body: { code: string; redirect_uri: string } }>(
"/v1/auth/sso/microsoft",
{
schema: {
body: {
type: "object",
required: ["code", "redirect_uri"],
properties: {
code: { type: "string" },
redirect_uri: { type: "string" },
},
},
},
},
async (req, reply) => {
if (!config.msClientId) {
return reply.code(501).send({ error: "microsoft_sso_not_configured" });
}
let profile: { sub: string; email: string; name: string };
try {
profile = await exchangeMicrosoftCode(req.body.code, req.body.redirect_uri);
} catch (err) {
req.log.warn({ err }, "Microsoft OAuth exchange failed");
return reply.code(400).send({ error: "oauth_exchange_failed" });
}
const { customer, virtualKey, isNew } = await upsertCustomer({
sso_provider: "microsoft",
sso_sub: profile.sub,
email: profile.email,
name: profile.name,
});
const jwt = app.jwt.sign(
{ sub: customer.id, email: customer.email },
{ expiresIn: config.jwtExpiresIn }
);
return reply.code(isNew ? 201 : 200).send({
jwt,
virtual_key: virtualKey,
customer_id: customer.id,
is_new: isNew,
});
}
);
}
async function exchangeMicrosoftCode(
code: string,
redirectUri: string
): Promise<{ sub: string; email: string; name: string }> {
const tenantId = config.msTenantId;
const params = new URLSearchParams({
client_id: config.msClientId,
code,
redirect_uri: redirectUri,
grant_type: "authorization_code",
scope: "openid email profile",
});
const tokenRes = await fetch(
`https://login.microsoftonline.com/${tenantId}/oauth2/v2.0/token`,
{
method: "POST",
headers: { "Content-Type": "application/x-www-form-urlencoded" },
body: params.toString(),
}
);
if (!tokenRes.ok) {
const text = await tokenRes.text();
throw new Error(`MS token endpoint failed ${tokenRes.status}: ${text}`);
}
const tokens = (await tokenRes.json()) as { id_token?: string; access_token: string };
// Decode the id_token (JWT) without verifying — we trust HTTPS to guarantee origin
if (tokens.id_token) {
const [, payload] = tokens.id_token.split(".");
const claims = JSON.parse(Buffer.from(payload, "base64url").toString());
return {
sub: claims.oid ?? claims.sub,
email: claims.email ?? claims.preferred_username,
name: claims.name ?? claims.email,
};
}
// Fallback: call /oidc/userinfo
const infoRes = await fetch("https://graph.microsoft.com/oidc/userinfo", {
headers: { Authorization: `Bearer ${tokens.access_token}` },
});
if (!infoRes.ok) throw new Error("MS userinfo fetch failed");
const info = (await infoRes.json()) as { sub: string; email: string; name: string };
return info;
}

26
src/utils/crypto.ts Normal file
View File

@@ -0,0 +1,26 @@
import { createCipheriv, createDecipheriv, randomBytes } from "crypto";
import { config } from "../config.js";
const ALGORITHM = "aes-256-gcm";
const KEY = Buffer.from(config.encryptionKey, "hex"); // must be 32 bytes
export function encrypt(plaintext: string): Buffer {
const iv = randomBytes(12);
const cipher = createCipheriv(ALGORITHM, KEY, iv);
const encrypted = Buffer.concat([
cipher.update(plaintext, "utf8"),
cipher.final(),
]);
const tag = cipher.getAuthTag();
// Layout: [12-byte IV][16-byte tag][ciphertext]
return Buffer.concat([iv, tag, encrypted]);
}
export function decrypt(buf: Buffer): string {
const iv = buf.subarray(0, 12);
const tag = buf.subarray(12, 28);
const encrypted = buf.subarray(28);
const decipher = createDecipheriv(ALGORITHM, KEY, iv);
decipher.setAuthTag(tag);
return decipher.update(encrypted) + decipher.final("utf8");
}

18
tsconfig.json Normal file
View File

@@ -0,0 +1,18 @@
{
"compilerOptions": {
"target": "ES2022",
"module": "NodeNext",
"moduleResolution": "NodeNext",
"outDir": "dist",
"rootDir": "src",
"strict": true,
"esModuleInterop": true,
"skipLibCheck": true,
"resolveJsonModule": true,
"declaration": true,
"declarationMap": true,
"sourceMap": true
},
"include": ["src/**/*"],
"exclude": ["node_modules", "dist"]
}

9
vitest.config.ts Normal file
View File

@@ -0,0 +1,9 @@
import { defineConfig } from "vitest/config";
export default defineConfig({
test: {
globals: true,
environment: "node",
include: ["src/__tests__/**/*.test.ts"],
},
});