Protect AI SDK tool usage

.github/workflows/end-to-end-tests.yml

@@ -66,3 +66,5 @@ jobs:
       - run: npm install
       - run: npm run build
       - run: npm run end2end
+        env:
+          GOOGLE_GENERATIVE_AI_API_KEY: ${{ secrets.GOOGLE_GENERATIVE_AI_API_KEY }}

.github/workflows/unit-test.yml

@@ -80,6 +80,8 @@ jobs:
       - run: npm run install-lib-only
       - run: npm run build
       - run: npm run test:ci
+        env:
+          GOOGLE_GENERATIVE_AI_API_KEY: ${{ secrets.GOOGLE_GENERATIVE_AI_API_KEY }}
       - name: "Upload coverage"
         uses: codecov/codecov-action@0565863a31f2c772f9f0395002a31e3f06189574 # v5
         with:

README.md

@@ -98,6 +98,8 @@ See list above for supported database drivers.
 
 * ✅ [`@koa/router`](https://www.npmjs.com/package/@koa/router) 13.x, 12.x, 11.x and 10.x
 
+### AI SDKs
+* ✅ [`ai`](https://www.npmjs.com/package/ai) 4.x
 
 ## Installation
 

end2end/tests/hono-sqlite-ai.test.js

@@ -0,0 +1,114 @@
+const t = require("tap");
+const { spawn } = require("child_process");
+const { resolve } = require("path");
+const timeout = require("../timeout");
+
+const pathToApp = resolve(
+  __dirname,
+  "../../sample-apps/hono-sqlite-ai",
+  "app.js"
+);
+
+t.test("it blocks in blocking mode", (t) => {
+  const server = spawn(`node`, [pathToApp, "4006"], {
+    env: { ...process.env, AIKIDO_DEBUG: "true", AIKIDO_BLOCK: "true" },
+  });
+
+  server.on("close", () => {
+    t.end();
+  });
+
+  server.on("error", (err) => {
+    t.fail(err.message);
+  });
+
+  let stdout = "";
+  server.stdout.on("data", (data) => {
+    stdout += data.toString();
+  });
+
+  let stderr = "";
+  server.stderr.on("data", (data) => {
+    stderr += data.toString();
+  });
+
+  // Wait for the server to start
+  timeout(2000)
+    .then(() => {
+      return Promise.all([
+        fetch(
+          `http://127.0.0.1:4006/weather?prompt=${encodeURIComponent('What is the weather in "Ghent\'; DELETE FROM weather; --" like?')}`,
+          {
+            signal: AbortSignal.timeout(5000),
+          }
+        ),
+      ]);
+    })
+    .then(async ([sqlInjection]) => {
+      t.equal(sqlInjection.status, 500);
+
+      const response = await sqlInjection.json();
+      t.equal(
+        response.error,
+        "Error executing tool weather: Zen has blocked an SQL injection: better-sqlite3.prepare(...) originating from aiToolParams.[0].location"
+      );
+    })
+    .catch((error) => {
+      t.fail(error.message);
+    })
+    .finally(() => {
+      server.kill();
+    });
+});
+
+t.test("it does not block in monitoring mode", (t) => {
+  const server = spawn(`node`, [pathToApp, "4007"], {
+    env: { ...process.env, AIKIDO_DEBUG: "true", AIKIDO_BLOCK: "false" },
+  });
+
+  server.on("close", () => {
+    t.end();
+  });
+
+  server.on("error", (err) => {
+    t.fail(err.message);
+  });
+
+  let stdout = "";
+  server.stdout.on("data", (data) => {
+    stdout += data.toString();
+  });
+
+  let stderr = "";
+  server.stderr.on("data", (data) => {
+    stderr += data.toString();
+  });
+
+  // Wait for the server to start
+  timeout(2000)
+    .then(() => {
+      return Promise.all([
+        fetch(
+          `http://127.0.0.1:4007/weather?prompt=${encodeURIComponent('What is the weather in "Ghent\'; DELETE FROM weather; --" like?')}`,
+          {
+            signal: AbortSignal.timeout(5000),
+          }
+        ),
+      ]);
+    })
+    .then(async ([sqlInjection]) => {
+      t.equal(sqlInjection.status, 500);
+
+      const response = await sqlInjection.json();
+      t.equal(
+        response.error,
+        "Error executing tool weather: The supplied SQL string contains more than one statement"
+      );
+    })
+    .catch((error) => {
+      t.fail(error.message);
+    })
+    .finally(() => {
+      server.kill();
+    });
+});

library/agent/Context.ts

@@ -30,6 +30,7 @@ export type Context = {
    */
   outgoingRequestRedirects?: { source: URL; destination: URL }[];
   executedMiddleware?: boolean;
+  aiToolParams?: unknown[]; // Parameters send to functions/tools that are called by a LLM
 };
 
 /**

library/agent/Source.ts

@@ -9,6 +9,7 @@ export const SOURCES = [
   "subdomains",
   "markUnsafe",
   "url",
+  "aiToolParams",
 ] as const;
 
 export type Source = (typeof SOURCES)[number];

library/agent/protect.ts

@@ -48,6 +48,7 @@ import { Fastify } from "../sources/Fastify";
 import { Koa } from "../sources/Koa";
 import { ClickHouse } from "../sinks/ClickHouse";
 import { Prisma } from "../sinks/Prisma";
+import { AiSDK } from "../sources/AiSDK";
 
 function getLogger(): Logger {
   if (isDebugging()) {
@@ -141,6 +142,7 @@ export function getWrappers() {
     new ClickHouse(),
     new Prisma(),
     // new Function(), Disabled because functionName.constructor === Function is false after patching global
+    new AiSDK(),
   ];
 }