feat(ai): multi-key BYOK, real server class, real entitlement enforcement

Three pieces built together tonight since they're naturally linked (the
server-class proxy is the real entitlement enforcement chokepoint):

1. Multi-key BYOK (public class): several named provider profiles
   (name/baseUrl/model), each with its own key in lib/ai/key-store.ts
   (keyed by profile id, not a single fixed 'public' slot). The "Try it"
   pane lets you pick which saved profile answers each question - not one
   fixed default.

2. `server` class, real: app/api/ai/server/{models,chat} proxy through
   this app's own backend to AI_SERVER_BASE_URL - same-origin from the
   browser, no CORS/OLLAMA_ORIGINS story at all, standing in tonight for
   VNC's EU/CH-hosted infra with the real Ollama on this Mac (swapping to
   the real instance tomorrow is a config change).

3. Real entitlement enforcement (lib/ai/entitlement.ts), scoped to `server`
   only (not local/public, per the 2026-08-05 decisions): checkAndAssignSeat()
   re-validates on every /api/ai/server/chat call - first use auto-assigns a
   seat if any remain, further calls from an unlicensed user get a 402 with
   a specific reason. recordUsage() appends to an append-only metering
   ledger (timestamp/user/model/tokens/latency) that IS the billing record.
   Admin data endpoints at /api/admin/ai/entitlement (seat total, revoke) -
   the visual admin console is a separate, not-yet-built task.

Two real bugs found and fixed during verification, not just claimed fixed:
- /api/ai/policy never actually added 'server' to entitlement.classes even
  when AI_SERVER_BASE_URL was set (only the type comment was updated) - the
  Server radio option silently never appeared until this was caught live.
- The new routes used readStalwartAuthContext(0) (hardcoded slot, SSO/reauth-
  specific) instead of getStalwartCredentials() (the general multi-slot
  session resolver every other authenticated route uses) - reachable but
  wrong, and would have hidden a real auth gap behind "works on my slot".

Verified end-to-end for real: built + ran the actual server, logged in via
the real (non-demo) auth flow, selected Server, listed the real Ollama
models through the proxy, asked "Reply with exactly the words: SERVER CLASS
WORKS" and got back exactly that - plus confirmed on disk (not just in the
UI) that data/admin-state/ai-entitlement.json recorded the seat assignment
and ai-metering.jsonl recorded real prompt/completion token counts and
latency from the actual model call. Rejection-path logic (seat limit
reached, zero seats configured, revocation) covered by 5 new unit tests
rather than a second live round trip. Full suite: typecheck clean, lint
clean, translations 48/48, production build succeeds.
This commit is contained in:
Bernd Rodler
2026-08-06 00:07:56 +02:00
parent bde8455df5
commit dda7adf565
11 changed files with 821 additions and 122 deletions
+92
View File
@@ -0,0 +1,92 @@
import { describe, expect, it, beforeEach, afterEach, vi } from 'vitest';
import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';
// Real end-to-end seat assignment against the real Ollama was verified live
// (see the commit this test ships with); this covers the rejection branch,
// which is deterministic and cheaper to prove with a unit test than another
// live round trip.
describe('lib/ai/entitlement', () => {
let stateDir: string;
beforeEach(async () => {
vi.resetModules();
stateDir = await mkdtemp(path.join(tmpdir(), 'ai-entitlement-test-'));
process.env.ADMIN_STATE_DIR = stateDir;
delete process.env.AI_SERVER_SEAT_TOTAL;
// Each test needs a fresh globalThis singleton, not just a fresh module -
// the module stashes cached state on globalThis specifically to survive
// HMR, so resetModules() alone doesn't clear it.
delete (globalThis as Record<symbol, unknown>)[Symbol.for('vncmail.ai.entitlement')];
});
afterEach(async () => {
delete process.env.ADMIN_STATE_DIR;
await rm(stateDir, { recursive: true, force: true });
});
it('assigns a seat on first use and allows the same user again', async () => {
const { checkAndAssignSeat, setSeatTotal } = await import('../entitlement');
await setSeatTotal(1);
const first = await checkAndAssignSeat('alice@example.com');
expect(first).toEqual({ allowed: true, seatJustAssigned: true });
const second = await checkAndAssignSeat('alice@example.com');
expect(second).toEqual({ allowed: true });
});
it('rejects a new user once all seats are assigned', async () => {
const { checkAndAssignSeat, setSeatTotal } = await import('../entitlement');
await setSeatTotal(1);
await checkAndAssignSeat('alice@example.com');
const rejected = await checkAndAssignSeat('bob@example.com');
expect(rejected.allowed).toBe(false);
expect(rejected.reason).toMatch(/already assigned/i);
});
it('rejects everyone when no seats are configured', async () => {
const { checkAndAssignSeat } = await import('../entitlement');
const result = await checkAndAssignSeat('anyone@example.com');
expect(result.allowed).toBe(false);
expect(result.reason).toMatch(/no licensed seats/i);
});
it('revoking a seat frees it for someone else', async () => {
const { checkAndAssignSeat, setSeatTotal, revokeSeat } = await import('../entitlement');
await setSeatTotal(1);
await checkAndAssignSeat('alice@example.com');
await revokeSeat('alice@example.com');
const result = await checkAndAssignSeat('bob@example.com');
expect(result).toEqual({ allowed: true, seatJustAssigned: true });
});
it('persists usage to the metering ledger, append-only', async () => {
const { recordUsage, readMeteringLedger } = await import('../entitlement');
await recordUsage({
timestamp: new Date(0).toISOString(),
username: 'alice@example.com',
model: 'qwen2.5:32b',
promptTokens: 10,
completionTokens: 5,
latencyMs: 123,
});
await recordUsage({
timestamp: new Date(0).toISOString(),
username: 'alice@example.com',
model: 'qwen2.5:32b',
promptTokens: 8,
completionTokens: 3,
latencyMs: 90,
});
const ledger = await readMeteringLedger();
expect(ledger).toHaveLength(2);
expect(ledger[0].promptTokens).toBe(10);
expect(ledger[1].promptTokens).toBe(8);
});
});
+162
View File
@@ -0,0 +1,162 @@
// Real entitlement + metering enforcement for the AI Assistant's `server`
// class (docs/AI-ASSISTANT-CONCEPT.md §9/§10 — per-seat licensing, a
// metering ledger that doubles as the billing record).
//
// Deliberately scoped to `server` only, not `local`/`public`, per the
// 2026-08-05 decisions: `local` ships free (never reaches a server this app
// controls, so it can't be metered — see the doc's own §9 reasoning) and
// `public` is explicitly unmonitored for now. `server` is the one class that
// (a) proxies through this app's own backend (see app/api/ai/server/*) and
// (b) has a real marginal cost (shared GPU time) worth gating — so it is the
// one place enforcement is both possible and worth building tonight.
//
// Persistence follows the existing admin state-dir convention
// (lib/admin/paths.ts): STATE, not CONFIG, because this is runtime-mutated
// data (seat assignments, usage), not operator-authored config.
import { readFile, writeFile, rename, appendFile } from 'node:fs/promises';
import { existsSync } from 'node:fs';
import { getStatePath, ensureStateDir } from '@/lib/admin/paths';
import { logger } from '@/lib/logger';
export interface AiEntitlementState {
subject: 'tenant' | 'user';
tier: 'base' | 'standard' | 'pro';
/** Total seats licensed. 0 = server class entirely unlicensed (default). */
seatsTotal: number;
/** Usernames who have consumed a seat (first successful use assigns one,
* matching real per-seat licensing — not deallocated by idling). */
assignedTo: string[];
}
export interface EntitlementCheck {
allowed: boolean;
reason?: string;
/** True the moment this call consumed a previously-unassigned seat. */
seatJustAssigned?: boolean;
}
export interface MeteringEntry {
timestamp: string;
username: string;
model: string;
/** Ollama reports these as prompt_eval_count / eval_count. */
promptTokens: number;
completionTokens: number;
latencyMs: number;
}
const STATE_FILE = 'ai-entitlement.json';
const LEDGER_FILE = 'ai-metering.jsonl';
const DEFAULT_STATE: AiEntitlementState = {
subject: 'tenant',
tier: 'base',
seatsTotal: Number.parseInt(process.env.AI_SERVER_SEAT_TOTAL ?? '0', 10) || 0,
assignedTo: [],
};
// Stash on globalThis like config-manager.ts — HMR/dev re-evaluates this
// module, and in-memory seat state must survive that or every hot reload
// would silently re-grant seats.
const SINGLETON_KEY = Symbol.for('vncmail.ai.entitlement');
type GlobalWithState = typeof globalThis & { [SINGLETON_KEY]?: Promise<AiEntitlementState> | undefined };
async function readState(): Promise<AiEntitlementState> {
try {
const raw = await readFile(getStatePath(STATE_FILE), 'utf-8');
return { ...DEFAULT_STATE, ...JSON.parse(raw) };
} catch (error) {
if ((error as NodeJS.ErrnoException).code !== 'ENOENT') {
logger.warn('ai-entitlement: failed to read state, using defaults', {
error: error instanceof Error ? error.message : String(error),
});
}
return { ...DEFAULT_STATE };
}
}
async function writeState(state: AiEntitlementState): Promise<void> {
await ensureStateDir();
const target = getStatePath(STATE_FILE);
const tmp = target + '.tmp';
await writeFile(tmp, JSON.stringify(state, null, 2), 'utf-8');
await rename(tmp, target);
}
let cached: AiEntitlementState | null = null;
async function loadCached(): Promise<AiEntitlementState> {
if (cached) return cached;
const g = globalThis as GlobalWithState;
if (!g[SINGLETON_KEY]) g[SINGLETON_KEY] = readState();
cached = await g[SINGLETON_KEY];
return cached;
}
/**
* The real enforcement point (doc §10 point 2): re-validated on every call,
* never trusts anything the client sent. Auto-assigns a seat on first use
* when seats remain — that's what "per-seat" means for a subject that
* hasn't been explicitly provisioned by an admin yet.
*/
export async function checkAndAssignSeat(username: string): Promise<EntitlementCheck> {
const state = await loadCached();
if (state.assignedTo.includes(username)) {
return { allowed: true };
}
if (state.assignedTo.length >= state.seatsTotal) {
return {
allowed: false,
reason: state.seatsTotal === 0
? 'The server-hosted AI class has no licensed seats configured.'
: `All ${state.seatsTotal} licensed seat(s) are already assigned to other users.`,
};
}
const next: AiEntitlementState = { ...state, assignedTo: [...state.assignedTo, username] };
await writeState(next);
cached = next;
const g = globalThis as GlobalWithState;
g[SINGLETON_KEY] = Promise.resolve(next);
return { allowed: true, seatJustAssigned: true };
}
/** The metering write IS the billing record — see module header. Append-only,
* never rewritten, so it stays valid as an audit trail even if this process
* crashes mid-write (worst case: one truncated trailing line). */
export async function recordUsage(entry: MeteringEntry): Promise<void> {
await ensureStateDir();
await appendFile(getStatePath(LEDGER_FILE), JSON.stringify(entry) + '\n', 'utf-8');
}
export async function getEntitlementState(): Promise<AiEntitlementState> {
return loadCached();
}
export async function setSeatTotal(total: number): Promise<AiEntitlementState> {
const state = await loadCached();
const next: AiEntitlementState = { ...state, seatsTotal: Math.max(0, Math.trunc(total)) };
await writeState(next);
cached = next;
(globalThis as GlobalWithState)[SINGLETON_KEY] = Promise.resolve(next);
return next;
}
export async function revokeSeat(username: string): Promise<AiEntitlementState> {
const state = await loadCached();
const next: AiEntitlementState = { ...state, assignedTo: state.assignedTo.filter((u) => u !== username) };
await writeState(next);
cached = next;
(globalThis as GlobalWithState)[SINGLETON_KEY] = Promise.resolve(next);
return next;
}
/** Read-only summary, no PII beyond usernames already visible to any admin. */
export async function readMeteringLedger(limit = 200): Promise<MeteringEntry[]> {
const path = getStatePath(LEDGER_FILE);
if (!existsSync(path)) return [];
const raw = await readFile(path, 'utf-8');
const lines = raw.trim().split('\n').filter(Boolean);
return lines.slice(-limit).map((line) => JSON.parse(line) as MeteringEntry);
}
+16 -10
View File
@@ -1,26 +1,32 @@
// Client-held storage for the user's own public-provider API key (BYOK).
// Client-held storage for the user's own public-provider API keys (BYOK).
//
// Decision 2026-08-05 (reverses docs/AI-ASSISTANT-CONCEPT.md decision #1's
// server-side-custody design): the user brings and holds their own key,
// server-side-custody design): the user brings and holds their own keys,
// client-side, not VNC. This is the same custody model as
// vncmail-native's lib/ai-key-store.ts (expo-secure-store there; this repo
// has no OS keychain access from a browser tab, so localStorage is the
// honest equivalent here — plain, not hidden behind a false sense of
// "secure storage"). A fuller Paperclip-style key-management UI (multiple
// providers, masking, rotation) is good follow-up work, not built tonight.
// "secure storage").
//
// Decision 2026-08-05 (later same night): several keys, not one — a user may
// hold multiple named provider profiles (different models, different
// providers) and pick which one answers a given question. Keys are stored
// separately from `lib/ai/local-settings.ts`'s profile metadata (name, base
// URL, model) so a profile can be exported/shared without its secret, and so
// clearing one key can't accidentally corrupt the profile list.
const KEY_PREFIX = 'vncmail:ai:key:';
export function getAiApiKey(provider: 'public'): string | null {
export function getAiApiKey(profileId: string): string | null {
if (typeof window === 'undefined') return null;
return window.localStorage.getItem(KEY_PREFIX + provider);
return window.localStorage.getItem(KEY_PREFIX + profileId);
}
export function setAiApiKey(provider: 'public', key: string): void {
export function setAiApiKey(profileId: string, key: string): void {
if (typeof window === 'undefined') return;
window.localStorage.setItem(KEY_PREFIX + provider, key);
window.localStorage.setItem(KEY_PREFIX + profileId, key);
}
export function clearAiApiKey(provider: 'public'): void {
export function clearAiApiKey(profileId: string): void {
if (typeof window === 'undefined') return;
window.localStorage.removeItem(KEY_PREFIX + provider);
window.localStorage.removeItem(KEY_PREFIX + profileId);
}
+87 -18
View File
@@ -1,11 +1,18 @@
// The AI assistant's wire client mirrors vncmail-native's src/api/ai.ts
// (same prototype scope: local Ollama + BYOK public, no VNC-hosted `server`
// class, no streaming) so the two clients stay in lockstep. Runs entirely
// client-side (`'use client'` callers only) — a direct loopback/provider
// fetch, matching docs/AI-ASSISTANT-CONCEPT.md §2's "local"/"public" rows,
// not proxied through this app's own Next.js server. That distinction
// matters once this app is hosted remotely: a server-side proxy would reach
// the *server's* loopback, not the user's own laptop running Ollama.
// The AI assistant's wire client. `local`/`public` mirror vncmail-native's
// src/api/ai.ts (direct loopback/provider fetch, no streaming) so the two
// clients stay in lockstep — matching docs/AI-ASSISTANT-CONCEPT.md §2's
// "local"/"public" rows, not proxied through this app's own Next.js server.
// That distinction matters once this app is hosted remotely: a server-side
// proxy would reach the *server's* loopback, not the user's own laptop
// running Ollama.
//
// `server` (added 2026-08-05 night) is the opposite by design: it DOES
// proxy through this app's own backend (app/api/ai/server/*), because it's
// centrally-hosted infra (VNC's EU/CH stack — standing in tonight for a real
// Ollama on this Mac, see lib/ai/entitlement.ts), not a user's own machine.
// That server-side hop is also the one real entitlement enforcement point
// (§10 point 2) — `local`/`public` never reach it, by design, and so cannot
// be metered or billed the same way.
export interface ChatMessage {
role: 'system' | 'user' | 'assistant';
@@ -74,6 +81,42 @@ export async function chatLocal(
return content;
}
// ── Server: centrally-hosted, proxied through this app's own backend
// (app/api/ai/server/*). Unlike `local`, this is same-origin from the
// browser's perspective — no CORS/OLLAMA_ORIGINS story at all — and unlike
// both `local` and `public`, every call is entitlement-checked server-side. ──
export async function listServerModels(): Promise<string[]> {
const res = await fetch('/api/ai/server/models');
if (!res.ok) {
const body = (await res.json().catch(() => null)) as { error?: string } | null;
throw new Error(body?.error ?? `AI server returned ${res.status}`);
}
const body = (await res.json()) as { models?: string[] };
return body.models ?? [];
}
export interface ServerChatResult {
answer: string;
/** True the moment this call consumed a previously-unassigned licensed seat
* (lib/ai/entitlement.ts) — surfaced so the UI can say so once, not left
* to happen silently the first time someone uses this class. */
seatJustAssigned: boolean;
}
export async function chatServer(model: string, messages: ChatMessage[]): Promise<ServerChatResult> {
const res = await fetch('/api/ai/server/chat', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ model, messages }),
});
const body = (await res.json().catch(() => null)) as { answer?: string; error?: string; seatJustAssigned?: boolean } | null;
if (!res.ok || !body?.answer) {
throw new Error(body?.error ?? `AI server returned ${res.status}`);
}
return { answer: body.answer, seatJustAssigned: body.seatJustAssigned === true };
}
// ── Public: OpenAI-compatible chat-completions. OpenRouter by default, but any
// endpoint speaking this shape works unmodified (self-hosted vLLM, LiteLLM, etc). ──
@@ -118,6 +161,10 @@ export interface AskResult {
sources: AskSource[];
/** True when the question was answered without any retrieved context. */
unaugmented: boolean;
/** True the moment this call consumed a previously-unassigned licensed
* seat on the `server` class (lib/ai/entitlement.ts). Always false for
* `local`/`public`, which aren't entitlement-gated. */
seatJustAssigned: boolean;
}
interface OfflineSearchHit {
@@ -152,21 +199,34 @@ export function buildPrompt(question: string, contextBlock: string): ChatMessage
];
}
/**
* One saved BYOK profile, resolved to an actual key — the caller picks which
* profile answers *this* question (docs decision 2026-08-05: several keys,
* selected case by case, not one fixed "the" public provider).
*/
export interface ResolvedPublicProfile {
baseUrl: string;
model: string;
apiKey: string;
}
export interface AskConfig {
provider: 'local' | 'public';
provider: 'local' | 'server' | 'public';
localBaseUrl: string;
localModel: string | null;
publicBaseUrl: string;
publicModel: string;
publicApiKey: string | null;
serverModel: string | null;
publicProfile: ResolvedPublicProfile | null;
}
export async function askMail(question: string, config: AskConfig): Promise<AskResult> {
if (config.provider === 'local' && !config.localModel) {
throw new Error('No local model selected');
}
if (config.provider === 'public' && !config.publicApiKey) {
throw new Error('No public API key saved');
if (config.provider === 'server' && !config.serverModel) {
throw new Error('No server model selected');
}
if (config.provider === 'public' && !config.publicProfile) {
throw new Error('No provider profile selected');
}
const retrieved = await retrieveContext(question);
@@ -174,14 +234,23 @@ export async function askMail(question: string, config: AskConfig): Promise<AskR
? buildPrompt(question, retrieved.contextBlock)
: [{ role: 'user' as const, content: question }];
const answer =
config.provider === 'public'
? await chatPublic(config.publicBaseUrl, config.publicApiKey as string, config.publicModel, messages)
: await chatLocal(config.localBaseUrl, config.localModel as string, messages);
let answer: string;
let seatJustAssigned = false;
if (config.provider === 'public') {
const profile = config.publicProfile as ResolvedPublicProfile;
answer = await chatPublic(profile.baseUrl, profile.apiKey, profile.model, messages);
} else if (config.provider === 'server') {
const result = await chatServer(config.serverModel as string, messages);
answer = result.answer;
seatJustAssigned = result.seatJustAssigned;
} else {
answer = await chatLocal(config.localBaseUrl, config.localModel as string, messages);
}
return {
answer,
sources: (retrieved?.hits ?? []).map((h) => ({ id: h.id, subject: h.title })),
unaugmented: !retrieved,
seatJustAssigned,
};
}
+48 -6
View File
@@ -4,14 +4,31 @@
// (docs/AI-ASSISTANT-CONCEPT.md §12's P0/P5), and migrating it into the
// shared store belongs with whichever phase makes these settings real
// product config rather than a local-AI test harness.
export type AiProvider = 'local' | 'public';
export type AiProvider = 'local' | 'server' | 'public';
/**
* A named public-provider configuration (BYOK). Decision 2026-08-05: several
* of these, not one — different models/providers for different questions,
* picked case by case at Ask time (see `activeProfileId`). The API key
* itself lives in `lib/ai/key-store.ts`, keyed by `id`, not here — so a
* profile's metadata can be listed/edited without ever handling the secret.
*/
export interface AiProviderProfile {
id: string;
name: string;
baseUrl: string;
model: string;
}
export interface AiLocalSettings {
provider: AiProvider | null;
localBaseUrl: string;
localModel: string | null;
publicBaseUrl: string;
publicModel: string;
serverModel: string | null;
publicProfiles: AiProviderProfile[];
/** Which saved profile answers the next question. Not a permanent default —
* the "Try it" UI lets this be changed per question. */
activeProfileId: string | null;
publicConsentAccepted: boolean;
}
@@ -21,17 +38,38 @@ export const DEFAULT_AI_SETTINGS: AiLocalSettings = {
provider: null,
localBaseUrl: 'http://127.0.0.1:11434',
localModel: null,
publicBaseUrl: 'https://openrouter.ai/api/v1',
publicModel: '',
serverModel: null,
publicProfiles: [],
activeProfileId: null,
publicConsentAccepted: false,
};
function newProfileId(): string {
return `profile-${Math.random().toString(36).slice(2, 10)}-${Math.random().toString(36).slice(2, 10)}`;
}
/** One-time upgrade from the earlier single-profile shape (a bare
* publicBaseUrl/publicModel pair) into the profile list, so a browser that
* already saved settings before profiles existed doesn't just lose them. */
function migrate(raw: Record<string, unknown>): Partial<AiLocalSettings> {
if (Array.isArray(raw.publicProfiles)) return raw as Partial<AiLocalSettings>;
if (typeof raw.publicBaseUrl === 'string' && typeof raw.publicModel === 'string' && raw.publicModel) {
const id = newProfileId();
return {
...raw,
publicProfiles: [{ id, name: 'Default', baseUrl: raw.publicBaseUrl, model: raw.publicModel }],
activeProfileId: id,
};
}
return raw as Partial<AiLocalSettings>;
}
export function loadAiSettings(): AiLocalSettings {
if (typeof window === 'undefined') return { ...DEFAULT_AI_SETTINGS };
try {
const raw = window.localStorage.getItem(STORAGE_KEY);
if (!raw) return { ...DEFAULT_AI_SETTINGS };
return { ...DEFAULT_AI_SETTINGS, ...JSON.parse(raw) };
return { ...DEFAULT_AI_SETTINGS, ...migrate(JSON.parse(raw)) };
} catch {
return { ...DEFAULT_AI_SETTINGS };
}
@@ -41,3 +79,7 @@ export function saveAiSettings(settings: AiLocalSettings): void {
if (typeof window === 'undefined') return;
window.localStorage.setItem(STORAGE_KEY, JSON.stringify(settings));
}
export function createProfile(name: string, baseUrl: string, model: string): AiProviderProfile {
return { id: newProfileId(), name, baseUrl, model };
}
+7 -3
View File
@@ -12,9 +12,13 @@
// not built yet). The client-side "this leaves the organisation"
// acknowledgement still shows (cheap, honest), it just isn't
// server-enforced yet.
// - `server` (VNC-hosted, EU/CH) isn't wired up client-side yet — infra is
// "this MacBook tonight, the dev k8s cluster tomorrow" per that
// decision, sequenced after `local` rather than before it.
// - `server` (VNC-hosted, EU/CH) is now wired up for real too (added later
// the same night, per "do it this night - no stop"): a real server-side
// proxy (app/api/ai/server/*) to AI_SERVER_BASE_URL, which stands in for
// the dev-k8s-hosted instance until that exists tomorrow. Unlike
// `local`/`public`, `server` IS entitlement-enforced for real —
// lib/ai/entitlement.ts — since it's the one class with a real,
// centrally-borne cost.
export type AiClass = 'local' | 'server' | 'public';