modelos-secure 0.1.1 → 0.2.0

package/hook/clients.mjs

@@ -1,6 +1,8 @@
 // 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.
 
+import { sdkContext } from "./edge.mjs";
+
 function safely(work, fallback) {
   try {
     return work();
@@ -198,7 +200,9 @@ export function clientPatches({ canaryLine, seeModelCall }) {
     const original = prototype.create;
     prototype.create = function (body, options) {
       const prepared = safely(() => withCanary(body, canaryLine()), body);
-      const result = original.call(this, prepared, options);
+      const result = sdkContext.run(true, () =>
+        original.call(this, prepared, options),
+      );
       const see = (data) => {
         if (body?.stream && data?.[Symbol.asyncIterator])
           return safely(
@@ -229,7 +233,9 @@ export function clientPatches({ canaryLine, seeModelCall }) {
           () => genAiWithCanary(params, canaryLine()),
           params,
         );
-        const response = await original.call(this, prepared);
+        const response = await sdkContext.run(true, () =>
+          original.call(this, prepared),
+        );
         safely(() => report(params?.model, readGenAi(response)));
         return response;
       };
@@ -241,7 +247,9 @@ export function clientPatches({ canaryLine, seeModelCall }) {
           () => genAiWithCanary(params, canaryLine()),
           params,
         );
-        const stream = await original.call(this, prepared);
+        const stream = await sdkContext.run(true, () =>
+          original.call(this, prepared),
+        );
         return safely(
           () => watchStream(stream, params?.model, genAiStreamReader()),
           stream,

package/hook/edge.mjs

@@ -0,0 +1,405 @@
+// The edge: every framework ends at an HTTP call to a model, so the hook watches there too, below
+// whatever framework the app uses. During an attempt it sees each model call (the tools offered, the
+// tool calls the model makes and their arguments), adds the prompt canary, plants attacker text into
+// a tool result, and stops every other call that leaves the machine, answering it as if it worked and
+// recording it. A tool known to run, such as a read against a staging copy, keeps its calls. Outside
+// an attempt it does nothing. No dependencies, Node 20 and later. ADR 0062.
+import { AsyncLocalStorage } from "node:async_hooks";
+import http from "node:http";
+import https from "node:https";
+import { URL } from "node:url";
+import { allowedHost } from "./hosts.mjs";
+
+/** Which tool is running, so its own calls can be told apart from calls nobody declared. */
+export const toolContext = new AsyncLocalStorage();
+/** Set while a model client the hook already patches is calling, so the call is counted once. */
+export const sdkContext = new AsyncLocalStorage();
+
+const MODEL_PATHS = [
+  [/\/chat\/completions$/, "chat"],
+  [/\/responses$/, "responses"],
+  [/\/messages$/, "anthropic"],
+  [/:(stream)?[gG]enerateContent$/, "gemini"],
+  [/\/api\/chat$/, "ollama"],
+];
+
+function safely(work, fallback) {
+  try {
+    return work();
+  } catch {
+    return fallback;
+  }
+}
+
+const parseArgs = (value) =>
+  value && typeof value === "object"
+    ? value
+    : safely(() => JSON.parse(value), { raw: String(value ?? "") });
+
+/** The API a request speaks, from its path and body, or null when it is not a model call. */
+export function modelShape(url, body) {
+  if (!body || typeof body !== "object") return null;
+  for (const [pattern, shape] of MODEL_PATHS)
+    if (pattern.test(url.pathname)) return shape;
+  if (Array.isArray(body.messages) && typeof body.model === "string")
+    return "chat";
+  if (Array.isArray(body.contents)) return "gemini";
+  return null;
+}
+
+const withLine = (content, line) =>
+  typeof content === "string"
+    ? `${content}\n${line}`
+    : Array.isArray(content)
+      ? [...content, { type: "text", text: line }]
+      : content;
+
+/** The canary goes into the system prompt, whatever the API calls it, once. */
+function addCanary(shape, body, line) {
+  if (!line || JSON.stringify(body).includes(line)) return body;
+  if (shape === "chat" || shape === "ollama") {
+    const i = body.messages.findIndex(
+      (m) => m?.role === "system" || m?.role === "developer",
+    );
+    if (i === -1) return body;
+    const messages = [...body.messages];
+    messages[i] = {
+      ...messages[i],
+      content: withLine(messages[i].content, line),
+    };
+    return { ...body, messages };
+  }
+  if (shape === "responses" && typeof body.instructions === "string")
+    return { ...body, instructions: `${body.instructions}\n${line}` };
+  if (shape === "anthropic" && body.system)
+    return {
+      ...body,
+      system:
+        typeof body.system === "string"
+          ? `${body.system}\n${line}`
+          : [...body.system, { type: "text", text: line }],
+    };
+  if (shape === "gemini") {
+    const instruction = body.systemInstruction ?? body.system_instruction;
+    if (!instruction?.parts) return body;
+    return {
+      ...body,
+      systemInstruction: {
+        ...instruction,
+        parts: [...instruction.parts, { text: line }],
+      },
+    };
+  }
+  return body;
+}
+
+/** The tool results in a request, with the tool each answers, so a plant lands where it reads. */
+function plantIntoResults(shape, body, tool, text) {
+  let delivered = false;
+  const named = new Map();
+  const plantText = (value) => {
+    delivered = true;
+    return typeof value === "string"
+      ? `${value}\n\n${text}`
+      : Array.isArray(value)
+        ? [...value, { type: "text", text }]
+        : { ...value, note: text };
+  };
+  if (shape === "chat" || shape === "ollama") {
+    for (const m of body.messages)
+      for (const call of m?.tool_calls ?? [])
+        named.set(call.id, call.function?.name);
+    const messages = body.messages.map((m) =>
+      m?.role === "tool" &&
+      (named.get(m.tool_call_id) === tool || m.name === tool)
+        ? { ...m, content: plantText(m.content) }
+        : m,
+    );
+    return { body: { ...body, messages }, delivered };
+  }
+  if (shape === "anthropic") {
+    for (const m of body.messages ?? [])
+      for (const part of Array.isArray(m.content) ? m.content : [])
+        if (part.type === "tool_use") named.set(part.id, part.name);
+    const messages = (body.messages ?? []).map((m) =>
+      Array.isArray(m.content)
+        ? {
+            ...m,
+            content: m.content.map((part) =>
+              part.type === "tool_result" &&
+              named.get(part.tool_use_id) === tool
+                ? { ...part, content: plantText(part.content) }
+                : part,
+            ),
+          }
+        : m,
+    );
+    return { body: { ...body, messages }, delivered };
+  }
+  if (shape === "gemini") {
+    const contents = body.contents.map((c) => ({
+      ...c,
+      parts: (c.parts ?? []).map((p) =>
+        p.functionResponse?.name === tool
+          ? {
+              ...p,
+              functionResponse: {
+                ...p.functionResponse,
+                response: plantText(p.functionResponse.response ?? {}),
+              },
+            }
+          : p,
+      ),
+    }));
+    return { body: { ...body, contents }, delivered };
+  }
+  if (shape === "responses" && Array.isArray(body.input)) {
+    for (const item of body.input)
+      if (item.type === "function_call") named.set(item.call_id, item.name);
+    const input = body.input.map((item) =>
+      item.type === "function_call_output" && named.get(item.call_id) === tool
+        ? { ...item, output: plantText(item.output) }
+        : item,
+    );
+    return { body: { ...body, input }, delivered };
+  }
+  return { body, delivered };
+}
+
+/** The tools offered and the calls asked for, in any of the APIs above. */
+export function readModel(shape, body, answer) {
+  const offered = [
+    ...(body.tools ?? []).map((t) => t.function?.name ?? t.name),
+    ...(body.tools ?? []).flatMap((t) =>
+      (t.functionDeclarations ?? t.function_declarations ?? []).map(
+        (d) => d.name,
+      ),
+    ),
+  ].filter(Boolean);
+  const calls = [];
+  for (const choice of answer.choices ?? [])
+    for (const call of choice.message?.tool_calls ?? [])
+      calls.push({
+        tool: call.function?.name,
+        args: parseArgs(call.function?.arguments),
+      });
+  for (const call of answer.message?.tool_calls ?? [])
+    calls.push({
+      tool: call.function?.name,
+      args: parseArgs(call.function?.arguments),
+    });
+  for (const item of answer.output ?? [])
+    if (item.type === "function_call")
+      calls.push({ tool: item.name, args: parseArgs(item.arguments) });
+  for (const part of Array.isArray(answer.content) ? answer.content : [])
+    if (part.type === "tool_use")
+      calls.push({ tool: part.name, args: part.input ?? {} });
+  for (const candidate of answer.candidates ?? [])
+    for (const part of candidate.content?.parts ?? [])
+      if (part.functionCall)
+        calls.push({
+          tool: part.functionCall.name,
+          args: part.functionCall.args ?? {},
+        });
+  const usage = answer.usage ?? answer.usageMetadata ?? {};
+  return {
+    model: body.model ?? answer.model,
+    offered,
+    calls: calls.filter((c) => c.tool),
+    inputTokens:
+      usage.prompt_tokens ?? usage.input_tokens ?? usage.promptTokenCount ?? 0,
+    outputTokens:
+      usage.completion_tokens ??
+      usage.output_tokens ??
+      usage.candidatesTokenCount ??
+      0,
+  };
+}
+
+/** A streamed answer, read event by event into the same shape as a whole one. */
+export function readStream(shape, body, text) {
+  const events = text
+    .split("\n")
+    .filter((l) => l.startsWith("data:"))
+    .map((l) => safely(() => JSON.parse(l.slice(5).trim()), null))
+    .filter(Boolean);
+  const pieces = new Map();
+  const whole = { choices: [], content: [], candidates: [], output: [] };
+  for (const event of events) {
+    for (const choice of event.choices ?? [])
+      for (const delta of choice.delta?.tool_calls ?? []) {
+        const piece = pieces.get(`c${delta.index}`) ?? { name: "", args: "" };
+        piece.name += delta.function?.name ?? "";
+        piece.args += delta.function?.arguments ?? "";
+        pieces.set(`c${delta.index}`, piece);
+      }
+    if (
+      event.type === "content_block_start" &&
+      event.content_block?.type === "tool_use"
+    )
+      pieces.set(`a${event.index}`, {
+        name: event.content_block.name,
+        args: "",
+      });
+    if (event.delta?.type === "input_json_delta") {
+      const piece = pieces.get(`a${event.index}`);
+      if (piece) piece.args += event.delta.partial_json ?? "";
+    }
+    if (event.candidates) whole.candidates.push(...event.candidates);
+    if (event.type === "response.completed")
+      whole.output = event.response?.output ?? [];
+    if (event.usage || event.usageMetadata)
+      whole.usage = event.usage ?? event.usageMetadata;
+  }
+  for (const piece of pieces.values())
+    whole.choices.push({
+      message: {
+        tool_calls: [{ function: { name: piece.name, arguments: piece.args } }],
+      },
+    });
+  return readModel(shape, body, whole);
+}
+
+/**
+ * Installs the edge. `hook` gives the attempt's state, the config's tools and ways in, what the
+ * person allowed through, and where a held call is recorded.
+ */
+export function installEdge(hook) {
+  const allowed = (host) => allowedHost(host, hook.allow);
+  /** Inside a tool the person said may run, its calls go through, to a staging copy. */
+  const insideRunningTool = () => hook.mode(toolContext.getStore()) === "run";
+
+  function hold(method, url, bytes) {
+    hook
+      .state()
+      .edge.push({ method, host: url.hostname, path: url.pathname, bytes });
+  }
+
+  function seeModel(shape, body, answer, streamed, counted) {
+    const read = streamed
+      ? readStream(shape, body, answer)
+      : readModel(shape, body, answer);
+    const state = hook.state();
+    if (!counted)
+      state.modelCalls.push({
+        model: read.model,
+        toolCalls: read.calls,
+        offered: read.offered,
+        inputTokens: read.inputTokens,
+        outputTokens: read.outputTokens,
+        edge: true,
+      });
+    // A call to a tool marked held that no function-level hook caught: what the model asked for
+    // is the evidence, and the edge stops what the tool does.
+    for (const call of read.calls)
+      if (
+        hook.mode(call.tool) === "hold" &&
+        !state.held.some((h) => h.tool === call.tool)
+      )
+        state.held.push({ tool: call.tool, args: call.args, at: "edge" });
+  }
+
+  /** What a model call carries in, during an attempt: the canary, and a plant at a tool result. */
+  function prepare(shape, body) {
+    const state = hook.state();
+    let prepared = addCanary(shape, body, hook.canaryLine());
+    const plant = state.attempt?.plant;
+    const way = plant && hook.wayIn(plant.wayIn);
+    if (plant && way?.function && !state.plantDelivered) {
+      const planted = plantIntoResults(
+        shape,
+        prepared,
+        way.function,
+        plant.text,
+      );
+      prepared = planted.body;
+      if (planted.delivered) state.plantDelivered = true;
+    }
+    return prepared;
+  }
+
+  const original = globalThis.fetch;
+  if (original)
+    globalThis.fetch = async function modelosSecureFetch(input, init) {
+      if (!hook.state().attempt) return original(input, init);
+      const url = new URL(
+        typeof input === "string" || input instanceof URL ? input : input.url,
+      );
+      const text = typeof init?.body === "string" ? init.body : undefined;
+      const body = text ? safely(() => JSON.parse(text), null) : null;
+      const shape = modelShape(url, body);
+      if (shape) {
+        const counted = Boolean(sdkContext.getStore());
+        const response = await original(input, {
+          ...init,
+          body: JSON.stringify(prepare(shape, body)),
+        });
+        const streamed =
+          body.stream === true || /stream/i.test(url.pathname + url.search);
+        void response
+          .clone()
+          .text()
+          .then((answer) =>
+            safely(() =>
+              seeModel(
+                shape,
+                body,
+                streamed ? answer : JSON.parse(answer),
+                streamed,
+                counted,
+              ),
+            ),
+          )
+          .catch(() => undefined);
+        return response;
+      }
+      if (allowed(url.hostname) || insideRunningTool())
+        return original(input, init);
+      hold(init?.method ?? "GET", url, text?.length ?? 0);
+      return new globalThis.Response(
+        JSON.stringify({ status: "ok", id: `held-${Date.now()}` }),
+        {
+          status: 200,
+          headers: { "content-type": "application/json" },
+        },
+      );
+    };
+
+  // Clients built on http and https, such as axios and older SDKs: a call that would leave goes to
+  // the hook's own port instead, which answers as if it worked.
+  for (const module of [http, https]) {
+    for (const name of ["request", "get"]) {
+      const real = module[name];
+      module[name] = function modelosSecureRequest(...args) {
+        const state = hook.state();
+        const port = hook.controlPort();
+        if (!state.attempt || !port) return real.apply(this, args);
+        const target = safely(() => requestUrl(module, args), null);
+        if (!target || allowed(target.hostname) || insideRunningTool())
+          return real.apply(this, args);
+        hold(requestMethod(args), target, 0);
+        return http[name](
+          `http://127.0.0.1:${port}/held`,
+          ...args.filter((a) => typeof a === "function"),
+        );
+      };
+    }
+  }
+}
+
+function requestUrl(module, args) {
+  const [first, second] = args;
+  if (typeof first === "string" || first instanceof URL) return new URL(first);
+  const options = first ?? second ?? {};
+  const protocol = options.protocol ?? (module === https ? "https:" : "http:");
+  return new URL(
+    `${protocol}//${options.hostname ?? options.host ?? "localhost"}${options.path ?? "/"}`,
+  );
+}
+
+function requestMethod(args) {
+  const options = args.find(
+    (a) => a && typeof a === "object" && !(a instanceof URL),
+  );
+  return options?.method ?? "GET";
+}

package/hook/hook.mjs

@@ -8,6 +8,7 @@ import nodeModule from "node:module";
 import { pathToFileURL, URL } from "node:url";
 import { isMainThread } from "node:worker_threads";
 import { clientPatches } from "./clients.mjs";
+import { installEdge, sdkContext, toolContext } from "./edge.mjs";
 import { rewriteLoaded, rewriteSource, urlToLoad } from "./rewrite.mjs";
 
 const raw = process.env.MODELOS_SECURE_HOOK;
@@ -19,6 +20,8 @@ if (raw && isMainThread) {
   const waysIn = new Map(
     config.waysIn.filter((w) => w.function).map((w) => [w.function, w]),
   );
+  const waysByName = new Map(config.waysIn.map((w) => [w.name, w]));
+  let controlPort = 0;
 
   let state = freshState();
   function freshState() {
@@ -29,6 +32,7 @@ if (raw && isMainThread) {
       ran: [],
       plantDelivered: false,
       retrieved: [],
+      edge: [],
     };
   }
 
@@ -138,7 +142,7 @@ if (raw && isMainThread) {
       return heldResult(name);
     }
     if (tool) state.ran.push(name);
-    const result = original.apply(self, args);
+    const result = toolContext.run(name, () => original.apply(self, args));
     if (result && typeof result.then === "function")
       return result.then((value) => afterRun(name, value));
     return afterRun(name, result);
@@ -181,11 +185,13 @@ if (raw && isMainThread) {
           ];
         }),
       );
-      const result = original({
-        ...options,
-        tools,
-        system: promptWithCanary(options.system),
-      });
+      const result = sdkContext.run(true, () =>
+        original({
+          ...options,
+          tools,
+          system: promptWithCanary(options.system),
+        }),
+      );
       if (result && typeof result.then === "function") {
         return result.then((value) => {
           for (const step of value.steps ?? []) {
@@ -210,6 +216,15 @@ if (raw && isMainThread) {
     ...clientPatches({ canaryLine, seeModelCall }),
   };
 
+  installEdge({
+    state: () => state,
+    canaryLine,
+    mode: (name) => toolsByName.get(name)?.mode,
+    wayIn: (name) => waysByName.get(name),
+    allow: config.edge?.allow ?? [],
+    controlPort: () => controlPort,
+  });
+
   // Source rewriting, so nothing in the app's code changes.
   const options = {
     held: config.tools
@@ -257,22 +272,35 @@ if (raw && isMainThread) {
           ran,
           plantDelivered,
           retrieved,
+          edge,
         } = state;
         state = freshState();
+        // The edge sees a tool call before its function runs. Where a hook then held the function
+        // itself, that record is the one kept, so a call counts once.
+        const byHook = new Set(
+          heldCalls.filter((h) => h.at !== "edge").map((h) => h.tool),
+        );
         return send({
-          heldCalls,
+          heldCalls: heldCalls.filter(
+            (h) => h.at !== "edge" || !byHook.has(h.tool),
+          ),
           modelCalls,
           ranTools: ran,
           plantDelivered,
           retrieved,
+          edgeHeld: edge,
         });
       }
+      // A call the edge stopped lands here, and is answered as if it worked.
+      if (request.url === "/held")
+        return send({ status: "ok", id: `held-${Date.now().toString()}` });
       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();
+      controlPort = port;
       server.unref();
       fetch(`${process.env.MODELOS_SECURE_RUNNER_URL}/ready`, {
         method: "POST",

package/hook/hosts.d.mts

@@ -0,0 +1,2 @@
+/** Whether a call to this host may leave during an attempt. */
+export function allowedHost(host: string, allow?: string[]): boolean;

package/hook/hosts.mjs

@@ -0,0 +1,26 @@
+// The hosts a review lets an app reach: its own machine, model providers and the sign-in hosts
+// their clients need, and what the person allowed in modelos-secure.yaml. Shared by the in-process
+// edge and the runner's proxy, so both stop the same calls. ADR 0062.
+
+const MODEL_HOSTS = [
+  /(^|\.)api\.openai\.com$/,
+  /\.openai\.azure\.com$/,
+  /(^|\.)api\.anthropic\.com$/,
+  /(^|\.)generativelanguage\.googleapis\.com$/,
+  /aiplatform\.googleapis\.com$/,
+  /^(oauth2|sts|iamcredentials)\.googleapis\.com$/,
+  /^bedrock-runtime\.[\w-]+\.amazonaws\.com$/,
+  /^(api\.mistral\.ai|api\.groq\.com|api\.cohere\.com|openrouter\.ai|api\.together\.xyz|api\.deepseek\.com|api\.x\.ai|api\.fireworks\.ai|router\.huggingface\.co|api-inference\.huggingface\.co)$/,
+  /^login\.microsoftonline\.com$/,
+];
+const LOOPBACK = /^(localhost|127\.\d+\.\d+\.\d+|\[?::1\]?|0\.0\.0\.0)$/;
+
+/** Whether a call to this host may leave during an attempt. */
+export function allowedHost(host, allow = []) {
+  const name = String(host).toLowerCase();
+  return (
+    LOOPBACK.test(name) ||
+    MODEL_HOSTS.some((p) => p.test(name)) ||
+    allow.some((a) => name === a || name.endsWith(`.${a}`))
+  );
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "modelos-secure",
-  "version": "0.1.1",
+  "version": "0.2.0",
   "description": "Tests AI systems for Security, Safety and Quality. One command connects a system from its code.",
   "license": "UNLICENSED",
   "repository": {

package/python/modelos_secure/clients.py

@@ -3,6 +3,7 @@
 Whatever happens here, the app's own call goes ahead as if the hook were not there.
 """
 
+from modelos_secure import edge
 import contextvars
 import functools
 import inspect
@@ -142,7 +143,11 @@ def _watch_call(original, prepare, read, stream_reader=None, always_streams=Fals
         @functools.wraps(original)
         async def wrapper(*args, **kwargs):
             args, kwargs, watching = before(args, kwargs)
-            result = await original(*args, **kwargs)
+            counted = edge.in_sdk.set(True)
+            try:
+                result = await original(*args, **kwargs)
+            finally:
+                edge.in_sdk.reset(counted)
             return after(args, kwargs, result) if watching else result
 
     else:
@@ -150,7 +155,11 @@ def _watch_call(original, prepare, read, stream_reader=None, always_streams=Fals
         @functools.wraps(original)
         def wrapper(*args, **kwargs):
             args, kwargs, watching = before(args, kwargs)
-            result = original(*args, **kwargs)
+            counted = edge.in_sdk.set(True)
+            try:
+                result = original(*args, **kwargs)
+            finally:
+                edge.in_sdk.reset(counted)
             if not watching:
                 return result
             if inspect.isawaitable(result):
@@ -406,9 +415,11 @@ def _patch_langchain(module):
         def wrapper(self, messages, *args, **kwargs):
             messages = _safely(_prepare_langchain_messages, messages, fallback=messages)
             token = _inside_framework.set(True)
+            counted = edge.in_sdk.set(True)
             try:
                 result = original(self, messages, *args, **kwargs)
             finally:
+                edge.in_sdk.reset(counted)
                 _inside_framework.reset(token)
             _safely(_read_langchain, _langchain_model_name(self), result)
             return result
@@ -421,9 +432,11 @@ def _patch_langchain(module):
         async def wrapper(self, messages, *args, **kwargs):
             messages = _safely(_prepare_langchain_messages, messages, fallback=messages)
             token = _inside_framework.set(True)
+            counted = edge.in_sdk.set(True)
             try:
                 result = await original(self, messages, *args, **kwargs)
             finally:
+                edge.in_sdk.reset(counted)
                 _inside_framework.reset(token)
             _safely(_read_langchain, _langchain_model_name(self), result)
             return result