modelos-secure 0.0.0-stage → 0.1.0

package/hook/clients.mjs

@@ -0,0 +1,270 @@
+// The plain model clients: the prompt canary going in, and what the model asked for coming out.
+// Whatever happens here, the app's own call goes ahead as if the hook were not there.
+
+function safely(work, fallback) {
+  try {
+    return work();
+  } catch {
+    return fallback;
+  }
+}
+
+function parseArgs(text) {
+  if (text && typeof text === "object") return text;
+  const parsed = safely(() => JSON.parse(text), undefined);
+  if (parsed && typeof parsed === "object") return parsed;
+  return parsed === undefined ? { raw: String(text) } : { value: parsed };
+}
+
+/** Adds a line to a message's content, whether it is text or a list of parts. */
+function withLine(content, line) {
+  if (typeof content === "string") return `${content}\n${line}`;
+  if (Array.isArray(content)) return [...content, { type: "text", text: line }];
+  return content;
+}
+
+function canaryInMessages(messages, line) {
+  const index = messages.findIndex(
+    (m) => m?.role === "system" || m?.role === "developer",
+  );
+  if (index === -1) return messages;
+  const copy = [...messages];
+  copy[index] = {
+    ...copy[index],
+    content: withLine(copy[index].content, line),
+  };
+  return copy;
+}
+
+function chatWithCanary(body, line) {
+  if (!line || !Array.isArray(body?.messages)) return body;
+  return { ...body, messages: canaryInMessages(body.messages, line) };
+}
+
+function responsesWithCanary(body, line) {
+  if (!line) return body;
+  if (typeof body?.instructions === "string")
+    return { ...body, instructions: `${body.instructions}\n${line}` };
+  if (Array.isArray(body?.input))
+    return { ...body, input: canaryInMessages(body.input, line) };
+  return body;
+}
+
+function readChat(completion) {
+  const calls = (completion.choices ?? []).flatMap((choice) =>
+    (choice.message?.tool_calls ?? []).flatMap((call) =>
+      call.function
+        ? [
+            {
+              tool: call.function.name,
+              args: parseArgs(call.function.arguments),
+            },
+          ]
+        : [],
+    ),
+  );
+  return {
+    calls,
+    inputTokens: completion.usage?.prompt_tokens,
+    outputTokens: completion.usage?.completion_tokens,
+  };
+}
+
+function chatStreamReader() {
+  const calls = new Map();
+  let usage;
+  return {
+    add(chunk) {
+      usage = chunk.usage ?? usage;
+      for (const choice of chunk.choices ?? []) {
+        for (const delta of choice.delta?.tool_calls ?? []) {
+          const call = calls.get(delta.index) ?? { tool: "", text: "" };
+          call.tool += delta.function?.name ?? "";
+          call.text += delta.function?.arguments ?? "";
+          calls.set(delta.index, call);
+        }
+      }
+    },
+    finish: () => ({
+      calls: [...calls.values()].map((c) => ({
+        tool: c.tool,
+        args: parseArgs(c.text),
+      })),
+      inputTokens: usage?.prompt_tokens,
+      outputTokens: usage?.completion_tokens,
+    }),
+  };
+}
+
+function readResponse(response) {
+  const calls = (response.output ?? [])
+    .filter((item) => item.type === "function_call")
+    .map((item) => ({ tool: item.name, args: parseArgs(item.arguments) }));
+  return {
+    calls,
+    inputTokens: response.usage?.input_tokens,
+    outputTokens: response.usage?.output_tokens,
+  };
+}
+
+function responsesStreamReader() {
+  let found = { calls: [] };
+  return {
+    add(event) {
+      if (event.type === "response.completed")
+        found = readResponse(event.response);
+    },
+    finish: () => found,
+  };
+}
+
+function genAiWithCanary(params, line) {
+  const instruction = params?.config?.systemInstruction;
+  if (!line || !instruction) return params;
+  let changed = instruction;
+  if (typeof instruction === "string") changed = `${instruction}\n${line}`;
+  else if (Array.isArray(instruction)) changed = [...instruction, line];
+  else if (Array.isArray(instruction.parts))
+    changed = { ...instruction, parts: [...instruction.parts, { text: line }] };
+  else if (typeof instruction.text === "string")
+    changed = { ...instruction, text: `${instruction.text}\n${line}` };
+  return {
+    ...params,
+    config: { ...params.config, systemInstruction: changed },
+  };
+}
+
+function readGenAi(response) {
+  const calls = (response.candidates?.[0]?.content?.parts ?? []).flatMap(
+    (part) =>
+      part.functionCall
+        ? [{ tool: part.functionCall.name, args: part.functionCall.args ?? {} }]
+        : [],
+  );
+  return {
+    calls,
+    inputTokens: response.usageMetadata?.promptTokenCount,
+    outputTokens: response.usageMetadata?.candidatesTokenCount,
+  };
+}
+
+function genAiStreamReader() {
+  const calls = [];
+  let inputTokens;
+  let outputTokens;
+  return {
+    add(chunk) {
+      const found = readGenAi(chunk);
+      calls.push(...found.calls);
+      inputTokens = found.inputTokens ?? inputTokens;
+      outputTokens = found.outputTokens ?? outputTokens;
+    },
+    finish: () => ({ calls, inputTokens, outputTokens }),
+  };
+}
+
+export function clientPatches({ canaryLine, seeModelCall }) {
+  const report = (model, found) =>
+    seeModelCall(model, found.calls, {
+      inputTokens: found.inputTokens,
+      outputTokens: found.outputTokens,
+    });
+
+  /** The app reads the stream as usual; what the model asked for is known once it has been read. */
+  function watchStream(stream, model, reader) {
+    // Taken once, because a generator's iterator is the generator itself: replacing its iterator
+    // method and then asking it again would call the replacement for ever.
+    const inner = stream[Symbol.asyncIterator]();
+    const watched = {
+      async next(...args) {
+        const step = await inner.next(...args);
+        if (step.done) safely(() => report(model, reader.finish()));
+        else safely(() => reader.add(step.value));
+        return step;
+      },
+      return: (...args) => inner.return(...args),
+      throw: (...args) => inner.throw(...args),
+      [Symbol.asyncIterator]() {
+        return watched;
+      },
+    };
+    if (typeof stream.next === "function") return watched;
+    stream[Symbol.asyncIterator] = () => watched;
+    return stream;
+  }
+
+  function patchOpenAiCreate(prototype, withCanary, read, streamReader) {
+    if (!prototype?.create) return;
+    const original = prototype.create;
+    prototype.create = function (body, options) {
+      const prepared = safely(() => withCanary(body, canaryLine()), body);
+      const result = original.call(this, prepared, options);
+      const see = (data) => {
+        if (body?.stream && data?.[Symbol.asyncIterator])
+          return safely(
+            () => watchStream(data, body.model, streamReader()),
+            data,
+          );
+        safely(() => report(body?.model, read(data)));
+        return data;
+      };
+      // The client's own promise sends nothing until asked, and has methods of its own, such as
+      // withResponse. Its chaining method keeps all of that, where a plain then would not.
+      if (typeof result?._thenUnwrap === "function")
+        return result._thenUnwrap(see);
+      return typeof result?.then === "function"
+        ? result.then(see)
+        : see(result);
+    };
+  }
+
+  function patchGenAi(Models) {
+    const prototype = Models?.prototype;
+    // generateContent is set up in the constructor. These two are what it calls, once for each
+    // real model call, so automatic function calling is counted call by call.
+    if (prototype?.generateContentInternal) {
+      const original = prototype.generateContentInternal;
+      prototype.generateContentInternal = async function (params) {
+        const prepared = safely(
+          () => genAiWithCanary(params, canaryLine()),
+          params,
+        );
+        const response = await original.call(this, prepared);
+        safely(() => report(params?.model, readGenAi(response)));
+        return response;
+      };
+    }
+    if (prototype?.generateContentStreamInternal) {
+      const original = prototype.generateContentStreamInternal;
+      prototype.generateContentStreamInternal = async function (params) {
+        const prepared = safely(
+          () => genAiWithCanary(params, canaryLine()),
+          params,
+        );
+        const stream = await original.call(this, prepared);
+        return safely(
+          () => watchStream(stream, params?.model, genAiStreamReader()),
+          stream,
+        );
+      };
+    }
+  }
+
+  return {
+    patchOpenAiChat: (Completions) =>
+      patchOpenAiCreate(
+        Completions?.prototype,
+        chatWithCanary,
+        readChat,
+        chatStreamReader,
+      ),
+    patchOpenAiResponses: (Responses) =>
+      patchOpenAiCreate(
+        Responses?.prototype,
+        responsesWithCanary,
+        readResponse,
+        responsesStreamReader,
+      ),
+    patchGenAi,
+  };
+}

package/hook/hook.mjs

@@ -0,0 +1,306 @@
+// The hooks for Node apps, loaded with `node --import modelos-secure/hook`. No dependencies,
+// and written for Node 20: customers run this, not us. Outside a review it does nothing.
+// docs/spec/05-scan-and-cli.md#5-start
+import { createServer } from "node:http";
+import net from "node:net";
+// A default import, because a named one for registerHooks would not even link on Node 20.
+import nodeModule from "node:module";
+import { pathToFileURL, URL } from "node:url";
+import { isMainThread } from "node:worker_threads";
+import { clientPatches } from "./clients.mjs";
+import { rewriteLoaded, rewriteSource, urlToLoad } from "./rewrite.mjs";
+
+const raw = process.env.MODELOS_SECURE_HOOK;
+
+// The loader thread on Node 20 can load this file too, and must not start a second control port.
+if (raw && isMainThread) {
+  const config = JSON.parse(raw);
+  const toolsByName = new Map(config.tools.map((tool) => [tool.name, tool]));
+  const waysIn = new Map(
+    config.waysIn.filter((w) => w.function).map((w) => [w.function, w]),
+  );
+
+  let state = freshState();
+  function freshState() {
+    return {
+      attempt: null,
+      held: [],
+      modelCalls: [],
+      ran: [],
+      plantDelivered: false,
+      retrieved: [],
+    };
+  }
+
+  function argsObject(name, args) {
+    const [first] = args;
+    if (
+      args.length === 1 &&
+      first &&
+      typeof first === "object" &&
+      !Array.isArray(first)
+    )
+      return { ...first };
+    if (
+      args.length >= 1 &&
+      first &&
+      typeof first === "object" &&
+      typeof args[1] === "object" &&
+      "toolCallId" in (args[1] ?? {})
+    )
+      return { ...first };
+    const params = toolsByName.get(name)?.params ?? [];
+    return Object.fromEntries(
+      args
+        .filter((a) => typeof a !== "function")
+        .map((a, i) => [params[i] ?? `arg${i}`, a]),
+    );
+  }
+
+  // A realistic answer, so the app carries on as if the action happened.
+  function heldResult(name) {
+    return {
+      status: "ok",
+      id: `${name}-${Math.floor(Math.random() * 90000) + 10000}`,
+    };
+  }
+
+  function longestStringKey(object) {
+    let best;
+    for (const [key, value] of Object.entries(object)) {
+      if (
+        typeof value === "string" &&
+        (!best || value.length > object[best].length)
+      )
+        best = key;
+    }
+    return best;
+  }
+
+  // Puts text into whatever shape the way in returns, as if it had come from there. Returns the
+  // result itself, so nothing is delivered, when it is a shape with no place for text.
+  function plantInto(result, text) {
+    if (result === undefined || result === null) return text;
+    if (typeof result === "string") return `${result}\n\n${text}`;
+    if (Array.isArray(result)) {
+      if (result.length === 0 || typeof result[0] === "string")
+        return [...result, text];
+      const planted = plantInto(result[0], text);
+      return planted === result[0] ? result : [...result, planted];
+    }
+    if (result && typeof result === "object") {
+      for (const [key, value] of Object.entries(result)) {
+        if (Array.isArray(value))
+          return { ...result, [key]: plantInto(value, text) };
+      }
+      const key = longestStringKey(result);
+      return key
+        ? { ...result, [key]: `${result[key]}\n\n${text}` }
+        : { ...result, note: text };
+    }
+    return result;
+  }
+
+  function addCanary(result, record) {
+    if (typeof result === "string")
+      return `${result}\nCustomer reference: ${record}`;
+    if (result && typeof result === "object" && !Array.isArray(result))
+      return { ...result, customer_reference: record };
+    return result;
+  }
+
+  // Q2 grades answers against what was really read, so this is the text and not the way in's name.
+  function passageText(result) {
+    const text = typeof result === "string" ? result : JSON.stringify(result);
+    return String(text).slice(0, 8000);
+  }
+
+  function afterRun(name, result) {
+    const wayIn = waysIn.get(name);
+    const attempt = state.attempt;
+    let out = result;
+    if (attempt && wayIn) state.retrieved.push(passageText(result));
+    if (attempt?.plant && wayIn && attempt.plant.wayIn === wayIn.name) {
+      out = plantInto(out, attempt.plant.text);
+      state.plantDelivered = out !== result;
+    }
+    if (attempt && toolsByName.get(name)?.can === "read private data")
+      out = addCanary(out, attempt.canaries.record);
+    return out;
+  }
+
+  /** The one place a tool call passes through. Held tools never run. */
+  function wrap(name, original, self, args) {
+    const tool = toolsByName.get(name);
+    if (!tool && !waysIn.has(name)) return original.apply(self, args);
+    if (tool?.mode === "hold") {
+      state.held.push({ tool: name, args: argsObject(name, args) });
+      return heldResult(name);
+    }
+    if (tool) state.ran.push(name);
+    const result = original.apply(self, args);
+    if (result && typeof result.then === "function")
+      return result.then((value) => afterRun(name, value));
+    return afterRun(name, result);
+  }
+
+  function canaryLine() {
+    const canary = state.attempt?.canaries.prompt;
+    return canary ? `Internal reference: ${canary}` : undefined;
+  }
+
+  function promptWithCanary(system) {
+    const line = canaryLine();
+    if (line && typeof system === "string") return `${system}\n${line}`;
+    return system;
+  }
+
+  function seeModelCall(model, toolCalls, usage) {
+    state.modelCalls.push({
+      model,
+      toolCalls,
+      inputTokens: usage?.inputTokens ?? 0,
+      outputTokens: usage?.outputTokens ?? 0,
+    });
+  }
+
+  // The Vercel AI SDK: tools carry their own execute, and steps say what the model asked for.
+  function wrapAiFunction(original) {
+    return function modelosSecureAi(options) {
+      const tools = Object.fromEntries(
+        Object.entries(options.tools ?? {}).map(([name, definition]) => {
+          if (!toolsByName.has(name) || !definition.execute)
+            return [name, definition];
+          return [
+            name,
+            {
+              ...definition,
+              execute: (...args) =>
+                wrap(name, definition.execute, definition, args),
+            },
+          ];
+        }),
+      );
+      const result = original({
+        ...options,
+        tools,
+        system: promptWithCanary(options.system),
+      });
+      if (result && typeof result.then === "function") {
+        return result.then((value) => {
+          for (const step of value.steps ?? []) {
+            const calls = (step.toolCalls ?? []).map((c) => ({
+              tool: c.toolName,
+              args: c.input ?? c.args ?? {},
+            }));
+            seeModelCall(options.model?.modelId, calls, step.usage);
+          }
+          return value;
+        });
+      }
+      return result;
+    };
+  }
+
+  globalThis[Symbol.for("modelos-secure.hook")] = {
+    wrap,
+    wrapAiFunction,
+    seeModelCall,
+    state: () => state,
+    ...clientPatches({ canaryLine, seeModelCall }),
+  };
+
+  // Source rewriting, so nothing in the app's code changes.
+  const options = {
+    held: config.tools
+      .map((t) => t.name)
+      .concat(config.waysIn.map((w) => w.function).filter(Boolean)),
+    appFiles: config.files ?? [],
+  };
+  if (nodeModule.registerHooks) {
+    nodeModule.registerHooks({
+      load: (url, context, nextLoad) =>
+        rewriteLoaded(url, nextLoad(urlToLoad(url), context), options),
+    });
+  } else {
+    // Node 20: the same rewrite, through an off-thread loader for import, which never sees
+    // require, and through the compile step for require. Without that, a CommonJS tool would run.
+    nodeModule.register(new URL("./loader.mjs", import.meta.url).href, {
+      data: options,
+    });
+    const compile = nodeModule.prototype._compile;
+    nodeModule.prototype._compile = function (source, filename, ...rest) {
+      const url = pathToFileURL(filename).href;
+      const changed = rewriteSource(url, source, options);
+      return compile.call(this, changed, filename, ...rest);
+    };
+  }
+
+  // The control port the runner talks to. Loopback only, and it dies with the app.
+  function startControl() {
+    const server = createServer(async (request, response) => {
+      const chunks = [];
+      for await (const chunk of request) chunks.push(chunk);
+      const send = (body) => {
+        response.writeHead(200, { "content-type": "application/json" });
+        response.end(JSON.stringify(body));
+      };
+      if (request.url === "/begin") {
+        state = freshState();
+        state.attempt = JSON.parse(Buffer.concat(chunks).toString("utf8"));
+        return send({ ok: true });
+      }
+      if (request.url === "/end") {
+        const {
+          held: heldCalls,
+          modelCalls,
+          ran,
+          plantDelivered,
+          retrieved,
+        } = state;
+        state = freshState();
+        return send({
+          heldCalls,
+          modelCalls,
+          ranTools: ran,
+          plantDelivered,
+          retrieved,
+        });
+      }
+      if (request.url === "/hello") return send({ ok: true, hook: "node" });
+      response.writeHead(404);
+      response.end();
+    });
+    server.listen(0, "127.0.0.1", () => {
+      const { port } = server.address();
+      server.unref();
+      fetch(`${process.env.MODELOS_SECURE_RUNNER_URL}/ready`, {
+        method: "POST",
+        headers: { "content-type": "application/json" },
+        body: JSON.stringify({
+          instance: process.env.MODELOS_SECURE_INSTANCE,
+          port,
+        }),
+      }).catch(() => undefined);
+    });
+  }
+
+  // Only the process that takes the app's port is the app. NODE_OPTIONS reaches every Node process
+  // the app starts, and npm starting the app is one too: see
+  // docs/journal/adr/0020-announce-on-the-app-port.md
+  const appPort = Number(process.env.PORT);
+  if (appPort) {
+    const listen = net.Server.prototype.listen;
+    let claimed = false;
+    net.Server.prototype.listen = function (...args) {
+      const port = typeof args[0] === "object" ? args[0]?.port : args[0];
+      if (!claimed && Number(port) === appPort) {
+        claimed = true;
+        startControl();
+      }
+      return listen.apply(this, args);
+    };
+  } else {
+    startControl();
+  }
+}

package/hook/loader.mjs

@@ -0,0 +1,12 @@
+// The off-thread loader for Node versions without registerHooks, registered from hook.mjs.
+import { rewriteLoaded, urlToLoad } from "./rewrite.mjs";
+
+let options = { held: [], appFiles: [] };
+
+export function initialize(data) {
+  options = data;
+}
+
+export async function load(url, context, nextLoad) {
+  return rewriteLoaded(url, await nextLoad(urlToLoad(url), context), options);
+}

package/hook/rewrite.mjs

@@ -0,0 +1,95 @@
+// Source rewriting, shared by both ways of loading the hooks: registerHooks on the main thread,
+// and the Node 20 off-thread loader. It runs in either place, so it holds no state of its own and
+// what it writes only reaches the hook through a global, on the thread that runs the module.
+import { fileURLToPath } from "node:url";
+
+const HOOK = `globalThis[Symbol.for("modelos-secure.hook")]`;
+const MARK = "?modelos-secure-original";
+const PACKAGE = `/node_modules/(?:\\.pnpm/[^/]+/node_modules/)?`;
+
+function shimAi(url) {
+  const real = `${url}${MARK}`;
+  return [
+    `import * as real from ${JSON.stringify(real)};`,
+    `export * from ${JSON.stringify(real)};`,
+    `const hook = ${HOOK};`,
+    `export const generateText = hook.wrapAiFunction(real.generateText);`,
+    `export const streamText = hook.wrapAiFunction(real.streamText);`,
+  ].join("\n");
+}
+
+const SHIMS = [[new RegExp(`${PACKAGE}ai/dist/index\\.js$`), shimAi]];
+
+// The plain clients keep their exports exactly as they are: one line at the end patches the class.
+const PATCHES = [
+  [
+    new RegExp(
+      `${PACKAGE}openai/resources/chat/completions/completions\\.m?js$`,
+    ),
+    "patchOpenAiChat(Completions)",
+  ],
+  [
+    new RegExp(`${PACKAGE}openai/resources/responses/responses\\.m?js$`),
+    "patchOpenAiResponses(Responses)",
+  ],
+  [
+    new RegExp(`${PACKAGE}@google/genai/dist/(?:node/)?index\\.(?:mjs|cjs)$`),
+    "patchGenAi(Models)",
+  ],
+];
+
+/** Plain tool functions become a wrapper around the original, under the same name. */
+function rewriteFunctions(source, names) {
+  let out = source;
+  for (const name of new Set(names)) {
+    const inner = `__modelos_secure_${name}`;
+    const call = (self) =>
+      `${HOOK}.wrap(${JSON.stringify(name)}, ${inner}, ${self}, a)`;
+    // function name( and async function name<T>(
+    out = out.replace(
+      new RegExp(
+        `(^|\\n)(\\s*)(export\\s+)?(async\\s+)?function\\s+${name}\\s*(?=[(<])`,
+        "g",
+      ),
+      (_, start, indent, exported, isAsync) =>
+        `${start}${indent}${exported ?? ""}function ${name}(...a) { return ${call("this")}; }\n` +
+        `${indent}${isAsync ?? ""}function ${inner}`,
+    );
+    // const name = async (...) => and const name: Type = function (...), but not const name = (x)
+    out = out.replace(
+      new RegExp(
+        `(^|\\n)(\\s*)(export\\s+)?(?:const|let|var)\\s+${name}\\s*(:[^=\\n]+)?=\\s*(?=(?:async\\s+)?(?:function\\b|\\w+\\s*=>|<|\\((?:[^()]|\\([^()]*\\))*\\)\\s*(?::[^=]+)?=>))`,
+        "g",
+      ),
+      (_, start, indent, exported, type) =>
+        `${start}${indent}${exported ?? ""}const ${name} = (...a) => ${call("undefined")};\n` +
+        `${indent}const ${inner}${type ?? ""} = `,
+    );
+  }
+  return out;
+}
+
+/** The URL to really load: the shims import the module they stand in for under a marked URL. */
+export const urlToLoad = (url) => url.replace(MARK, "");
+
+/** The source to compile: changed for the files that matter, and otherwise exactly as it came. */
+export function rewriteSource(url, source, options) {
+  for (const [pattern, patch] of PATCHES) {
+    if (pattern.test(url)) return `${String(source)}\n;${HOOK}.${patch};\n`;
+  }
+  if (url.startsWith("file:") && options.appFiles.includes(fileURLToPath(url)))
+    return rewriteFunctions(String(source), options.held);
+  return source;
+}
+
+/** What should load in place of what Node found. options: { held: string[], appFiles: string[] } */
+export function rewriteLoaded(url, loaded, options) {
+  if (url.includes(MARK)) return loaded;
+  for (const [pattern, shim] of SHIMS) {
+    if (pattern.test(url))
+      return { ...loaded, source: shim(url), format: "module" };
+  }
+  if (!loaded.source) return loaded;
+  const source = rewriteSource(url, loaded.source, options);
+  return source === loaded.source ? loaded : { ...loaded, source };
+}