diff --git a/control-plane/package-lock.json b/control-plane/package-lock.json index 2dfa119d..2b581d80 100644 --- a/control-plane/package-lock.json +++ b/control-plane/package-lock.json @@ -32,7 +32,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", @@ -16080,9 +16080,9 @@ } }, "node_modules/inferable": { - "version": "0.30.50", - "resolved": "https://registry.npmjs.org/inferable/-/inferable-0.30.50.tgz", - "integrity": "sha512-UBI0PYrdcpaqQY8q+Wx2rZoL/UxPlfcQ2dx1B1H07jBaWvZkOHL4OTxXRicZnERHCbbNCVXFgGlbEEOFOrA7vw==", + "version": "0.30.59", + "resolved": "https://registry.npmjs.org/inferable/-/inferable-0.30.59.tgz", + "integrity": "sha512-PSYX0un1iLYjCckGy77uj2Qby0uFIWiMpTzCzJKj/VHKc2syO5nxHUPRB886iUo6kZgX0V63XUgjLTI15Z9vWg==", "license": "MIT", "dependencies": { "@ts-rest/core": "^3.28.0", diff --git a/control-plane/package.json b/control-plane/package.json index 280cbd20..bd598b83 100644 --- 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 index 12e744b0..ee8cd06d 100644 --- 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 index 8e75bba1..f566da29 100644 --- a/control-plane/src/modules/service-definitions.test.ts +++ b/control-plane/src/modules/service-definitions.test.ts @@ -1,5 +1,5 @@ import { dereferenceSync, JSONSchema } from "dereference-json-schema"; -import { InvalidJobArgumentsError } from "../utilities/errors"; +import { InvalidJobArgumentsError, InvalidServiceRegistrationError } from "../utilities/errors"; import { packer } from "./packer"; import { deserializeFunctionSchema, @@ -7,9 +7,11 @@ import { 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 index f062b8bd..80056ad1 100644 --- 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)}` + ); + } + } } };