ag-ai 0.1.0

package/runtime/server.mjs

@@ -0,0 +1,621 @@
+import { timingSafeEqual, randomUUID } from 'node:crypto';
+import { createServer } from 'node:http';
+import net from 'node:net';
+/* eslint-disable import/extensions -- Native Node ESM requires file extensions. */
+
+import {
+  GenerationProtocolError,
+  normalizeGenerationRequest,
+  normalizeIdempotencyKey,
+  publicModelCatalog,
+  requestDigest,
+  idempotencyDigest,
+} from './protocol.mjs';
+import { GenerationStoreError } from './store.mjs';
+
+const DEFAULT_GENERATION_PORT = 4501;
+const MAX_REQUEST_BYTES = 64 * 1024 * 1024;
+const MAX_RESPONSE_BYTES = 64 * 1024 * 1024;
+const MAX_MODEL_RESPONSE_BYTES = 512 * 1024;
+const JOB_LEASE_MS = 31 * 60 * 1000;
+const LOCAL_RESPONSE_PREFIX = 'generation/';
+const MAX_CONCURRENT_BODY_READS = 2;
+
+const noopAnalytics = { emit: () => Promise.resolve(true) };
+
+const asObject = (value) => (value && typeof value === 'object' ? value : null);
+
+const httpError = (status, code, message) => new GenerationProtocolError(status, code, message);
+
+const errorPayload = (error) => {
+  if (error instanceof GenerationProtocolError) {
+    return { status: error.status, body: { error: { code: error.code, message: error.message } } };
+  }
+  if (error instanceof GenerationStoreError) {
+    const status =
+      error.code === 'not_found'
+        ? 404
+        : error.code === 'result_expired'
+          ? 410
+          : error.code === 'idempotency_conflict'
+            ? 409
+            : error.code === 'queue_full'
+              ? 429
+              : 500;
+    return { status, body: { error: { code: error.code, message: error.message } } };
+  }
+  return {
+    status: 500,
+    body: { error: { code: 'internal_error', message: 'generation request failed' } },
+  };
+};
+
+const isPrivateIPv4 = (address) => {
+  if (net.isIP(address) !== 4) return false;
+  const octets = address.split('.').map(Number);
+  return (
+    octets[0] === 127 ||
+    octets[0] === 10 ||
+    (octets[0] === 172 && octets[1] >= 16 && octets[1] <= 31) ||
+    (octets[0] === 192 && octets[1] === 168)
+  );
+};
+
+const normalizeIPv4 = (address) => {
+  if (typeof address !== 'string') return undefined;
+  if (address.startsWith('::ffff:')) return address.slice('::ffff:'.length);
+  return address;
+};
+
+const isAllowedPeer = (address, allowedPrefix) => {
+  const normalized = normalizeIPv4(address);
+  if (!normalized || net.isIP(normalized) !== 4) return address === '::1';
+  if (normalized.startsWith('127.')) return true;
+  if (!isPrivateIPv4(normalized)) return false;
+  return allowedPrefix === undefined ? true : normalized.startsWith(allowedPrefix);
+};
+
+const normalizeBindHost = (host) => {
+  if (typeof host !== 'string' || net.isIP(host) !== 4 || !isPrivateIPv4(host)) {
+    throw new GenerationProtocolError(
+      500,
+      'unsafe_configuration',
+      'generation host must be a private IPv4 address',
+    );
+  }
+  return host;
+};
+
+const normalizePort = (port) => {
+  if (port === undefined || port === null || port === '') return DEFAULT_GENERATION_PORT;
+  const value = Number(port);
+  if (!Number.isInteger(value) || value < 1 || value > 65535) {
+    throw new GenerationProtocolError(500, 'unsafe_configuration', 'generation port is invalid');
+  }
+  return value;
+};
+
+const tokenMatches = (header, token) => {
+  if (typeof header !== 'string' || !header.startsWith('Bearer ')) return false;
+  const supplied = Buffer.from(header.slice('Bearer '.length));
+  const expected = Buffer.from(token);
+  return supplied.length === expected.length && timingSafeEqual(supplied, expected);
+};
+
+const writeJson = (response, status, value, maximum = MAX_RESPONSE_BYTES) => {
+  const bytes = Buffer.from(JSON.stringify(value));
+  if (bytes.length > maximum) {
+    response.statusCode = 500;
+    response.setHeader('content-type', 'application/json');
+    response.end(
+      JSON.stringify({
+        error: { code: 'response_too_large', message: 'response exceeds its size limit' },
+      }),
+    );
+    return;
+  }
+  response.statusCode = status;
+  response.setHeader('cache-control', 'no-store');
+  response.setHeader('content-type', 'application/json');
+  response.setHeader('content-length', bytes.length);
+  response.end(bytes);
+};
+
+const readBody = (request) =>
+  new Promise((resolve, reject) => {
+    const declared = Number(request.headers['content-length']);
+    if (Number.isFinite(declared) && declared > MAX_REQUEST_BYTES) {
+      reject(httpError(413, 'request_too_large', 'request exceeds its size limit'));
+      request.resume();
+      return;
+    }
+    const chunks = [];
+    let length = 0;
+    let settled = false;
+    const fail = (error) => {
+      if (settled) return;
+      settled = true;
+      reject(error);
+      request.destroy();
+    };
+    request.on('data', (chunk) => {
+      length += chunk.length;
+      if (length > MAX_REQUEST_BYTES) {
+        fail(httpError(413, 'request_too_large', 'request exceeds its size limit'));
+        return;
+      }
+      chunks.push(chunk);
+    });
+    request.on('end', () => {
+      if (settled) return;
+      settled = true;
+      resolve(Buffer.concat(chunks, length));
+    });
+    request.on('error', (error) => fail(error));
+  });
+
+const parseBody = async (request) => {
+  const bytes = await readBody(request);
+  if (bytes.length === 0) {
+    throw httpError(400, 'invalid_request', 'request body is required');
+  }
+  try {
+    return JSON.parse(bytes.toString('utf8'));
+  } catch {
+    throw httpError(400, 'invalid_request', 'request body is not valid JSON');
+  }
+};
+
+const safeHeader = (value) => (Array.isArray(value) ? value[0] : value);
+
+const requestRoute = (request) => {
+  try {
+    const segments = new URL(request.url ?? '/', 'http://generation.local').pathname
+      .split('/')
+      .filter(Boolean);
+    if (segments.length === 2 && segments[0] === 'v1' && segments[1] === 'models') {
+      return 'models';
+    }
+    if (segments.length === 2 && segments[0] === 'v1' && segments[1] === 'jobs') {
+      return 'jobs-create';
+    }
+    if (segments.length === 3 && segments[0] === 'v1' && segments[1] === 'jobs') {
+      return 'jobs-read';
+    }
+    if (
+      segments.length === 4 &&
+      segments[0] === 'v1' &&
+      segments[1] === 'jobs' &&
+      segments[3] === 'cancel'
+    ) {
+      return 'jobs-cancel';
+    }
+  } catch {
+    return 'unknown';
+  }
+  return 'unknown';
+};
+
+const validUsage = (value) => {
+  const usage = asObject(value);
+  if (!usage || Array.isArray(usage)) return undefined;
+  const inputTokens = usage.inputTokens;
+  const outputTokens = usage.outputTokens;
+  const totalTokens = usage.totalTokens;
+  if (
+    ![inputTokens, outputTokens, totalTokens].every(
+      (item) => Number.isInteger(item) && item >= 0 && item <= 10_000_000,
+    )
+  ) {
+    throw httpError(500, 'invalid_response', 'generation executor returned invalid usage metadata');
+  }
+  return { inputTokens, outputTokens, totalTokens };
+};
+
+const genericResult = (record, responseBytes, imageBytes) => {
+  let response;
+  try {
+    response = JSON.parse(responseBytes.toString('utf8'));
+  } catch {
+    throw httpError(
+      500,
+      'invalid_response',
+      'generation executor returned invalid result metadata',
+    );
+  }
+  if (
+    !response ||
+    typeof response !== 'object' ||
+    !Number.isInteger(response.generatedAt) ||
+    response.generatedAt < 0 ||
+    typeof response.model !== 'string' ||
+    response.model !== record.model ||
+    !Array.isArray(response.output) ||
+    response.output.length > 64
+  ) {
+    throw httpError(
+      500,
+      'invalid_response',
+      'generation executor returned invalid result metadata',
+    );
+  }
+  if (
+    Object.keys(response).some((key) => !['generatedAt', 'model', 'output', 'usage'].includes(key))
+  ) {
+    throw httpError(
+      500,
+      'invalid_response',
+      'generation executor returned unsupported result metadata',
+    );
+  }
+  const output = [];
+  const referencedImages = new Set();
+  const requestedOutput = new Set(record.request.output);
+  for (const part of response.output) {
+    if (!part || typeof part !== 'object') {
+      throw httpError(
+        500,
+        'invalid_response',
+        'generation executor returned an invalid output part',
+      );
+    }
+    if (part.type === 'text') {
+      if (Object.keys(part).some((key) => !['type', 'text'].includes(key))) {
+        throw httpError(
+          500,
+          'invalid_response',
+          'generation executor returned an invalid text output',
+        );
+      }
+      if (!requestedOutput.has('text')) {
+        throw httpError(
+          500,
+          'invalid_response',
+          'generation executor returned unrequested text output',
+        );
+      }
+      if (typeof part.text !== 'string' || Buffer.byteLength(part.text) > 2 * 1024 * 1024) {
+        throw httpError(
+          500,
+          'invalid_response',
+          'generation executor returned invalid text output',
+        );
+      }
+      output.push({ type: 'text', text: part.text });
+      continue;
+    }
+    if (part.type !== 'image' || part.mimeType !== 'image/png') {
+      throw httpError(
+        500,
+        'invalid_response',
+        'generation executor returned an unsupported output part',
+      );
+    }
+    if (Object.keys(part).some((key) => !['type', 'mimeType', 'name'].includes(key))) {
+      throw httpError(
+        500,
+        'invalid_response',
+        'generation executor returned an invalid image output',
+      );
+    }
+    if (!requestedOutput.has('image')) {
+      throw httpError(
+        500,
+        'invalid_response',
+        'generation executor returned unrequested image output',
+      );
+    }
+    if (
+      typeof part.name !== 'string' ||
+      !/^image-[0-3]\.png$/.test(part.name) ||
+      referencedImages.has(part.name)
+    ) {
+      throw httpError(
+        500,
+        'invalid_response',
+        'generation executor returned an invalid image reference',
+      );
+    }
+    const bytes = imageBytes.get(part.name);
+    if (!bytes) {
+      throw httpError(
+        500,
+        'invalid_response',
+        'generation executor returned a missing image artifact',
+      );
+    }
+    referencedImages.add(part.name);
+    output.push({
+      type: 'image',
+      mimeType: 'image/png',
+      data: bytes.toString('base64'),
+      name: part.name,
+    });
+  }
+  for (const name of imageBytes.keys()) {
+    if (!referencedImages.has(name)) {
+      throw httpError(
+        500,
+        'invalid_response',
+        'generation executor uploaded an unreferenced image artifact',
+      );
+    }
+  }
+  const result = {
+    output,
+    model: response.model,
+    generatedAt: response.generatedAt,
+  };
+  const usage = validUsage(response.usage);
+  if (usage) result.usage = usage;
+  return result;
+};
+
+const createLocalRunner = ({ store, record }) => {
+  const leaseToken = ['local', randomUUID()].join('-');
+  const leaseExpiresAt = Date.now() + JOB_LEASE_MS;
+  const uploads = new Map();
+  const job = {
+    id: record.id,
+    runId: record.id,
+    leaseToken,
+    leaseExpiresAt,
+    model: record.model,
+    reasoningEffort: record.reasoningEffort,
+  };
+  return {
+    job,
+    input() {
+      return Promise.resolve(record.request);
+    },
+    leaseHeartbeat() {
+      return store.isCancellationRequested(record.id).then((cancelRequested) => ({
+        cancelRequested,
+        leaseExpiresAt,
+      }));
+    },
+    upload(_jobId, _token, bytes, name) {
+      if (!(bytes instanceof Buffer))
+        return Promise.reject(new Error('generation upload must be bytes'));
+      if (name !== 'response.json' && !/^image-[0-3]\.png$/.test(name)) {
+        return Promise.reject(new Error('generation upload name is invalid'));
+      }
+      uploads.set(name, Buffer.from(bytes));
+      return Promise.resolve({ key: `${LOCAL_RESPONSE_PREFIX}${record.id}/${name}` });
+    },
+    complete() {
+      const response = uploads.get('response.json');
+      if (!response) return Promise.reject(new Error('generation response upload is missing'));
+      const images = new Map(
+        [...uploads.entries()]
+          .filter(([name]) => /^image-[0-3]\.png$/.test(name))
+          .map(([name, bytes]) => [name, bytes]),
+      );
+      return store.complete(record.id, genericResult(record, response, images));
+    },
+    fail(_jobId, _token, message, outcome) {
+      return store.fail(record.id, message, outcome);
+    },
+  };
+};
+
+class GenerationServer {
+  constructor({
+    host,
+    port,
+    token,
+    store,
+    getModels,
+    isReady = () => true,
+    onQueued,
+    analytics,
+    allowedPeerPrefix,
+  } = {}) {
+    if (typeof token !== 'string' || token.length === 0) {
+      throw new GenerationProtocolError(
+        500,
+        'unsafe_configuration',
+        'generation token is required',
+      );
+    }
+    this.host = normalizeBindHost(host);
+    this.port = normalizePort(port);
+    this.token = String(token);
+    this.store = store;
+    this.getModels = getModels ?? (() => []);
+    this.isReady = isReady;
+    this.onQueued = onQueued;
+    this.analytics = analytics ?? noopAnalytics;
+    this.allowedPeerPrefix = allowedPeerPrefix;
+    this.activeBodyReads = 0;
+    this.server = createServer((request, response) => {
+      void this.handle(request, response);
+    });
+    this.listening = false;
+  }
+
+  async handle(request, response) {
+    const route = requestRoute(request);
+    void this.analytics.emit('http', route, 'received');
+    response.once('finish', () => {
+      void this.analytics.emit('http', route, 'responded', response.statusCode);
+    });
+    response.setHeader('cache-control', 'no-store');
+    response.setHeader('x-content-type-options', 'nosniff');
+    if (!isAllowedPeer(request.socket.remoteAddress, this.allowedPeerPrefix)) {
+      void this.analytics.emit('http', route, 'peer-rejected');
+      writeJson(response, 403, {
+        error: { code: 'forbidden', message: 'peer is outside the private LAN' },
+      });
+      return;
+    }
+    void this.analytics.emit('http', route, 'peer-accepted');
+    if (!tokenMatches(request.headers.authorization, this.token)) {
+      void this.analytics.emit('http', route, 'authentication-rejected');
+      writeJson(response, 401, {
+        error: { code: 'unauthorized', message: 'invalid generation credential' },
+      });
+      return;
+    }
+    void this.analytics.emit('http', route, 'authenticated');
+    try {
+      const parsed = new URL(request.url ?? '/', 'http://generation.local');
+      if (parsed.search || parsed.hash) {
+        throw httpError(404, 'not_found', 'endpoint was not found');
+      }
+      const segments = parsed.pathname.split('/').filter(Boolean);
+      if (segments.length === 2 && segments[0] === 'v1' && segments[1] === 'models') {
+        if (request.method !== 'GET')
+          throw httpError(405, 'method_not_allowed', 'method is not allowed');
+        const models = await this.getModels();
+        void this.analytics.emit('http', route, 'model-catalog-loaded');
+        writeJson(
+          response,
+          200,
+          publicModelCatalog(models, { ready: this.isReady() }),
+          MAX_MODEL_RESPONSE_BYTES,
+        );
+        return;
+      }
+      if (segments.length === 2 && segments[0] === 'v1' && segments[1] === 'jobs') {
+        if (request.method !== 'POST')
+          throw httpError(405, 'method_not_allowed', 'method is not allowed');
+        await this.createJob(request, response);
+        return;
+      }
+      if (segments.length === 3 && segments[0] === 'v1' && segments[1] === 'jobs') {
+        if (request.method !== 'GET')
+          throw httpError(405, 'method_not_allowed', 'method is not allowed');
+        const record = await this.store.get(segments[2]);
+        if (!record) throw httpError(404, 'not_found', 'generation job was not found');
+        if (record.bodyPresent === false && record.status === 'completed') {
+          throw new GenerationStoreError('generation result has expired', 'result_expired');
+        }
+        void this.analytics.emit('http', route, 'job-loaded', record.status);
+        writeJson(response, 200, this.store.publicJob(record));
+        return;
+      }
+      if (
+        segments.length === 4 &&
+        segments[0] === 'v1' &&
+        segments[1] === 'jobs' &&
+        segments[3] === 'cancel'
+      ) {
+        if (request.method !== 'POST')
+          throw httpError(405, 'method_not_allowed', 'method is not allowed');
+        const record = await this.store.cancel(segments[2]);
+        if (!record) throw httpError(404, 'not_found', 'generation job was not found');
+        void this.analytics.emit('http', route, 'cancellation-recorded', record.status);
+        writeJson(response, 200, this.store.publicJob(record));
+        return;
+      }
+      throw httpError(404, 'not_found', 'endpoint was not found');
+    } catch (error) {
+      const payload = errorPayload(error);
+      void this.analytics.emit('http', route, 'failed', payload.body.error.code);
+      writeJson(response, payload.status, payload.body);
+    }
+  }
+
+  async createJob(request, response) {
+    const key = normalizeIdempotencyKey(safeHeader(request.headers['idempotency-key']));
+    if (this.activeBodyReads >= MAX_CONCURRENT_BODY_READS) {
+      request.resume();
+      throw new GenerationProtocolError(429, 'overloaded', 'generation request capacity is full');
+    }
+    this.activeBodyReads += 1;
+    try {
+      void this.analytics.emit('job', 'local', 'body-read-started');
+      const body = await parseBody(request);
+      void this.analytics.emit('job', 'local', 'body-read');
+      const idempotencyKeyHash = idempotencyDigest(key);
+      const digest = requestDigest(body);
+      const existing = await this.store.findByIdempotencyKeyHash(idempotencyKeyHash);
+      void this.analytics.emit('job', 'local', 'idempotency-checked');
+      if (existing) {
+        if (existing.requestHash !== digest) {
+          throw new GenerationStoreError(
+            'idempotency key was already used for a different request',
+            'idempotency_conflict',
+          );
+        }
+        void this.analytics.emit('job', 'local', 'idempotency-reused', existing.status);
+        writeJson(response, 200, { id: existing.id, status: existing.status });
+        return;
+      }
+      const models = await this.getModels();
+      void this.analytics.emit('job', 'local', 'model-catalog-loaded');
+      const normalized = normalizeGenerationRequest(body, models);
+      void this.analytics.emit('job', 'local', 'request-validated');
+      const created = await this.store.create({
+        idempotencyKeyHash,
+        requestHash: digest,
+        request: normalized,
+        model: normalized.model,
+        reasoningEffort: normalized.reasoningEffort,
+      });
+      void this.analytics.emit(
+        'job',
+        'local',
+        created.created ? 'queued' : 'idempotency-reused-after-race',
+        created.record.status,
+      );
+      if (created.created) this.onQueued?.();
+      writeJson(response, created.created ? 202 : 200, {
+        id: created.record.id,
+        status: created.record.status,
+      });
+    } finally {
+      this.activeBodyReads -= 1;
+    }
+  }
+
+  async listen() {
+    if (this.listening) return this.address();
+    await new Promise((resolve, reject) => {
+      const onError = (error) => {
+        this.server.off('listening', onListening);
+        reject(error);
+      };
+      const onListening = () => {
+        this.server.off('error', onError);
+        this.listening = true;
+        resolve();
+      };
+      this.server.once('error', onError);
+      this.server.once('listening', onListening);
+      this.server.listen(this.port, this.host);
+    });
+    return this.address();
+  }
+
+  address() {
+    const address = this.server.address();
+    return typeof address === 'object' && address
+      ? { address: address.address, port: address.port }
+      : undefined;
+  }
+
+  async close() {
+    if (!this.listening) return;
+    await new Promise((resolve) => this.server.close(() => resolve()));
+    this.listening = false;
+  }
+}
+
+const startGenerationServer = async (options) => {
+  const server = new GenerationServer(options);
+  await server.store.init();
+  await server.listen();
+  return server;
+};
+
+export {
+  DEFAULT_GENERATION_PORT,
+  GenerationServer,
+  MAX_REQUEST_BYTES,
+  createLocalRunner,
+  isAllowedPeer,
+  normalizeBindHost,
+  startGenerationServer,
+};