@danceiny/gotry 0.2.0-rc.32 → 0.2.0-rc.33

package/dist/src/tenant-seam/llm-broker.js

@@ -0,0 +1,508 @@
+/**
+ * The LLM credential broker (docs/rfc/aaas-rfc.md §4.8, threat S-18).
+ *
+ * A planner child is the one place on a replica where model output, which is untrusted, steers a process that holds a
+ * bearer credential. Today that credential is the deployment-wide gateway key in the child's environment: valid from
+ * anywhere, for as long as nobody rotates it, and shared by every tenant. The broker moves the credential out of the
+ * child. The upstream key and base URL stay in this process; each planner process tree gets a *ticket* instead, a random
+ * bearer that is good only against this broker, on the loopback interface, for a short time, for a bounded number of
+ * requests, and that dies with the port it was issued for.
+ *
+ * What the broker is, and is not:
+ * - It is a credential-injecting forward proxy for one upstream. The child can ask it for exactly one thing, a POST to
+ *   an allow-listed path; the upstream host and path are fixed here, the child's `authorization` header is dropped and
+ *   the real key is attached, redirects are refused, and a response comes back with a closed set of headers.
+ * - It is not isolation from this process. Whoever can run code here can read the key. It narrows what a *compromised
+ *   child* gets: a ticket that expires, works from nowhere else, and names no key.
+ * - It meters. Every completed response is parsed for its `usage`, so a deployment can charge measured spend to the
+ *   tenant whose ticket asked, instead of the configured worst case (`TenantPlannerContext.reportUsage`).
+ *
+ * The broker carries the chat-completions route (`DEEPSEEK_PROVIDER=gateway-chat`, `dsh-llm-gateway-chat.js`). Another
+ * wire dialect needs its own path in `allowedPaths` and a live check before it is added.
+ */
+import { createHash, randomBytes } from 'node:crypto';
+import { createServer } from 'node:http';
+import { once } from 'node:events';
+const DEFAULT_TTL_MS = 15 * 60_000;
+const DEFAULT_MAX_REQUESTS = 64;
+const DEFAULT_MAX_CONCURRENT = 4;
+const DEFAULT_MAX_REQUEST_BYTES = 4 * 1024 * 1024;
+const DEFAULT_UPSTREAM_TIMEOUT_MS = 60_000;
+const DEFAULT_MAX_CONNECTIONS = 256;
+const DEFAULT_HEADERS_TIMEOUT_MS = 10_000;
+const DEFAULT_CONNECTIONS_CHECKING_INTERVAL_MS = 5_000;
+const MAX_USAGE_LINE_BYTES = 1024 * 1024;
+const TOKEN_PREFIX = 'gtk-';
+const LOOPBACK_ADDRESSES = new Set(['127.0.0.1', '::1', '::ffff:127.0.0.1']);
+function digest(token) {
+    return createHash('sha256').update(token).digest('hex');
+}
+function positive(name, value) {
+    if (!Number.isFinite(value) || value <= 0)
+        throw new RangeError(`${name} must be a positive number`);
+    return value;
+}
+function jsonError(res, status, code, headers = {}) {
+    if (res.headersSent) {
+        res.destroy();
+        return;
+    }
+    const body = JSON.stringify({ error: { code } });
+    res.writeHead(status, { 'content-type': 'application/json', 'content-length': Buffer.byteLength(body), 'cache-control': 'no-store', ...headers });
+    res.end(body);
+}
+/** The upstream URL for one allowed path, built the way the chat-completions adapter builds its own endpoint. */
+function upstreamUrlFor(base, path) {
+    const root = base.href.replace(/\/+$/u, '');
+    return base.pathname.replace(/\/+$/u, '').endsWith('/v1') ? `${root}${path}` : `${root}/v1${path}`;
+}
+/**
+ * What a child is told when the upstream answers with an error. The upstream's own body is never relayed: an upstream
+ * that echoes the credential it was sent (a misconfigured gateway does) would hand it to the very process the broker
+ * keeps it from. The status and a closed classification are all the adapter needs.
+ */
+function upstreamErrorBody(status) {
+    const type = status === 401 || status === 403 ? 'auth'
+        : status === 429 ? 'rate_limit'
+            : status === 408 || status >= 500 ? 'server'
+                : 'invalid_request';
+    return JSON.stringify({ error: { code: `upstream_${status}`, type } });
+}
+/** Resolves when the response drains or its connection closes, whichever is first, and leaves no listener behind. */
+function drainOrClose(res) {
+    return new Promise((resolve) => {
+        const done = () => {
+            res.off('drain', done);
+            res.off('close', done);
+            resolve();
+        };
+        res.once('drain', done);
+        res.once('close', done);
+    });
+}
+/** Reads `usage` out of a response as it streams past, without holding on to the response. Never throws. */
+class UsageTap {
+    sse;
+    carry = '';
+    carryOverflow = false;
+    found;
+    done = false;
+    decoder = new TextDecoder();
+    constructor(contentType) {
+        this.sse = contentType.includes('text/event-stream');
+    }
+    push(chunk) {
+        try {
+            this.carry += this.decoder.decode(chunk, { stream: true });
+            if (this.sse) {
+                for (let newline = this.carry.indexOf('\n'); newline >= 0; newline = this.carry.indexOf('\n')) {
+                    this.line(this.carry.slice(0, newline));
+                    this.carry = this.carry.slice(newline + 1);
+                }
+            }
+            if (this.carry.length > MAX_USAGE_LINE_BYTES) {
+                this.carry = '';
+                this.carryOverflow = true;
+            }
+        }
+        catch { /* metering must never touch the response */ }
+    }
+    /** True once the stream's `[DONE]` sentinel has gone past: nothing after it matters to a reader that stops there. */
+    get sawDone() {
+        return this.done;
+    }
+    finish() {
+        try {
+            this.carry += this.decoder.decode();
+            if (this.sse)
+                this.line(this.carry);
+            else if (!this.carryOverflow)
+                this.consider(JSON.parse(this.carry));
+        }
+        catch { /* an unparseable body is simply unmetered */ }
+        return this.found;
+    }
+    line(raw) {
+        const line = raw.trim();
+        if (!line.startsWith('data:'))
+            return;
+        const data = line.slice('data:'.length).trim();
+        if (data === '[DONE]') {
+            this.done = true;
+            return;
+        }
+        if (data === '')
+            return;
+        try {
+            this.consider(JSON.parse(data));
+        }
+        catch { /* a frame that is not JSON carries no usage */ }
+    }
+    consider(value) {
+        if (typeof value !== 'object' || value === null)
+            return;
+        const usage = value.usage;
+        if (typeof usage === 'object' && usage !== null)
+            this.found = usage;
+    }
+}
+function count(value) {
+    return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0 ? value : undefined;
+}
+/** The same mapping the chat-completions adapter applies: cache reads are subtracted from the prompt count. */
+function measuredUsage(raw) {
+    const prompt = count(raw.prompt_tokens);
+    const completion = count(raw.completion_tokens);
+    if (prompt === undefined || completion === undefined)
+        return undefined;
+    const details = typeof raw.prompt_tokens_details === 'object' && raw.prompt_tokens_details !== null ? raw.prompt_tokens_details : {};
+    const cacheRead = count(details.cached_tokens) ?? count(raw.prompt_cache_hit_tokens);
+    if (cacheRead !== undefined && cacheRead > prompt)
+        return undefined;
+    return { inputTokens: prompt - (cacheRead ?? 0), outputTokens: completion, ...(cacheRead === undefined ? {} : { cacheReadTokens: cacheRead }) };
+}
+export class LlmCredentialBroker {
+    upstream;
+    upstreamKey;
+    ttlMs;
+    maxRequests;
+    maxConcurrent;
+    maxRequestBytes;
+    upstreamTimeoutMs;
+    allowedPaths;
+    maxConnections;
+    headersTimeoutMs;
+    connectionsCheckingIntervalMs;
+    now;
+    transport;
+    tickets = new Map();
+    refused = new Map();
+    server;
+    baseUrl = '';
+    issued = 0;
+    forwarded = 0;
+    upstreamFailures = 0;
+    closed = false;
+    constructor(options) {
+        let upstream;
+        try {
+            upstream = new URL(options.upstreamBaseUrl);
+        }
+        catch {
+            throw new RangeError('upstreamBaseUrl must be an absolute URL');
+        }
+        if (upstream.protocol !== 'https:' && upstream.protocol !== 'http:')
+            throw new RangeError('upstreamBaseUrl must be http or https');
+        if (upstream.username !== '' || upstream.password !== '')
+            throw new RangeError('upstreamBaseUrl must not carry credentials');
+        if (typeof options.upstreamApiKey !== 'string' || options.upstreamApiKey.length === 0)
+            throw new RangeError('upstreamApiKey must be a non-empty string');
+        this.upstream = upstream;
+        this.upstreamKey = options.upstreamApiKey;
+        this.ttlMs = positive('ticketTtlMs', options.ticketTtlMs ?? DEFAULT_TTL_MS);
+        this.maxRequests = positive('maxRequestsPerTicket', options.maxRequestsPerTicket ?? DEFAULT_MAX_REQUESTS);
+        this.maxConcurrent = positive('maxConcurrentPerTicket', options.maxConcurrentPerTicket ?? DEFAULT_MAX_CONCURRENT);
+        this.maxRequestBytes = positive('maxRequestBytes', options.maxRequestBytes ?? DEFAULT_MAX_REQUEST_BYTES);
+        this.upstreamTimeoutMs = positive('upstreamTimeoutMs', options.upstreamTimeoutMs ?? DEFAULT_UPSTREAM_TIMEOUT_MS);
+        const paths = options.allowedPaths ?? ['/chat/completions'];
+        for (const path of paths) {
+            if (!/^\/[A-Za-z0-9._~/-]*$/u.test(path) || path.includes('..') || path.includes('//'))
+                throw new RangeError(`allowed path ${JSON.stringify(path)} is not a plain path`);
+        }
+        this.allowedPaths = new Set(paths);
+        this.maxConnections = positive('maxConnections', options.maxConnections ?? DEFAULT_MAX_CONNECTIONS);
+        this.headersTimeoutMs = positive('headersTimeoutMs', options.headersTimeoutMs ?? DEFAULT_HEADERS_TIMEOUT_MS);
+        this.connectionsCheckingIntervalMs = positive('connectionsCheckingIntervalMs', options.connectionsCheckingIntervalMs ?? DEFAULT_CONNECTIONS_CHECKING_INTERVAL_MS);
+        this.now = options.now ?? Date.now;
+        this.transport = options.fetch ?? fetch;
+    }
+    /** Starts listening on a loopback ephemeral port. Tickets can be issued once this resolves. */
+    async start() {
+        if (this.server || this.closed)
+            throw new Error('llm_broker_already_started');
+        // This port is handed to a process that model output steers, in the same process as everything else on the replica:
+        // it must not be a way to use up its descriptors. Connections are capped, and one that does not send a request is dropped
+        // (Node starts the headers clock when a connection is accepted, not at its first byte; the suite checks it on whatever Node runs it).
+        const server = createServer({ headersTimeout: this.headersTimeoutMs, requestTimeout: Math.max(60_000, this.headersTimeoutMs), connectionsCheckingInterval: this.connectionsCheckingIntervalMs }, (req, res) => { void this.handle(req, res); });
+        server.maxConnections = this.maxConnections;
+        // A planner run is long-lived, but nothing here should linger on an idle connection.
+        server.keepAliveTimeout = 5_000;
+        server.listen(0, '127.0.0.1');
+        await once(server, 'listening');
+        this.server = server;
+        this.baseUrl = `http://127.0.0.1:${server.address().port}/v1`;
+    }
+    issue(options = {}) {
+        if (!this.server || this.closed)
+            throw new Error('llm_broker_not_running');
+        // Expired tickets leave with the next issue, so a ticket nobody revoked cannot sit in the map for long.
+        this.sweep();
+        const ttlMs = positive('ttlMs', options.ttlMs ?? this.ttlMs);
+        const token = `${TOKEN_PREFIX}${randomBytes(32).toString('base64url')}`;
+        const key = digest(token);
+        const record = { onUsage: options.onUsage, inflight: new Set(), expiresAtMs: this.now() + ttlMs, requests: 0, inFlight: 0, revoked: false };
+        this.tickets.set(key, record);
+        this.issued += 1;
+        return {
+            baseUrl: this.baseUrl,
+            token,
+            expiresAtMs: record.expiresAtMs,
+            revoke: () => {
+                record.revoked = true;
+                this.tickets.delete(key);
+                // A revoked ticket also stops what it has under way: a stream already open would otherwise outlive it.
+                for (const controller of [...record.inflight])
+                    controller.abort('revoked');
+            },
+        };
+    }
+    stats() {
+        const now = this.now();
+        let active = 0;
+        for (const record of this.tickets.values())
+            if (!record.revoked && record.expiresAtMs > now)
+                active += 1;
+        return { issued: this.issued, active, tracked: this.tickets.size, forwarded: this.forwarded, refused: Object.fromEntries(this.refused), upstreamFailures: this.upstreamFailures };
+    }
+    async close() {
+        if (this.closed)
+            return;
+        this.closed = true;
+        this.tickets.clear();
+        const server = this.server;
+        if (!server)
+            return;
+        server.closeAllConnections();
+        await new Promise((resolve) => server.close(() => resolve()));
+    }
+    sweep() {
+        const now = this.now();
+        for (const [key, record] of this.tickets)
+            if (record.expiresAtMs <= now)
+                this.tickets.delete(key);
+    }
+    refuse(res, status, code, headers = {}) {
+        this.refused.set(code, (this.refused.get(code) ?? 0) + 1);
+        jsonError(res, status, code, headers);
+    }
+    /** A refusal whose request body is not going to be read: the answer goes out, and only then is the connection cut. */
+    refuseAndClose(req, res, status, code) {
+        res.once('finish', () => req.destroy());
+        this.refuse(res, status, code, { connection: 'close' });
+    }
+    ticketOf(req) {
+        const header = req.headers.authorization;
+        if (typeof header !== 'string' || !header.startsWith('Bearer '))
+            return undefined;
+        const token = header.slice('Bearer '.length).trim();
+        if (!token.startsWith(TOKEN_PREFIX) || token.length > 128)
+            return undefined;
+        const record = this.tickets.get(digest(token));
+        if (!record || record.revoked || record.expiresAtMs <= this.now())
+            return undefined;
+        return record;
+    }
+    async handle(req, res) {
+        const ticket = {};
+        try {
+            // Defense in depth: the socket is bound to loopback, and a remote peer is refused even so.
+            if (!LOOPBACK_ADDRESSES.has(req.socket.remoteAddress ?? '')) {
+                this.refuseAndClose(req, res, 403, 'not_loopback');
+                return;
+            }
+            const url = new URL(req.url ?? '/', 'http://broker.invalid');
+            const prefix = '/v1';
+            const path = url.pathname.startsWith(`${prefix}/`) ? url.pathname.slice(prefix.length) : '';
+            if (!this.allowedPaths.has(path)) {
+                this.refuse(res, 404, 'path_not_allowed');
+                req.resume();
+                return;
+            }
+            if (req.method !== 'POST') {
+                this.refuse(res, 405, 'method_not_allowed', { allow: 'POST' });
+                req.resume();
+                return;
+            }
+            const record = this.ticketOf(req);
+            if (!record) {
+                this.refuse(res, 401, 'ticket_invalid');
+                req.resume();
+                return;
+            }
+            if (record.requests >= this.maxRequests) {
+                this.refuse(res, 429, 'ticket_exhausted');
+                req.resume();
+                return;
+            }
+            if (record.inFlight >= this.maxConcurrent) {
+                this.refuse(res, 429, 'ticket_busy', { 'retry-after': '1' });
+                req.resume();
+                return;
+            }
+            record.requests += 1;
+            record.inFlight += 1;
+            ticket.record = record;
+            const declared = Number(req.headers['content-length']);
+            if (Number.isFinite(declared) && declared > this.maxRequestBytes) {
+                this.refuseAndClose(req, res, 413, 'request_too_large');
+                return;
+            }
+            const body = await this.readBody(req);
+            if (body === null) {
+                this.refuseAndClose(req, res, 413, 'request_too_large');
+                return;
+            }
+            await this.forward(req, res, record, path, body);
+        }
+        catch {
+            if (!res.writableEnded)
+                jsonError(res, 500, 'broker_error');
+        }
+        finally {
+            if (ticket.record)
+                ticket.record.inFlight -= 1;
+        }
+    }
+    async readBody(req) {
+        const chunks = [];
+        let size = 0;
+        for await (const chunk of req) {
+            size += chunk.length;
+            if (size > this.maxRequestBytes)
+                return null;
+            chunks.push(chunk);
+        }
+        return Buffer.concat(chunks);
+    }
+    async forward(req, res, record, path, body) {
+        const abort = new AbortController();
+        record.inflight.add(abort);
+        try {
+            await this.relay(req, res, record, path, body, abort);
+        }
+        finally {
+            record.inflight.delete(abort);
+        }
+    }
+    async relay(req, res, record, path, body, abort) {
+        res.on('close', () => { if (!res.writableFinished)
+            abort.abort('client'); });
+        const headers = {
+            authorization: `Bearer ${this.upstreamKey}`,
+            'content-type': typeof req.headers['content-type'] === 'string' ? req.headers['content-type'] : 'application/json',
+        };
+        if (typeof req.headers.accept === 'string')
+            headers.accept = req.headers.accept;
+        // The only requests the broker carries are JSON; the text form is what goes upstream.
+        const text = body.toString('utf8');
+        let model = '';
+        try {
+            const parsed = JSON.parse(text);
+            if (typeof parsed.model === 'string')
+                model = parsed.model;
+        }
+        catch { /* an unparseable body is the upstream's to refuse; it is simply unmetered here */ }
+        const timer = setTimeout(() => abort.abort('timeout'), this.upstreamTimeoutMs);
+        let response;
+        try {
+            response = await this.transport(upstreamUrlFor(this.upstream, path), { method: 'POST', headers, body: text, redirect: 'manual', signal: abort.signal });
+        }
+        catch {
+            this.upstreamFailures += 1;
+            if (res.writableEnded || res.destroyed)
+                return;
+            const reason = abort.signal.reason;
+            if (reason === 'revoked')
+                jsonError(res, 401, 'ticket_invalid');
+            else
+                jsonError(res, abort.signal.aborted ? 504 : 502, abort.signal.aborted ? 'upstream_timeout' : 'upstream_unreachable');
+            return;
+        }
+        finally {
+            clearTimeout(timer);
+        }
+        // A redirect would send the injected key wherever the upstream points; it is a failure, not something to follow.
+        if (response.status >= 300 && response.status < 400) {
+            this.upstreamFailures += 1;
+            await response.body?.cancel().catch(() => undefined);
+            jsonError(res, 502, 'upstream_redirect_refused');
+            return;
+        }
+        const contentType = response.headers.get('content-type') ?? 'application/json';
+        const outHeaders = { 'content-type': contentType, 'cache-control': 'no-store' };
+        const retryAfter = response.headers.get('retry-after');
+        if (retryAfter !== null && /^\d{1,6}$/u.test(retryAfter))
+            outHeaders['retry-after'] = retryAfter;
+        if (response.status < 200 || response.status >= 300) {
+            // Never the upstream's own words (see upstreamErrorBody); its body is not read.
+            await response.body?.cancel().catch(() => undefined);
+            const mapped = upstreamErrorBody(response.status);
+            const errorHeaders = { 'content-type': 'application/json', 'content-length': String(Buffer.byteLength(mapped)), 'cache-control': 'no-store' };
+            if (outHeaders['retry-after'] !== undefined)
+                errorHeaders['retry-after'] = outHeaders['retry-after'];
+            res.writeHead(response.status, errorHeaders);
+            res.end(mapped);
+            this.forwarded += 1;
+            return;
+        }
+        res.writeHead(response.status, outHeaders);
+        const tap = new UsageTap(contentType);
+        // Measured spend is reported once, and before the child can see that the stream is over: at the `[DONE]` frame, which
+        // is where the chat-completions adapter stops reading (it then cancels the request, whatever the upstream still does),
+        // or, for a body with no such frame, when the body has ended.
+        let metered = false;
+        const meter = () => {
+            if (metered)
+                return;
+            metered = true;
+            const usage = tap.finish();
+            const measured = usage ? measuredUsage(usage) : undefined;
+            if (measured && model !== '' && record.onUsage) {
+                try {
+                    record.onUsage({ model, usage: measured });
+                }
+                catch { /* the caller's meter never affects the response */ }
+            }
+        };
+        try {
+            if (response.body) {
+                const reader = response.body.getReader();
+                for (;;) {
+                    const { done, value } = await reader.read();
+                    if (done)
+                        break;
+                    tap.push(value);
+                    if (tap.sawDone)
+                        meter();
+                    if (!res.write(value)) {
+                        // A child that went away while the pipe was full never drains it: stop waiting when the connection closes.
+                        await drainOrClose(res);
+                        if (res.destroyed) {
+                            await reader.cancel().catch(() => undefined);
+                            break;
+                        }
+                    }
+                }
+            }
+            meter();
+            res.end();
+            this.forwarded += 1;
+        }
+        catch {
+            // A stream that had already delivered `[DONE]` was complete; the child may simply have stopped reading and cancelled.
+            // Anything else is the upstream failing mid-stream, and the child must see a broken connection, never a truncated
+            // stream that looks whole.
+            if (tap.sawDone)
+                this.forwarded += 1;
+            else
+                this.upstreamFailures += 1;
+            res.destroy();
+        }
+    }
+}
+/** A broker that is listening. */
+export async function startLlmCredentialBroker(options) {
+    const broker = new LlmCredentialBroker(options);
+    await broker.start();
+    return broker;
+}