fix: support JSON Schema extension keywords and properly infer required fields

Author: Kamiel Wanrooij

runtimes/runtimes/agent.test.ts

@@ -17,6 +17,21 @@ describe('Agent Tools', () => {
             required: ['test'],
         },
     } as const
+
+    const SOME_TOOL_SPEC_WITH_EXTENSIONS = {
+        name: 'test',
+        description: 'test',
+        inputSchema: {
+            type: 'object',
+            properties: {
+                test: {
+                    type: 'string',
+                },
+            },
+            required: ['test'],
+            someExtension: true,
+        },
+    } as const
     const SOME_TOOL_HANDLER = async (_: { test: string }) => true
 
     beforeEach(() => {
@@ -91,4 +106,13 @@ describe('Agent Tools', () => {
         assert.equal(tools[0].description, SOME_TOOL_SPEC.description)
         assert.equal(tools[0].inputSchema, SOME_TOOL_SPEC.inputSchema)
     })
+
+    it('should support JSON Schema extension keywords', () => {
+        AGENT.addTool(SOME_TOOL_SPEC_WITH_EXTENSIONS, SOME_TOOL_HANDLER)
+        const tools = AGENT.getTools({ format: 'bedrock' })
+        assert.equal(
+            tools[0].toolSpecification.inputSchema.json['someExtension'],
+            SOME_TOOL_SPEC_WITH_EXTENSIONS.inputSchema.someExtension
+        )
+    })
 })

runtimes/runtimes/agent.ts

@@ -11,7 +11,7 @@ type Tool<T, R> = {
 
 export const newAgent = (): Agent => {
     const tools: Record<string, Tool<any, any>> = {}
-    const ajv = new Ajv()
+    const ajv = new Ajv({ strictSchema: false })
 
     return {
         addTool: <T extends InferSchema<S['inputSchema']>, S extends ToolSpec>(

runtimes/server-interface/agent.ts

@@ -6,6 +6,7 @@ interface BaseSchema {
     $id?: string
     $schema?: string
     definitions?: Record<string, JSONSchema>
+    [extensionKeywords: string]: any
 }
 
 interface StringSchema extends BaseSchema {
@@ -74,9 +75,17 @@ type InferArray<T extends { type: 'array'; items: any }> = T['items'] extends {
     ? InferSchema<T['items']>[]
     : never
 
-type InferObject<T extends { type: 'object'; properties: Record<string, any> }> = {
-    [K in keyof T['properties']]?: InferSchema<T['properties'][K]>
-} & (T extends { required: string[] } ? { [K in T['required'][number]]: InferSchema<T['properties'][K]> } : unknown)
+type InferObject<T extends { type: 'object'; properties: Record<string, any> }> = T extends {
+    required: readonly string[]
+}
+    ? {
+          [K in keyof T['properties'] as K extends T['required'][number] ? K : never]: InferSchema<T['properties'][K]>
+      } & {
+          [K in keyof T['properties'] as K extends T['required'][number] ? never : K]?: InferSchema<T['properties'][K]>
+      }
+    : {
+          [K in keyof T['properties']]?: InferSchema<T['properties'][K]>
+      }
 
 export type InferSchema<T> = T extends { type: 'array'; items: any }
     ? InferArray<T>