@vaultgradient/pq-agent 0.1.0

package/dist/index.js

@@ -0,0 +1,645 @@
+#!/usr/bin/env node
+
+// src/index.ts
+import { Command } from "commander";
+import { homedir } from "os";
+import { join, resolve as resolve2 } from "path";
+
+// src/client.ts
+import { randomUUID as randomUUID2 } from "crypto";
+import { EventEmitter } from "events";
+import WebSocket from "ws";
+
+// ../protocol/dist/errors.js
+var WS_CLOSE_CODES = {
+  GOING_AWAY: 1001,
+  INTERNAL_ERROR: 1011,
+  INVALID_AGENT_TOKEN: 4001,
+  TENANT_NOT_ENROLLED: 4002,
+  PROTOCOL_VERSION_MISMATCH: 4003
+};
+var NO_RECONNECT_CLOSE_CODES = /* @__PURE__ */ new Set([
+  WS_CLOSE_CODES.INVALID_AGENT_TOKEN,
+  WS_CLOSE_CODES.TENANT_NOT_ENROLLED,
+  WS_CLOSE_CODES.PROTOCOL_VERSION_MISMATCH
+]);
+
+// ../protocol/dist/version.js
+var PROTOCOL_VERSION = 1;
+var PROTOCOL_VERSION_HEADER = "x-pq-protocol-version";
+
+// src/audit.ts
+import { createHash } from "crypto";
+var ZERO_HASH = "0".repeat(64);
+function canonicalJson(value) {
+  if (value === null) return "null";
+  if (typeof value === "number" || typeof value === "boolean") {
+    return JSON.stringify(value);
+  }
+  if (typeof value === "string") return JSON.stringify(value);
+  if (Array.isArray(value)) {
+    return "[" + value.map(canonicalJson).join(",") + "]";
+  }
+  if (typeof value === "object") {
+    const obj = value;
+    const keys = Object.keys(obj).filter((k) => obj[k] !== void 0).sort();
+    const parts = keys.map(
+      (k) => `${JSON.stringify(k)}:${canonicalJson(obj[k])}`
+    );
+    return "{" + parts.join(",") + "}";
+  }
+  throw new Error(`canonicalJson: unsupported value of type ${typeof value}`);
+}
+function computeHash(prevHash, entry) {
+  return createHash("sha256").update(prevHash).update(canonicalJson(entry)).digest("hex");
+}
+var AuditEmitter = class {
+  constructor(cfg) {
+    this.cfg = cfg;
+  }
+  cfg;
+  prevHash = ZERO_HASH;
+  emit(partial) {
+    const { ts: tsOverride, ...rest } = partial;
+    const entry = {
+      tenant_id: this.cfg.tenantId,
+      agent_id: this.cfg.agentId,
+      ts: tsOverride ?? (/* @__PURE__ */ new Date()).toISOString(),
+      ...rest
+    };
+    const prev_hash = this.prevHash;
+    const hash = computeHash(prev_hash, entry);
+    this.prevHash = hash;
+    return { entry, prev_hash, hash };
+  }
+  // For tests / observability.
+  currentHash() {
+    return this.prevHash;
+  }
+};
+
+// src/dispatch.ts
+import { randomUUID } from "crypto";
+var REQUEST_TIMEOUT_MS = 3e4;
+var MAX_ROWS = 1e5;
+var MAX_BYTES = 10 * 1024 * 1024;
+async function dispatch(inbound, opts) {
+  switch (inbound.type) {
+    case "ping":
+      return { outbound: [{ type: "pong", id: inbound.id, payload: {} }] };
+    case "request.list_sources":
+      return wrapRequest(inbound.id, opts, () => listSources(opts.runtime));
+    case "request.describe_source":
+      return wrapRequest(
+        inbound.id,
+        opts,
+        () => describeSource(opts.runtime, inbound.payload.name, inbound.payload.sample_size)
+      );
+    case "request.list_endpoints":
+      return wrapRequest(inbound.id, opts, () => listEndpoints(opts.runtime));
+    case "request.call_endpoint":
+      return wrapRequest(
+        inbound.id,
+        opts,
+        () => callEndpoint(opts.runtime, inbound.payload.path)
+      );
+    case "request.query":
+      return wrapRequest(
+        inbound.id,
+        opts,
+        () => runQuery(opts.runtime, inbound.payload.expression)
+      );
+    case "control.rotate_token":
+    case "control.update_config":
+    case "control.shutdown":
+      return { outbound: [] };
+  }
+}
+async function wrapRequest(requestId, opts, handler) {
+  const timeoutMs = opts.timeoutMs ?? REQUEST_TIMEOUT_MS;
+  const startedAt = Date.now();
+  let result;
+  try {
+    result = await Promise.race([
+      handler(),
+      timeoutAfter(timeoutMs)
+    ]);
+  } catch (err) {
+    if (err instanceof TimeoutError) {
+      return {
+        outbound: [
+          {
+            type: "response.error",
+            id: requestId,
+            payload: {
+              code: "timeout",
+              message: `request exceeded ${timeoutMs}ms`
+            }
+          }
+        ]
+      };
+    }
+    return {
+      outbound: [
+        {
+          type: "response.error",
+          id: requestId,
+          payload: {
+            code: "execution",
+            message: err instanceof Error ? err.message : String(err)
+          }
+        }
+      ]
+    };
+  }
+  const latency_ms = Date.now() - startedAt;
+  if (!result.ok) {
+    return {
+      outbound: [
+        {
+          type: "response.error",
+          id: requestId,
+          payload: { code: result.code, message: result.message }
+        }
+      ]
+    };
+  }
+  const bytes = byteLengthOfJson(result.rows);
+  if (result.rows.length > MAX_ROWS || bytes > MAX_BYTES) {
+    return {
+      outbound: [
+        {
+          type: "response.error",
+          id: requestId,
+          payload: {
+            code: "too_large",
+            message: `result exceeded cap (rows=${result.rows.length}, bytes=${bytes})`,
+            details: {
+              row_count: result.rows.length,
+              bytes,
+              max_rows: MAX_ROWS,
+              max_bytes: MAX_BYTES
+            }
+          }
+        }
+      ]
+    };
+  }
+  return {
+    outbound: [
+      {
+        type: "response.chunk",
+        id: requestId,
+        payload: { rows: result.rows, chunk_index: 0 }
+      },
+      {
+        type: "response.end",
+        id: requestId,
+        payload: {
+          total_rows: result.rows.length,
+          latency_ms,
+          // The published SourceManager.runQuery() returns rows without
+          // exposing whether push-down was used. Slice 4 reports false until
+          // the OSS API surfaces this; the wire field is still in the spec
+          // and consumed by future audit/billing observability.
+          pushed_down: false
+        }
+      }
+    ]
+  };
+}
+var TimeoutError = class extends Error {
+};
+function timeoutAfter(ms) {
+  return new Promise(
+    (_, reject) => setTimeout(() => reject(new TimeoutError()), ms)
+  );
+}
+function byteLengthOfJson(value) {
+  return Buffer.byteLength(JSON.stringify(value), "utf8");
+}
+async function listSources(runtime) {
+  const names = runtime.sources.getSourceNames();
+  const statuses = runtime.sources.getAllStatuses();
+  const rows = names.map((name) => ({
+    name,
+    status: statuses[name] ?? null
+  }));
+  return { ok: true, rows };
+}
+async function describeSource(runtime, name, sampleSize) {
+  const status = runtime.sources.getSourceStatus(name);
+  if (!status) {
+    return { ok: false, code: "not_found", message: `source "${name}" is not configured` };
+  }
+  const data = runtime.sources.getSourceData(name) ?? [];
+  const size = Math.min(data.length, Math.max(1, sampleSize ?? 5));
+  const sample = data.slice(0, size);
+  const fields = /* @__PURE__ */ new Set();
+  for (const row of sample) {
+    if (row && typeof row === "object") {
+      for (const k of Object.keys(row)) fields.add(k);
+    }
+  }
+  return {
+    ok: true,
+    rows: [
+      {
+        name,
+        status,
+        fields: [...fields],
+        sample
+      }
+    ]
+  };
+}
+async function listEndpoints(runtime) {
+  const rows = [...runtime.endpoints.values()];
+  return { ok: true, rows };
+}
+async function callEndpoint(runtime, path) {
+  const endpoint = runtime.endpoints.get(path);
+  if (!endpoint) {
+    return { ok: false, code: "not_found", message: `endpoint "${path}" is not registered` };
+  }
+  return runQuery(runtime, endpoint.query);
+}
+async function runQuery(runtime, expression) {
+  try {
+    const result = await runtime.sources.runQuery(expression);
+    const rows = Array.isArray(result) ? result : [result];
+    return { ok: true, rows };
+  } catch (err) {
+    const message = err instanceof Error ? err.message : String(err);
+    return { ok: false, code: "execution", message };
+  }
+}
+function isInboundEnvelope(value) {
+  if (typeof value !== "object" || value === null) return false;
+  const env = value;
+  return typeof env.type === "string" && typeof env.id === "string" && typeof env.payload === "object" && env.payload !== null;
+}
+
+// src/reconnect.ts
+var Backoff = class {
+  initialMs;
+  ceilingMs;
+  resetAfterStableMs;
+  currentMs;
+  constructor(opts = {}) {
+    this.initialMs = opts.initialMs ?? 1e3;
+    this.ceilingMs = opts.ceilingMs ?? 6e4;
+    this.resetAfterStableMs = opts.resetAfterStableMs ?? 5 * 6e4;
+    this.currentMs = this.initialMs;
+  }
+  // Returns the delay (ms) to wait before the next reconnect attempt, then
+  // doubles the cap for the following attempt (up to ceiling).
+  nextDelayMs(random = Math.random) {
+    const delay = Math.floor(random() * this.currentMs);
+    this.currentMs = Math.min(this.currentMs * 2, this.ceilingMs);
+    return delay;
+  }
+  reset() {
+    this.currentMs = this.initialMs;
+  }
+};
+
+// src/client.ts
+function createClient(opts) {
+  const emitter = new EventEmitter();
+  const backoff = new Backoff();
+  const audit = new AuditEmitter({
+    tenantId: opts.tenantId ?? "dev-tenant",
+    agentId: opts.agentId ?? "dev-agent"
+  });
+  let socket = null;
+  let stopped = false;
+  let reconnectTimer = null;
+  let stableTimer = null;
+  function assertSafeUrl(url) {
+    if (url.startsWith("wss://")) return;
+    if (url.startsWith("ws://")) {
+      if (!opts.allowInsecure) {
+        throw new Error(
+          `Refusing to connect: URL must use wss:// (got ${url}). Pass allowInsecure for local dev only.`
+        );
+      }
+      let host;
+      try {
+        host = new URL(url).hostname;
+      } catch {
+        throw new Error(`Refusing to connect: invalid URL (${url}).`);
+      }
+      const localish = host === "localhost" || host === "127.0.0.1" || host === "::1" || host === "0.0.0.0" || host.endsWith(".local");
+      if (!localish) {
+        throw new Error(
+          `Refusing to connect: --insecure is only allowed for localhost targets, got host="${host}". Use wss:// for ${url}.`
+        );
+      }
+      return;
+    }
+    throw new Error(
+      `Refusing to connect: URL must be ws:// or wss:// (got ${url}).`
+    );
+  }
+  async function connect() {
+    assertSafeUrl(opts.url);
+    const ws = new WebSocket(opts.url, {
+      headers: {
+        Authorization: `Bearer ${opts.token}`,
+        [PROTOCOL_VERSION_HEADER]: String(PROTOCOL_VERSION)
+      }
+    });
+    socket = ws;
+    ws.on("open", () => {
+      stableTimer = setTimeout(
+        () => backoff.reset(),
+        backoff.resetAfterStableMs
+      );
+      emitter.emit("open");
+    });
+    ws.on("message", (raw) => {
+      void handleFrame(ws, raw.toString());
+    });
+    ws.on("close", (code) => {
+      if (stableTimer) {
+        clearTimeout(stableTimer);
+        stableTimer = null;
+      }
+      socket = null;
+      emitter.emit("close", code);
+      if (stopped) return;
+      if (NO_RECONNECT_CLOSE_CODES.has(code)) {
+        emitter.emit(
+          "error",
+          new Error(
+            `WS closed ${code} \u2014 non-recoverable per \xA77. Re-enrollment or operator action required.`
+          )
+        );
+        return;
+      }
+      scheduleReconnect();
+    });
+    ws.on("error", (err) => {
+      emitter.emit("error", err);
+    });
+  }
+  async function handleFrame(ws, raw) {
+    let parsed;
+    try {
+      parsed = JSON.parse(raw);
+    } catch (err) {
+      emitter.emit("error", new Error(`Invalid JSON frame: ${String(err)}`));
+      return;
+    }
+    if (!isInboundEnvelope(parsed)) {
+      emitter.emit("error", new Error("Frame missing required envelope fields"));
+      return;
+    }
+    const inbound = parsed;
+    const shouldAudit = inbound.type.startsWith("request.");
+    const start = Date.now();
+    const result = await dispatch(inbound, {
+      runtime: opts.runtime,
+      ...opts.requestTimeoutMs !== void 0 && {
+        timeoutMs: opts.requestTimeoutMs
+      }
+    });
+    const latency_ms = Date.now() - start;
+    for (const outbound of result.outbound) {
+      send(ws, outbound);
+    }
+    if (shouldAudit) {
+      const auditEnvelope = buildAuditEvent(
+        inbound,
+        result.outbound,
+        latency_ms
+      );
+      send(ws, auditEnvelope);
+    }
+  }
+  function buildAuditEvent(inbound, outbound, latency_ms) {
+    const summary = summarizeOutbound(outbound);
+    let expression;
+    let endpoint_path;
+    if (inbound.type === "request.query") expression = inbound.payload.expression;
+    if (inbound.type === "request.call_endpoint") endpoint_path = inbound.payload.path;
+    const auditEvent = audit.emit({
+      request_id: inbound.id,
+      caller_id: "unknown",
+      // §13 spec gap — see audit.ts header
+      message_type: inbound.type,
+      ...expression !== void 0 && { expression },
+      ...endpoint_path !== void 0 && { endpoint_path },
+      pushed_down: summary.pushed_down,
+      row_count: summary.row_count,
+      bytes: summary.bytes,
+      latency_ms,
+      outcome: summary.outcome,
+      ...summary.error_code !== void 0 && { error_code: summary.error_code }
+    });
+    return {
+      type: "event.audit",
+      id: randomUUID2(),
+      payload: auditEvent
+    };
+  }
+  function send(ws, envelope) {
+    if (ws.readyState !== WebSocket.OPEN) return;
+    ws.send(JSON.stringify(envelope));
+  }
+  function scheduleReconnect() {
+    const delay = backoff.nextDelayMs();
+    reconnectTimer = setTimeout(() => {
+      reconnectTimer = null;
+      if (stopped) return;
+      void connect().catch((err) => {
+        emitter.emit("error", err);
+        scheduleReconnect();
+      });
+    }, delay);
+  }
+  return {
+    async start() {
+      stopped = false;
+      await connect();
+    },
+    async stop() {
+      stopped = true;
+      if (reconnectTimer) {
+        clearTimeout(reconnectTimer);
+        reconnectTimer = null;
+      }
+      if (stableTimer) {
+        clearTimeout(stableTimer);
+        stableTimer = null;
+      }
+      if (socket) {
+        socket.close();
+        socket = null;
+      }
+    },
+    on(event, listener) {
+      emitter.on(event, listener);
+    }
+  };
+}
+function summarizeOutbound(outbound) {
+  let row_count = 0;
+  let bytes = 0;
+  let pushed_down = false;
+  let outcome = "success";
+  let error_code;
+  for (const env of outbound) {
+    if (env.type === "response.error") {
+      outcome = "error";
+      error_code = env.payload.code;
+    } else if (env.type === "response.chunk") {
+      bytes += JSON.stringify(env.payload).length;
+    } else if (env.type === "response.end") {
+      row_count = env.payload.total_rows;
+      pushed_down = env.payload.pushed_down;
+    }
+  }
+  return error_code !== void 0 ? { outcome, error_code, row_count, bytes, pushed_down } : { outcome, row_count, bytes, pushed_down };
+}
+
+// src/enrollment.ts
+import { promises as fs } from "fs";
+import { dirname } from "path";
+async function readAgentToken(path) {
+  try {
+    const raw = await fs.readFile(path, "utf8");
+    return JSON.parse(raw);
+  } catch (err) {
+    if (err.code === "ENOENT") return null;
+    throw err;
+  }
+}
+function resolveBearerToken(stored, enrollmentTokenEnv) {
+  if (stored) return { token: stored.agentToken, kind: "agent" };
+  if (enrollmentTokenEnv) {
+    return { token: enrollmentTokenEnv, kind: "enrollment" };
+  }
+  throw new Error(
+    "No agent token on disk and no PQ_AGENT_ENROLLMENT_TOKEN set. Run `pq-agent enroll <token>` first, or set the env var for first connect."
+  );
+}
+
+// src/runtime.ts
+import { dirname as dirname2, resolve } from "path";
+import { SourceManager } from "@vaultgradient/pipequery-cli/sources";
+
+// src/config.ts
+import { readFile } from "fs/promises";
+import { parse as parseYaml } from "yaml";
+async function loadPipequeryConfig(path) {
+  const raw = await readFile(path, "utf8");
+  const parsed = parseYaml(raw);
+  if (typeof parsed !== "object" || parsed === null) {
+    throw new Error(`pipequery config at ${path} did not parse to an object`);
+  }
+  return parsed;
+}
+
+// src/runtime.ts
+var AgentRuntime = class _AgentRuntime {
+  sources;
+  endpoints;
+  constructor(sources, endpoints) {
+    this.sources = sources;
+    this.endpoints = endpoints;
+  }
+  static async fromConfigFile(configPath) {
+    const absConfig = resolve(configPath);
+    const cwd = dirname2(absConfig);
+    const cfg = await loadPipequeryConfig(absConfig);
+    return _AgentRuntime.fromConfig(cfg, cwd);
+  }
+  static async fromConfig(config, cwd) {
+    const sources = new SourceManager(cwd);
+    for (const [name, sourceConfig] of Object.entries(config.sources ?? {})) {
+      await sources.addSource(name, sourceConfig);
+    }
+    const endpoints = /* @__PURE__ */ new Map();
+    for (const [path, ep] of Object.entries(config.endpoints ?? {})) {
+      endpoints.set(path, { path, query: ep.query });
+    }
+    return new _AgentRuntime(sources, endpoints);
+  }
+  async dispose() {
+    await this.sources.dispose();
+  }
+};
+
+// src/index.ts
+var DEFAULT_TOKEN_PATH = join(
+  process.env.XDG_DATA_HOME ?? join(homedir(), ".local", "share"),
+  "pipequery",
+  "agent.token"
+);
+var program = new Command();
+program.name("pq-agent").description(
+  "PipeQuery hosted agent \u2014 dials the control plane over WSS and serves MCP traffic from inside your network."
+).version("0.0.0");
+program.command("run").description("Connect to the control plane and serve requests.").option(
+  "--url <url>",
+  "Control-plane WSS URL",
+  process.env.PQ_AGENT_CONTROL_PLANE_URL ?? "ws://127.0.0.1:8787"
+).option(
+  "--config <path>",
+  "Path to pipequery.yaml (sources + endpoints)",
+  process.env.PQ_AGENT_CONFIG ?? "./pipequery.yaml"
+).option(
+  "--token-file <path>",
+  "Where to read/write the agent token",
+  process.env.PQ_AGENT_TOKEN_FILE ?? DEFAULT_TOKEN_PATH
+).option(
+  "--insecure",
+  "Allow ws:// (dev only \u2014 refuses against prod control plane in real impl)",
+  false
+).action(
+  async (opts) => {
+    const stored = await readAgentToken(opts.tokenFile);
+    const bearer = resolveBearerToken(
+      stored,
+      process.env.PQ_AGENT_ENROLLMENT_TOKEN
+    );
+    console.log(`[agent] loading config from ${resolve2(opts.config)}`);
+    const runtime = await AgentRuntime.fromConfigFile(opts.config);
+    console.log(
+      `[agent] loaded ${runtime.sources.getSourceNames().length} source(s), ${runtime.endpoints.size} endpoint(s)`
+    );
+    const client = createClient({
+      url: opts.url,
+      token: bearer.token,
+      runtime,
+      allowInsecure: opts.insecure
+    });
+    client.on("open", () => {
+      console.log(`[agent] connected to ${opts.url} (auth=${bearer.kind})`);
+    });
+    client.on("close", (code) => {
+      console.log(`[agent] disconnected (code=${code})`);
+    });
+    client.on("error", (err) => {
+      console.error(`[agent] error: ${err.message}`);
+    });
+    const shutdown = async (sig) => {
+      console.log(`[agent] ${sig} received, draining\u2026`);
+      await client.stop();
+      await runtime.dispose();
+      process.exit(0);
+    };
+    process.on("SIGINT", () => void shutdown("SIGINT"));
+    process.on("SIGTERM", () => void shutdown("SIGTERM"));
+    await client.start();
+  }
+);
+program.command("status").description("Print the current agent token + control-plane config.").action(() => {
+  console.log("status command not implemented yet");
+});
+program.parseAsync(process.argv).catch((err) => {
+  console.error(err.message);
+  process.exit(1);
+});

package/dist/runtime.d.ts

@@ -0,0 +1,23 @@
+import { SourceManager } from '@vaultgradient/pipequery-cli/sources';
+
+interface PipequeryYaml {
+    sources?: Record<string, unknown>;
+    endpoints?: Record<string, {
+        query: string;
+    }>;
+}
+
+interface EndpointEntry {
+    path: string;
+    query: string;
+}
+declare class AgentRuntime {
+    readonly sources: SourceManager;
+    readonly endpoints: Map<string, EndpointEntry>;
+    private constructor();
+    static fromConfigFile(configPath: string): Promise<AgentRuntime>;
+    static fromConfig(config: PipequeryYaml, cwd: string): Promise<AgentRuntime>;
+    dispose(): Promise<void>;
+}
+
+export { AgentRuntime, type EndpointEntry };

package/dist/runtime.js

@@ -0,0 +1,48 @@
+// src/runtime.ts
+import { dirname, resolve } from "path";
+import { SourceManager } from "@vaultgradient/pipequery-cli/sources";
+
+// src/config.ts
+import { readFile } from "fs/promises";
+import { parse as parseYaml } from "yaml";
+async function loadPipequeryConfig(path) {
+  const raw = await readFile(path, "utf8");
+  const parsed = parseYaml(raw);
+  if (typeof parsed !== "object" || parsed === null) {
+    throw new Error(`pipequery config at ${path} did not parse to an object`);
+  }
+  return parsed;
+}
+
+// src/runtime.ts
+var AgentRuntime = class _AgentRuntime {
+  sources;
+  endpoints;
+  constructor(sources, endpoints) {
+    this.sources = sources;
+    this.endpoints = endpoints;
+  }
+  static async fromConfigFile(configPath) {
+    const absConfig = resolve(configPath);
+    const cwd = dirname(absConfig);
+    const cfg = await loadPipequeryConfig(absConfig);
+    return _AgentRuntime.fromConfig(cfg, cwd);
+  }
+  static async fromConfig(config, cwd) {
+    const sources = new SourceManager(cwd);
+    for (const [name, sourceConfig] of Object.entries(config.sources ?? {})) {
+      await sources.addSource(name, sourceConfig);
+    }
+    const endpoints = /* @__PURE__ */ new Map();
+    for (const [path, ep] of Object.entries(config.endpoints ?? {})) {
+      endpoints.set(path, { path, query: ep.query });
+    }
+    return new _AgentRuntime(sources, endpoints);
+  }
+  async dispose() {
+    await this.sources.dispose();
+  }
+};
+export {
+  AgentRuntime
+};