chore: Validate handleCustomerAuth registration

inferablehq · Dec 10, 2024 · 2593832 · 2593832
1 parent 4261e78
commit 2593832
Show file tree

Hide file tree

Showing 5 changed files with 120 additions and 9 deletions.
diff --git a/control-plane/package-lock.json b/control-plane/package-lock.json
diff --git a/control-plane/package.json b/control-plane/package.json
@@ -42,7 +42,7 @@
     "fastify": "^4.21.0",
     "fastify-plugin": "^4.5.1",
     "flagsmith-nodejs": "^3.3.3",
-    "inferable": "^0.30.39",
+    "inferable": "^0.30.59",
     "jest": "^29.6.4",
     "js-tiktoken": "^1.0.12",
     "json-schema-to-zod": "^2.1.0",

diff --git a/control-plane/src/modules/auth/customer-auth.ts b/control-plane/src/modules/auth/customer-auth.ts
@@ -7,8 +7,8 @@ import { packer } from "../packer";
 import * as jobs from "../jobs/jobs";
 import { getJobStatusSync } from "../jobs/jobs";
 
-const VERIFY_FUNCTION_NAME = "handleCustomerAuth";
-const VERIFY_FUNCTION_SERVICE = "default";
+export const VERIFY_FUNCTION_NAME = "handleCustomerAuth";
+export const VERIFY_FUNCTION_SERVICE = "default";
 const VERIFY_FUNCTION_ID = `${VERIFY_FUNCTION_SERVICE}_${VERIFY_FUNCTION_NAME}`;
 
 /**

diff --git a/control-plane/src/modules/service-definitions.test.ts b/control-plane/src/modules/service-definitions.test.ts
@@ -1,15 +1,17 @@
 import { dereferenceSync, JSONSchema } from "dereference-json-schema";
-import { InvalidJobArgumentsError } from "../utilities/errors";
+import { InvalidJobArgumentsError, InvalidServiceRegistrationError } from "../utilities/errors";
 import { packer } from "./packer";
 import {
   deserializeFunctionSchema,
   embeddableServiceFunction,
   parseJobArgs,
   serviceFunctionEmbeddingId,
   updateServiceEmbeddings,
+  validateServiceRegistration,
 } from "./service-definitions";
 import { createOwner } from "./test/util";
 import { zodToJsonSchema } from "zod-to-json-schema";
+import { z } from "zod";
 describe("updateServiceEmbeddings", () => {
   let owner: { clusterId: string };
   beforeAll(async () => {
@@ -281,3 +283,92 @@ describe("deserializeFunctionSchema", () => {
     });
   });
 });
+
+describe("validateServiceRegistration", () => {
+  it("should reject invalid schema", () => {
+    expect(() => {
+      validateServiceRegistration({
+        service: "default",
+        definition: {
+          name: "default",
+          functions: [
+            {
+              name: "someFn",
+              schema: JSON.stringify({
+                type: "wrong_type",
+              })
+            },
+          ],
+        },
+      });
+    }).toThrow(InvalidServiceRegistrationError);
+  })
+
+  it("should accept valid schema", () => {
+    expect(() => {
+      validateServiceRegistration({
+        service: "default",
+        definition: {
+          name: "default",
+          functions: [
+            {
+              name: "someFn",
+              description: "someFn",
+              schema: JSON.stringify(zodToJsonSchema(
+                z.object({
+                  test: z.string(),
+                })
+              ))
+            },
+          ],
+        },
+      });
+    }).not.toThrow();
+  })
+
+  it("should reject incorrect handleCustomerAuth registration", () => {
+    expect(() => {
+      validateServiceRegistration({
+        service: "default",
+        definition: {
+          name: "default",
+          functions: [
+            {
+              name: "handleCustomerAuth",
+              description: "handleCustomerAuth",
+              schema: JSON.stringify(zodToJsonSchema(
+                z.object({
+                  test: z.string(),
+                })
+              ))
+            },
+          ],
+        },
+      });
+    }).toThrow(InvalidServiceRegistrationError);
+  })
+
+  it("should accept valid handleCustomerAuth registration", () => {
+    expect(() => {
+      validateServiceRegistration({
+        service: "default",
+        definition: {
+          name: "default",
+          functions: [
+            {
+              name: "handleCustomerAuth",
+              description: "handleCustomerAuth",
+              schema: JSON.stringify(zodToJsonSchema(
+                z.object({
+                  token: z.string(),
+                })
+              ))
+            },
+          ],
+        },
+      });
+    }).not.toThrow();
+  })
+
+})
+
diff --git a/control-plane/src/modules/service-definitions.ts b/control-plane/src/modules/service-definitions.ts
@@ -1,5 +1,6 @@
 import { and, eq, lte } from "drizzle-orm";
 import {
+    handleCustomerAuthSchema,
   validateDescription,
   validateFunctionName,
   validateFunctionSchema,
@@ -19,6 +20,7 @@ import { embeddableEntitiy } from "./embeddings/embeddings";
 import { logger } from "./observability/logger";
 import { packer } from "./packer";
 import { withThrottle } from "./util";
+import { VERIFY_FUNCTION_NAME, VERIFY_FUNCTION_SERVICE } from "./auth/customer-auth";
 
 // The time without a ping before a service is considered expired
 const SERVICE_LIVE_THRESHOLD_MS = 30 * 60 * 1000; // 30 minutes
@@ -369,7 +371,7 @@ export const updateServiceEmbeddings = async ({
   );
 };
 
-const validateServiceRegistration = ({
+export const validateServiceRegistration = ({
   service,
   definition,
 }: {
@@ -398,6 +400,24 @@ const validateServiceRegistration = ({
         );
       }
     }
+
+    // Checks for customer auth handler
+    if (service === VERIFY_FUNCTION_SERVICE && fn.name === VERIFY_FUNCTION_NAME) {
+      if (!fn.schema) {
+        throw new InvalidServiceRegistrationError(
+          `${fn.name} must have a valid schema`
+        );
+      }
+
+      // Check that the schema accepts and expected value
+      const zodSchema = deserializeFunctionSchema(fn.schema);
+      const schema = zodSchema.safeParse({ token: "test" });
+      if (!schema.success) {
+        throw new InvalidServiceRegistrationError(
+          `${fn.name} schema invalid: ${JSON.stringify(schema.error.issues)}`
+        );
+      }
+    }
   }
 };