@takeshape/schema 11.40.3 → 11.43.0

package/dist/cjs/src/agents.js

@@ -3,8 +3,10 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.getInspectAgentSessionQueryName = exports.getAgentEndStates = exports.getAgentEndTransitions = exports.BUILT_IN_CHAT_ARG_NAMES = exports.BUILT_IN_CHAT_ARGS = exports.END_AGENT_EXECUTION = void 0;
+exports.addAiQueries = exports.createArgs = exports.removeBuiltInArgs = exports.getInspectAgentSessionQueryName = exports.getAgentEndStates = exports.getAgentEndTransitions = exports.BUILT_IN_CHAT_ARG_NAMES = exports.BUILT_IN_CHAT_ARGS = exports.END_AGENT_EXECUTION = void 0;
 const upperFirst_js_1 = __importDefault(require("lodash/upperFirst.js"));
+const uniq_js_1 = __importDefault(require("lodash/uniq.js"));
+const uniqBy_js_1 = __importDefault(require("lodash/uniqBy.js"));
 exports.END_AGENT_EXECUTION = 'endAgentExecution';
 exports.BUILT_IN_CHAT_ARGS = [
     {
@@ -56,3 +58,87 @@ const getInspectAgentSessionQueryName = (agentName) => {
     return `inspect${(0, upperFirst_js_1.default)(agentName)}`;
 };
 exports.getInspectAgentSessionQueryName = getInspectAgentSessionQueryName;
+const removeBuiltInArgs = (args) => {
+    return args.filter(arg => !exports.BUILT_IN_CHAT_ARG_NAMES.includes(arg.argName));
+};
+exports.removeBuiltInArgs = removeBuiltInArgs;
+const createArgs = (agent) => {
+    let apiArguments = (0, uniqBy_js_1.default)(agent.api.arguments ?? [], arg => arg.argName);
+    if (agent.api.type === 'chat') {
+        apiArguments = (0, exports.removeBuiltInArgs)(apiArguments).concat(exports.BUILT_IN_CHAT_ARGS);
+    }
+    return apiArguments.length > 0
+        ? {
+            type: 'object',
+            properties: apiArguments.reduce((acc, argument) => {
+                acc[argument.argName] = {
+                    type: argument.argType === 'sessionId' ? 'string' : argument.argType
+                };
+                return acc;
+            }, {}),
+            required: apiArguments.filter(arg => arg.required).map(arg => arg.argName)
+        }
+        : undefined;
+};
+exports.createArgs = createArgs;
+function addAiQueries(projectSchema) {
+    const agents = projectSchema['ai-experimental']?.agents;
+    if (!agents) {
+        return projectSchema;
+    }
+    const newSchema = {
+        ...projectSchema,
+        queries: { ...projectSchema.queries },
+        mutations: { ...projectSchema.mutations }
+    };
+    for (const [agentName, agent] of Object.entries(agents)) {
+        // Get valid return types based on states that could possibly be the end state
+        // TODO In the future the agent will have a return type defined and we will instead validate
+        // that all the return states will return the correct shape.
+        const returnStates = [...(0, exports.getAgentEndStates)(agent, true)];
+        const returnTypes = (0, uniq_js_1.default)(returnStates.map(stateId => {
+            const { execution } = agent.states[stateId];
+            if (execution.type === 'chat') {
+                return 'TSChatResponse';
+            }
+            if (execution.type === 'generate') {
+                return execution.outputShape ? execution.outputShape : 'string';
+            }
+            return 'JSON';
+        }));
+        const shape = returnTypes.length === 0 ? 'string' : returnTypes.length === 1 ? returnTypes[0] : 'JSON';
+        if (newSchema.mutations[agentName] !== undefined) {
+            throw new Error(`Schema already has a mutation with the name ${agentName}`);
+        }
+        newSchema.mutations[agentName] = {
+            description: agent.description,
+            args: (0, exports.createArgs)(agent),
+            shape,
+            resolver: {
+                name: 'ai:runAgent',
+                agentName
+            }
+        };
+        if (shape === 'TSChatResponse') {
+            newSchema.queries[(0, exports.getInspectAgentSessionQueryName)(agentName)] ||= {
+                description: `Inspect a session for the ${agentName} agent`,
+                args: {
+                    type: 'object',
+                    properties: {
+                        sessionId: {
+                            type: 'string'
+                        }
+                    },
+                    required: ['sessionId']
+                },
+                shape: 'TSAgentSession',
+                resolver: {
+                    name: 'ai:inspectAgentSession',
+                    agentName
+                }
+            };
+        }
+    }
+    return newSchema;
+}
+exports.addAiQueries = addAiQueries;

package/dist/cjs/src/index.js

@@ -58,3 +58,5 @@ __exportStar(require("./util/merge"), exports);
 __exportStar(require("./util/patch-schema"), exports);
 __exportStar(require("./util/shapes"), exports);
 __exportStar(require("./constants"), exports);
+__exportStar(require("./runtime-schema"), exports);
+__exportStar(require("./service-dependencies"), exports);

package/dist/cjs/src/relationships.js

@@ -3,9 +3,11 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.hasUnequalRelationships = exports.isEqualRelationship = exports.findExistingRelationships = exports.getRelationship = exports.getLegacyRelationship = exports.getRelationshipShapeIds = exports.getRelationshipShapes = exports.getRelationshipShapeRefs = exports.getRelationshipSchema = void 0;
+exports.addRelatedFields = exports.hasUnequalRelationships = exports.isEqualRelationship = exports.findExistingRelationships = exports.getRelationship = exports.getLegacyRelationship = exports.getRelationshipShapeIds = exports.getRelationshipShapes = exports.getRelationshipShapeRefs = exports.getRelationshipSchema = void 0;
 const util_1 = require("@takeshape/util");
 const find_js_1 = __importDefault(require("lodash/find.js"));
+const camelCase_js_1 = __importDefault(require("lodash/camelCase.js"));
+const uniq_js_1 = __importDefault(require("lodash/uniq.js"));
 const types_1 = require("./types");
 const refs_1 = require("./refs");
 const unions_1 = require("./unions");
@@ -235,3 +237,86 @@ function hasUnequalRelationships(relationships) {
     }));
 }
 exports.hasUnequalRelationships = hasUnequalRelationships;
+function getRelatedShapeIds(relationships) {
+    return (0, uniq_js_1.default)(relationships.map(rel => (rel.hasBackreference ? rel.shapeId : undefined)).filter(util_1.isDefined));
+}
+function getShapes(projectSchema, shapeIds) {
+    return shapeIds.map(shapeId => (0, shapes_1.getShapeById)(projectSchema, shapeId)).filter(util_1.isDefined);
+}
+/**
+ * Adds backreference fields to the schema.
+ */
+function addRelatedFields(projectSchema, allRelationships) {
+    for (const [shapeId, shapeRelationships] of Object.entries(allRelationships)) {
+        const shape = (0, shapes_1.getShapeById)(projectSchema, shapeId);
+        if (shape && (0, types_1.isObjectSchema)(shape.schema)) {
+            const relatedShapeIds = getRelatedShapeIds(shapeRelationships);
+            const relatedShapeNames = getShapes(projectSchema, relatedShapeIds).map(shape => shape.name);
+            if (relatedShapeNames.length) {
+                let shapeName;
+                if (relatedShapeNames.length === 1) {
+                    // If only one back reference exists, _references is a regular shape
+                    shapeName = relatedShapeNames[0];
+                }
+                else {
+                    // If many back references exist, _references will be a union of referring shapes
+                    shapeName = `${shape.name}Reference`;
+                    projectSchema.shapes[shapeName] = {
+                        id: shapeName,
+                        name: shapeName,
+                        title: shapeName,
+                        schema: {
+                            oneOf: relatedShapeNames.map(name => ({ '@ref': `local:${name}` }))
+                        }
+                    };
+                }
+                // _references, a list of all backreferences across all fields on this shape.
+                // It lists all the items that have references to the current item.
+                // It is for convenience.
+                if (shapeRelationships.some(rel => rel.hasBackreference)) {
+                    shape.schema.properties._references = {
+                        '@args': `TSListArgs<local:${shapeName}>`,
+                        '@ref': `PaginatedList<local:${shapeName}>`,
+                        '@resolver': {
+                            name: 'shapedb:list',
+                            service: 'shapedb',
+                            args: {
+                                ops: [
+                                    { path: '$', mapping: '$args' },
+                                    { path: `baseWhere._references.eq`, mapping: '$source._id' },
+                                    { path: `baseWhere._shapeId.in`, value: relatedShapeIds }
+                                ]
+                            }
+                        }
+                    };
+                }
+            }
+            // Create a schema for the backreference on the referred-to shape
+            for (const relationship of shapeRelationships) {
+                const relatedShape = (0, shapes_1.getShapeById)(projectSchema, relationship.shapeId);
+                if (relatedShape && relationship.hasBackreference) {
+                    const { relatedName } = relationship;
+                    const relatedFieldName = relatedName ? relatedName : `${(0, camelCase_js_1.default)(relatedShape.name)}Set`;
+                    const filterField = relatedName ? relationship.path.concat('_id').join('.') : '_references';
+                    shape.schema.properties[relatedFieldName] = {
+                        '@args': `TSListArgs<local:${relatedShape.name}>`,
+                        '@ref': `PaginatedList<local:${relatedShape.name}>`,
+                        '@resolver': {
+                            name: 'shapedb:list',
+                            service: 'shapedb',
+                            args: {
+                                ops: [
+                                    { path: '$', mapping: '$args' },
+                                    { path: `baseWhere.${filterField}.eq`, mapping: '$source._id' },
+                                    { path: `baseWhere._shapeId.eq`, value: relatedShape.id }
+                                ]
+                            }
+                        }
+                    };
+                }
+            }
+        }
+    }
+    return projectSchema;
+}
+exports.addRelatedFields = addRelatedFields;

package/dist/cjs/src/runtime-schema.js

@@ -0,0 +1,81 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.buildRuntimeSchema = exports.applyDefaultsAndFlatten = exports.applyLegacyCompatibilityTweaks = void 0;
+const errors_1 = require("@takeshape/errors");
+const compose_js_1 = __importDefault(require("lodash/fp/compose.js"));
+const util_1 = require("@takeshape/util");
+const set_js_1 = __importDefault(require("lodash/set.js"));
+const isError_js_1 = __importDefault(require("lodash/isError.js"));
+const types_1 = require("./types");
+const schema_util_1 = require("./schema-util");
+const relationships_1 = require("./relationships");
+const flatten_templates_1 = require("./flatten-templates");
+const agents_1 = require("./agents");
+const service_dependencies_1 = require("./service-dependencies");
+function applyLegacyCompatibilityTweaks(projectSchema) {
+    const newSchema = (0, util_1.deepClone)(projectSchema);
+    let hasSearchableShapes = false;
+    for (const [shapeName, shape] of Object.entries(newSchema.shapes)) {
+        if ((0, types_1.isModelShape)(shape)) {
+            hasSearchableShapes = true;
+            // Magic _contentTypeId field used for V1 compatibility
+            (0, set_js_1.default)(shape, ['schema', 'properties', '_contentTypeId'], { type: 'string' });
+            (0, set_js_1.default)(shape, ['schema', 'properties', '_contentTypeName'], { type: 'string' });
+            if (shape.model?.type !== 'single') {
+                newSchema.queries[`search${shapeName}Index`] = {
+                    shape: `SearchResults<${shapeName}>`,
+                    resolver: {
+                        name: 'takeshape:search',
+                        service: 'takeshape',
+                        shapeName
+                    },
+                    args: `TSSearchArgs<${shapeName}>`
+                };
+            }
+        }
+        if (shapeName === 'Asset') {
+            // Magic s3Key field to provide V1 compatibility with old projects until the field is
+            // formally removed and users are notified
+            (0, set_js_1.default)(shape, ['schema', 'properties', 's3Key'], {
+                title: 's3 key',
+                type: 'string',
+                '@mapping': 'shapedb:Asset.Hk6FQuz5',
+                '@deprecationReason': 'Use path instead'
+            });
+        }
+    }
+    if (hasSearchableShapes) {
+        // A placeholder shape for the TSSearchable interface
+        newSchema.shapes.TSSearchable = (0, schema_util_1.createShape)('TSSearchable', { type: 'object', properties: {} });
+        newSchema.queries.search = {
+            shape: 'SearchResults<TSSearchable>',
+            resolver: {
+                name: 'takeshape:search',
+                service: 'takeshape'
+            },
+            args: 'TSSearchArgs<TSSearchable>'
+        };
+    }
+    (0, relationships_1.addRelatedFields)(newSchema, (0, relationships_1.findExistingRelationships)(projectSchema, projectSchema.shapes));
+    return newSchema;
+}
+exports.applyLegacyCompatibilityTweaks = applyLegacyCompatibilityTweaks;
+exports.applyDefaultsAndFlatten = (0, compose_js_1.default)(flatten_templates_1.flattenTemplates, applyLegacyCompatibilityTweaks, schema_util_1.applyDefaultsToSchema, agents_1.addAiQueries);
+function buildRuntimeSchema(projectSchema, serviceLayers, log) {
+    try {
+        const projectSchemaWithDependencies = (0, service_dependencies_1.resolveSchemaShapeDependencies)(projectSchema, serviceLayers);
+        return (0, exports.applyDefaultsAndFlatten)(projectSchemaWithDependencies);
+    }
+    catch (err) {
+        const error = new errors_1.SchemaBuildError('An error occurred while building the schema', {
+            cause: (0, isError_js_1.default)(err) ? err : undefined,
+            schema: projectSchema
+        });
+        log('build runtime error', error);
+        throw error;
+    }
+}
+exports.buildRuntimeSchema = buildRuntimeSchema;

package/dist/cjs/src/service-dependencies.js

@@ -0,0 +1,167 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.resolveSchemaShapeDependencies = exports.collectReferencedShapeNames = void 0;
+const util_1 = require("@takeshape/util");
+const get_js_1 = __importDefault(require("lodash/get.js"));
+const set_js_1 = __importDefault(require("lodash/set.js"));
+const isObject_js_1 = __importDefault(require("lodash/isObject.js"));
+const pick_js_1 = __importDefault(require("lodash/pick.js"));
+const types_1 = require("./types");
+const refs_1 = require("./refs");
+const interfaces_1 = require("./interfaces");
+const schema_util_1 = require("./schema-util");
+function getGraphQLServiceShapeMap(layers) {
+    return Object.assign({}, ...Object.values(layers).map(layer => layer.schema?.shapes ?? {}));
+}
+function schemaExtendsShape(context, shapeName, schema) {
+    return ((0, types_1.isExtendsSchema)(schema) &&
+        schema.extends.some((item) => (0, refs_1.getRefShapeName)(context, item) === shapeName));
+}
+function getMissingPropertyRefs(projectSchema, getNamespace) {
+    const propertyRefs = (0, refs_1.getAllPropertyRefs)(projectSchema);
+    return propertyRefs.filter(ref => {
+        const path = (0, refs_1.propertyRefItemToPath)(getNamespace, ref);
+        return !(0, get_js_1.default)(projectSchema, path);
+    });
+}
+function refToQuery(serviceSchemas, ref) {
+    const serviceSchema = serviceSchemas[ref.serviceId].schema;
+    if (serviceSchema) {
+        if (ref.shapeName === 'Query' || ref.shapeName === 'Mutation') {
+            const operation = ref.shapeName === 'Query' ? 'queries' : 'mutations';
+            return serviceSchema[operation][ref.propertyName];
+        }
+    }
+}
+function getShapeRefs(services, serviceSchemas, propertyRefs) {
+    const results = [];
+    for (const ref of propertyRefs) {
+        const query = refToQuery(serviceSchemas, ref);
+        const layerSchema = serviceSchemas[ref.serviceId]?.schema;
+        if (query && layerSchema) {
+            const layerWithService = {
+                ...layerSchema,
+                services: (0, pick_js_1.default)(services, [ref.serviceId]) // include own service so refs have isValidService: true
+            };
+            results.push(...(0, schema_util_1.getAllRefsInQuery)(layerWithService, [], query));
+        }
+    }
+    return results;
+}
+function interfacesToShapeNames(projectSchema, shapeRefArray) {
+    return shapeRefArray.map(shapeRef => (0, refs_1.refItemToNamespacedShapeName)((0, refs_1.refExpressionToRefItem)(projectSchema, shapeRef)));
+}
+/**
+ * Add all referenced shape names to the refSet meant to be used as a callback to visitSchemaProperties
+ */
+function collectReferencedShapeNames(context, shapeNames) {
+    const addRef = (refItem) => {
+        shapeNames.add((0, refs_1.refItemToNamespacedShapeName)(refItem));
+    };
+    const collect = (schema) => {
+        const args = schema['@args'];
+        if (args && (0, isObject_js_1.default)(args) && (0, types_1.isObjectSchema)(args)) {
+            Object.values(args.properties).forEach(collect);
+        }
+        const argsRefItem = (0, schema_util_1.getArgsReference)(context, schema);
+        if (argsRefItem) {
+            addRef(argsRefItem);
+        }
+        const refItem = (0, refs_1.getRef)(context, schema);
+        if (refItem) {
+            addRef(refItem);
+        }
+    };
+    return collect;
+}
+exports.collectReferencedShapeNames = collectReferencedShapeNames;
+/**
+ * Returns a schema with all of its dependencies resolved.
+ * SHOULD only return the minimum set of remote dependencies required to enable
+ * the local shape definitions.
+ */
+function resolveSchemaShapeDependencies(projectSchema, serviceLayers) {
+    const newSchema = (0, util_1.deepClone)(projectSchema);
+    if ((0, types_1.isProjectSchemaWithServices)(projectSchema)) {
+        const newShapeMap = {};
+        const requiredShapeNames = new Set((0, schema_util_1.getAllNamespaceShapes)(projectSchema));
+        const getNamespace = (0, refs_1.createGetNamespace)(projectSchema);
+        const missingPropertyRefs = getMissingPropertyRefs(projectSchema, getNamespace);
+        if (requiredShapeNames.size || missingPropertyRefs.length) {
+            const remoteShapeMap = getGraphQLServiceShapeMap(serviceLayers);
+            // Add shape references from queries
+            (0, util_1.addAll)(requiredShapeNames, getShapeRefs(projectSchema.services, serviceLayers, missingPropertyRefs)
+                .filter(schema_util_1.isValidRefItem)
+                .map(refs_1.refItemToNamespacedShapeName));
+            // Add referenced queries and mutations
+            for (const ref of missingPropertyRefs) {
+                const query = refToQuery(serviceLayers, ref);
+                if (query) {
+                    (0, set_js_1.default)(newSchema, (0, refs_1.propertyRefItemToPath)(getNamespace, ref), query);
+                }
+            }
+            // Gather all the remote dependencies that our local schema requires.
+            // If the required shape exists in the remote shape map, visit the
+            // properties on the remote shape's schema and ensure they are required
+            // as well.
+            const collectDependencies = (shapeNames) => {
+                for (const name of shapeNames) {
+                    const localShape = newSchema.shapes[name];
+                    const remoteShape = remoteShapeMap[name];
+                    // Don't check for dependencies if the shape is already in our schema unless it extends a remote shape
+                    // getNamespacedShapeNameList already collected those dependencies
+                    if ((!localShape || (localShape && schemaExtendsShape(projectSchema, name, localShape.schema))) &&
+                        remoteShape) {
+                        (0, schema_util_1.visitSchemaProperties)(remoteShape.schema, [remoteShape.name, 'schema'], collectReferencedShapeNames({ services: newSchema.services }, shapeNames));
+                    }
+                }
+            };
+            collectDependencies(requiredShapeNames);
+            // Collect interfaces that requiredShapeNames implement
+            const interfaceNames = new Set();
+            const requiredShapes = (0, util_1.mapSet)(requiredShapeNames, shapeName => newSchema.shapes[shapeName] ?? remoteShapeMap[shapeName]);
+            for (const shape of requiredShapes) {
+                if (shape?.interfaces) {
+                    (0, util_1.addAll)(interfaceNames, interfacesToShapeNames(projectSchema, shape.interfaces));
+                }
+            }
+            // Collect implementations for interfaces in requiredShapeNames
+            const implementationNames = new Set();
+            const implementations = (0, interfaces_1.getImplementationShapeNameMap)({ ...newSchema.shapes, ...remoteShapeMap });
+            for (const shape of requiredShapes) {
+                if ((0, interfaces_1.isInterfaceShape)(shape) && implementations[shape.name]) {
+                    (0, util_1.addAll)(implementationNames, implementations[shape.name]);
+                }
+            }
+            // Collect the dependencies for the interfaces and implementations
+            collectDependencies(interfaceNames);
+            collectDependencies(implementationNames);
+            (0, util_1.addAll)(requiredShapeNames, interfaceNames);
+            (0, util_1.addAll)(requiredShapeNames, implementationNames);
+            // Populate a ShapeMap with the remote shapes we've gathered
+            // A combined context which prefers remote shapes for use when merging/inlining
+            const mergingContext = { services: newSchema.services, shapes: { ...newSchema.shapes, ...remoteShapeMap } };
+            for (const name of requiredShapeNames) {
+                const localShape = newSchema.shapes[name];
+                const remoteShape = remoteShapeMap[name];
+                if (localShape && remoteShape && schemaExtendsShape(mergingContext, name, localShape.schema)) {
+                    // Merge and inline combined shape when a local shape extends a remote shape
+                    newShapeMap[name] = {
+                        ...localShape,
+                        schema: (0, refs_1.dereferenceObjectSchema)(mergingContext, newSchema.shapes[name].schema)
+                    };
+                }
+                else if (localShape ?? remoteShape) {
+                    newShapeMap[name] = localShape ?? remoteShape;
+                }
+            }
+        }
+        // Assign remote shapes to the new schema, mutating the new schema
+        Object.assign(newSchema.shapes, (0, interfaces_1.pruneUnusedInterfaces)(newSchema, newShapeMap));
+    }
+    return newSchema;
+}
+exports.resolveSchemaShapeDependencies = resolveSchemaShapeDependencies;

package/dist/cjs/src/util/patch-schema.js

@@ -3,48 +3,53 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.patchSchema = exports.shallowMerge = void 0;
-const forIn_js_1 = __importDefault(require("lodash/forIn.js"));
+exports.patchSchema = void 0;
 const omit_js_1 = __importDefault(require("lodash/omit.js"));
-const schemaUpdateMergedKeys = ['queries', 'mutations', 'shapes', 'workflows', 'forms', 'services'];
+const get_js_1 = __importDefault(require("lodash/get.js"));
+const set_js_1 = __importDefault(require("lodash/fp/set.js"));
+const unset_js_1 = __importDefault(require("lodash/unset.js"));
+/**
+ * Array of lodash.get paths to keys that are merged when applying a schema update.
+ */
+const schemaUpdateMergedKeys = [
+    'queries',
+    'mutations',
+    'shapes',
+    'workflows',
+    'forms',
+    'services',
+    'ai-experimental.agents'
+];
 function shallowMerge(original, update) {
     const result = { ...original };
-    if (update) {
-        (0, forIn_js_1.default)(update, (obj, key) => {
-            // Null is a sentinel value to indicate delete
-            if (obj === null) {
-                // eslint-disable-next-line @typescript-eslint/no-dynamic-delete
-                delete result[key];
-            }
-            else {
-                result[key] = obj;
-            }
-        });
+    for (const [key, obj] of Object.entries(update)) {
+        // Null is a sentinel value to indicate delete
+        if (obj === null) {
+            (0, unset_js_1.default)(result, key);
+        }
+        else {
+            result[key] = obj;
+        }
     }
     return result;
 }
-exports.shallowMerge = shallowMerge;
 /**
- * Apply a schema update to a schema.
- * Resulting schema is not assumed to be valid and should be validated if appropriate.
+ * Apply a schema update to a schema. This can operate on a schema with frozen properties,
+ * such as one produced by immer. Resulting schema is not assumed to be valid and should
+ * be validated if appropriate.
  */
-function patchSchema(projectSchema, update) {
+function patchSchema(projectSchema, schemaUpdate) {
     // Patch old schema with update
-    const updatedSchema = {
+    let updatedSchema = {
         ...projectSchema,
-        ...(0, omit_js_1.default)(update, [...schemaUpdateMergedKeys, 'ai-experimental'])
+        ...(0, omit_js_1.default)(schemaUpdate, schemaUpdateMergedKeys)
     };
     for (const key of schemaUpdateMergedKeys) {
-        if (update[key]) {
-            updatedSchema[key] = shallowMerge(projectSchema[key], update[key]);
+        const update = (0, get_js_1.default)(schemaUpdate, key);
+        if (update) {
+            updatedSchema = (0, set_js_1.default)(key, shallowMerge((0, get_js_1.default)(projectSchema, key) ?? {}, update), updatedSchema);
         }
     }
-    if (update['ai-experimental']?.agents) {
-        const ai = {
-            agents: shallowMerge(projectSchema['ai-experimental']?.agents, update['ai-experimental'].agents)
-        };
-        updatedSchema['ai-experimental'] = ai;
-    }
     return updatedSchema;
 }
 exports.patchSchema = patchSchema;

package/dist/cjs/src/validate.js

@@ -141,6 +141,9 @@ function enumerateBasicResolvers(resolver, path) {
     visit(resolver, path);
     return results;
 }
+const isValidAgentMutation = (projectSchema, mutationName) => {
+    return Boolean(projectSchema['ai-experimental']?.agents?.[mutationName]);
+};
 const validateAIToolConfig = (projectSchema, getNamespace, tool, basePath) => {
     const toolRef = (0, ai_tools_1.getToolRef)(tool);
     const path = typeof tool === 'string' ? basePath : basePath.concat('ref');
@@ -164,6 +167,9 @@ const validateAIToolConfig = (projectSchema, getNamespace, tool, basePath) => {
         };
     }
     const property = (0, get_js_1.default)(projectSchema, propertyPath);
+    if (propertyPath[0] === 'mutations' && isValidAgentMutation(projectSchema, propertyPath[1])) {
+        return;
+    }
     if (!property) {
         return {
             type: 'notFound',

package/dist/esm/src/agents.js

@@ -1,4 +1,6 @@
 import upperFirst from 'lodash/upperFirst.js';
+import uniq from 'lodash/uniq.js';
+import uniqBy from 'lodash/uniqBy.js';
 export const END_AGENT_EXECUTION = 'endAgentExecution';
 export const BUILT_IN_CHAT_ARGS = [
     {
@@ -47,3 +49,84 @@ export const getAgentEndStates = (agent, includeSuspend = false) => {
 export const getInspectAgentSessionQueryName = (agentName) => {
     return `inspect${upperFirst(agentName)}`;
 };
+export const removeBuiltInArgs = (args) => {
+    return args.filter(arg => !BUILT_IN_CHAT_ARG_NAMES.includes(arg.argName));
+};
+export const createArgs = (agent) => {
+    let apiArguments = uniqBy(agent.api.arguments ?? [], arg => arg.argName);
+    if (agent.api.type === 'chat') {
+        apiArguments = removeBuiltInArgs(apiArguments).concat(BUILT_IN_CHAT_ARGS);
+    }
+    return apiArguments.length > 0
+        ? {
+            type: 'object',
+            properties: apiArguments.reduce((acc, argument) => {
+                acc[argument.argName] = {
+                    type: argument.argType === 'sessionId' ? 'string' : argument.argType
+                };
+                return acc;
+            }, {}),
+            required: apiArguments.filter(arg => arg.required).map(arg => arg.argName)
+        }
+        : undefined;
+};
+export function addAiQueries(projectSchema) {
+    const agents = projectSchema['ai-experimental']?.agents;
+    if (!agents) {
+        return projectSchema;
+    }
+    const newSchema = {
+        ...projectSchema,
+        queries: { ...projectSchema.queries },
+        mutations: { ...projectSchema.mutations }
+    };
+    for (const [agentName, agent] of Object.entries(agents)) {
+        // Get valid return types based on states that could possibly be the end state
+        // TODO In the future the agent will have a return type defined and we will instead validate
+        // that all the return states will return the correct shape.
+        const returnStates = [...getAgentEndStates(agent, true)];
+        const returnTypes = uniq(returnStates.map(stateId => {
+            const { execution } = agent.states[stateId];
+            if (execution.type === 'chat') {
+                return 'TSChatResponse';
+            }
+            if (execution.type === 'generate') {
+                return execution.outputShape ? execution.outputShape : 'string';
+            }
+            return 'JSON';
+        }));
+        const shape = returnTypes.length === 0 ? 'string' : returnTypes.length === 1 ? returnTypes[0] : 'JSON';
+        if (newSchema.mutations[agentName] !== undefined) {
+            throw new Error(`Schema already has a mutation with the name ${agentName}`);
+        }
+        newSchema.mutations[agentName] = {
+            description: agent.description,
+            args: createArgs(agent),
+            shape,
+            resolver: {
+                name: 'ai:runAgent',
+                agentName
+            }
+        };
+        if (shape === 'TSChatResponse') {
+            newSchema.queries[getInspectAgentSessionQueryName(agentName)] ||= {
+                description: `Inspect a session for the ${agentName} agent`,
+                args: {
+                    type: 'object',
+                    properties: {
+                        sessionId: {
+                            type: 'string'
+                        }
+                    },
+                    required: ['sessionId']
+                },
+                shape: 'TSAgentSession',
+                resolver: {
+                    name: 'ai:inspectAgentSession',
+                    agentName
+                }
+            };
+        }
+    }
+    return newSchema;
+}

package/dist/esm/src/index.js

@@ -42,3 +42,5 @@ export * from "./util/merge.js";
 export * from "./util/patch-schema.js";
 export * from "./util/shapes.js";
 export * from "./constants.js";
+export * from "./runtime-schema.js";
+export * from "./service-dependencies.js";