feat(backend): Require discriminator value on graph save (#9858)

If a node has a multi-credentials input (e.g. AI Text Generator block)
but the discriminator value (e.g. model choice) is missing, the input
can't be discriminated into a single-provider input. Discrimination into
a single-provider input is necessary to make a graph-level credentials
input for use in the Library.

### Changes 🏗️

- feat(backend): Require discriminator fields to always have a value

- dx(frontend): Improve typing of discriminator stuff
- dx(frontend): Fix typing in `NodeOneOfDiscriminatorField` component

### Checklist 📋

#### For code changes:
- [x] I have clearly listed my changes in the PR description
- [x] I have made a test plan
- [x] I have tested my changes according to the test plan:
- [x] Saving & running graphs with and without credentials works
normally
- Note: We don't have any blocks with a discriminator that doesn't have
a default value, so currently I don't think it's possible to produce a
case where this mechanism would be triggered.

autogpt_platform/backend/backend/data/graph.py

@@ -428,10 +428,6 @@ class GraphModel(Graph):
             if (block := get_block(node.block_id)) is not None
         }
 
-        for node in graph.nodes:
-            if (block := nodes_block.get(node.id)) is None:
-                raise ValueError(f"Invalid block {node.block_id} for node #{node.id}")
-
         input_links = defaultdict(list)
 
         for link in graph.links:
@@ -446,8 +442,8 @@ class GraphModel(Graph):
                 [sanitize(name) for name in node.input_default]
                 + [sanitize(link.sink_name) for link in input_links.get(node.id, [])]
             )
-            input_schema = block.input_schema
-            for name in (required_fields := input_schema.get_required_fields()):
+            InputSchema = block.input_schema
+            for name in (required_fields := InputSchema.get_required_fields()):
                 if (
                     name not in provided_inputs
                     # Webhook payload is passed in by ExecutionManager
@@ -457,7 +453,7 @@ class GraphModel(Graph):
                         in (BlockType.WEBHOOK, BlockType.WEBHOOK_MANUAL)
                     )
                     # Checking availability of credentials is done by ExecutionManager
-                    and name not in input_schema.get_credentials_fields()
+                    and name not in InputSchema.get_credentials_fields()
                     # Validate only I/O nodes, or validate everything when executing
                     and (
                         for_run
@@ -484,37 +480,43 @@ class GraphModel(Graph):
                     )
 
             # Get input schema properties and check dependencies
-            input_fields = input_schema.model_fields
+            input_fields = InputSchema.model_fields
 
-            def has_value(name):
+            def has_value(node: Node, name: str):
                 return (
-                    node is not None
-                    and name in node.input_default
+                    name in node.input_default
                     and node.input_default[name] is not None
                     and str(node.input_default[name]).strip() != ""
                 ) or (name in input_fields and input_fields[name].default is not None)
 
             # Validate dependencies between fields
-            for field_name, field_info in input_fields.items():
-                # Apply input dependency validation only on run & field with depends_on
-                json_schema_extra = field_info.json_schema_extra or {}
-                if not (
-                    for_run
-                    and isinstance(json_schema_extra, dict)
-                    and (
-                        dependencies := cast(
-                            list[str], json_schema_extra.get("depends_on", [])
-                        )
-                    )
-                ):
+            for field_name in input_fields.keys():
+                field_json_schema = InputSchema.get_field_schema(field_name)
+
+                dependencies: list[str] = []
+
+                # Check regular field dependencies (only pre graph execution)
+                if for_run:
+                    dependencies.extend(field_json_schema.get("depends_on", []))
+
+                # Require presence of credentials discriminator (always).
+                # The `discriminator` is either the name of a sibling field (str),
+                # or an object that discriminates between possible types for this field:
+                # {"propertyName": prop_name, "mapping": {prop_value: sub_schema}}
+                if (
+                    discriminator := field_json_schema.get("discriminator")
+                ) and isinstance(discriminator, str):
+                    dependencies.append(discriminator)
+
+                if not dependencies:
                     continue
 
                 # Check if dependent field has value in input_default
-                field_has_value = has_value(field_name)
+                field_has_value = has_value(node, field_name)
                 field_is_required = field_name in required_fields
 
                 # Check for missing dependencies when dependent field is present
-                missing_deps = [dep for dep in dependencies if not has_value(dep)]
+                missing_deps = [dep for dep in dependencies if not has_value(node, dep)]
                 if missing_deps and (field_has_value or field_is_required):
                     raise ValueError(
                         f"Node {block.name} #{node.id}: Field `{field_name}` requires [{', '.join(missing_deps)}] to be set"

autogpt_platform/frontend/src/components/node-input-components.tsx

@@ -14,6 +14,7 @@ import {
   BlockIOArraySubSchema,
   BlockIOBooleanSubSchema,
   BlockIOCredentialsSubSchema,
+  BlockIODiscriminatedOneOfSubSchema,
   BlockIOKVSubSchema,
   BlockIONumberSubSchema,
   BlockIOObjectSubSchema,
@@ -536,7 +537,7 @@ export const NodeGenericInputField: FC<{
 const NodeOneOfDiscriminatorField: FC<{
   nodeId: string;
   propKey: string;
-  propSchema: any;
+  propSchema: BlockIODiscriminatedOneOfSubSchema;
   currentValue?: any;
   defaultValue?: any;
   errors: { [key: string]: string | undefined };
@@ -564,25 +565,25 @@ const NodeOneOfDiscriminatorField: FC<{
     const oneOfVariants = propSchema.oneOf || [];
 
     return oneOfVariants
-      .map((variant: any) => {
-        const variantDiscValue =
-          variant.properties?.[discriminatorProperty]?.const;
+      .map((variant) => {
+        const variantDiscValue = variant.properties?.[discriminatorProperty]
+          ?.const as string; // NOTE: can discriminators only be strings?
 
         return {
           value: variantDiscValue,
-          schema: variant as BlockIOSubSchema,
+          schema: variant,
         };
       })
-      .filter((v: any) => v.value != null);
+      .filter((v) => v.value != null);
   }, [discriminatorProperty, propSchema.oneOf]);
 
   const initialVariant = defaultValue
     ? variantOptions.find(
-        (opt: any) => defaultValue[discriminatorProperty] === opt.value,
+        (opt) => defaultValue[discriminatorProperty] === opt.value,
       )
     : currentValue
       ? variantOptions.find(
-          (opt: any) => currentValue[discriminatorProperty] === opt.value,
+          (opt) => currentValue[discriminatorProperty] === opt.value,
         )
       : null;
 
@@ -603,9 +604,7 @@ const NodeOneOfDiscriminatorField: FC<{
 
   const handleVariantChange = (newType: string) => {
     setChosenType(newType);
-    const chosenVariant = variantOptions.find(
-      (opt: any) => opt.value === newType,
-    );
+    const chosenVariant = variantOptions.find((opt) => opt.value === newType);
     if (chosenVariant) {
       const initialValue = {
         [discriminatorProperty]: newType,
@@ -615,7 +614,7 @@ const NodeOneOfDiscriminatorField: FC<{
   };
 
   const chosenVariantSchema = variantOptions.find(
-    (opt: any) => opt.value === chosenType,
+    (opt) => opt.value === chosenType,
   )?.schema;
 
   function getEntryKey(key: string): string {
@@ -664,7 +663,7 @@ const NodeOneOfDiscriminatorField: FC<{
                 >
                   <NodeHandle
                     keyName={getEntryKey(someKey)}
-                    schema={childSchema as BlockIOSubSchema}
+                    schema={childSchema}
                     isConnected={isConnected(getEntryKey(someKey))}
                     isRequired={false}
                     side="left"
@@ -675,7 +674,7 @@ const NodeOneOfDiscriminatorField: FC<{
                       nodeId={nodeId}
                       key={propKey}
                       propKey={childKey}
-                      propSchema={childSchema as BlockIOSubSchema}
+                      propSchema={childSchema}
                       currentValue={
                         currentValue
                           ? currentValue[someKey]

autogpt_platform/frontend/src/lib/autogpt-server-api/types.ts

@@ -94,6 +94,7 @@ export type BlockIOSubSchemaMeta = {
 export type BlockIOObjectSubSchema = BlockIOSubSchemaMeta & {
   type: "object";
   properties: { [key: string]: BlockIOSubSchema };
+  const?: { [key: keyof BlockIOObjectSubSchema["properties"]]: any };
   default?: { [key: keyof BlockIOObjectSubSchema["properties"]]: any };
   required?: (keyof BlockIOObjectSubSchema["properties"])[];
   secret?: boolean;
@@ -102,6 +103,7 @@ export type BlockIOObjectSubSchema = BlockIOSubSchemaMeta & {
 export type BlockIOKVSubSchema = BlockIOSubSchemaMeta & {
   type: "object";
   additionalProperties?: { type: "string" | "number" | "integer" };
+  const?: { [key: string]: string | number };
   default?: { [key: string]: string | number };
   secret?: boolean;
 };
@@ -109,6 +111,7 @@ export type BlockIOKVSubSchema = BlockIOSubSchemaMeta & {
 export type BlockIOArraySubSchema = BlockIOSubSchemaMeta & {
   type: "array";
   items?: BlockIOSimpleTypeSubSchema;
+  const?: Array<string>;
   default?: Array<string>;
   secret?: boolean;
 };
@@ -117,6 +120,7 @@ export type BlockIOStringSubSchema = BlockIOSubSchemaMeta & {
   type: "string";
   enum?: string[];
   secret?: true;
+  const?: string;
   default?: string;
   format?: string;
   maxLength?: number;
@@ -124,12 +128,14 @@ export type BlockIOStringSubSchema = BlockIOSubSchemaMeta & {
 
 export type BlockIONumberSubSchema = BlockIOSubSchemaMeta & {
   type: "integer" | "number";
+  const?: number;
   default?: number;
   secret?: boolean;
 };
 
 export type BlockIOBooleanSubSchema = BlockIOSubSchemaMeta & {
   type: "boolean";
+  const?: boolean;
   default?: boolean;
   secret?: boolean;
 };
@@ -196,12 +202,16 @@ export type BlockIOCredentialsSubSchema = BlockIOObjectSubSchema & {
 
 export type BlockIONullSubSchema = BlockIOSubSchemaMeta & {
   type: "null";
+  const?: null;
   secret?: boolean;
 };
 
 // At the time of writing, combined schemas only occur on the first nested level in a
 // block schema. It is typed this way to make the use of these objects less tedious.
-type BlockIOCombinedTypeSubSchema = BlockIOSubSchemaMeta & { type: never } & (
+type BlockIOCombinedTypeSubSchema = BlockIOSubSchemaMeta & {
+  type: never;
+  const: never;
+} & (
     | {
         allOf: [BlockIOSimpleTypeSubSchema];
         secret?: boolean;
@@ -211,13 +221,26 @@ type BlockIOCombinedTypeSubSchema = BlockIOSubSchemaMeta & { type: never } & (
         default?: string | number | boolean | null;
         secret?: boolean;
       }
-    | {
-        oneOf: BlockIOSimpleTypeSubSchema[];
-        default?: string | number | boolean | null;
-        secret?: boolean;
-      }
+    | BlockIOOneOfSubSchema
+    | BlockIODiscriminatedOneOfSubSchema
   );
 
+export type BlockIOOneOfSubSchema = {
+  oneOf: BlockIOSimpleTypeSubSchema[];
+  default?: string | number | boolean | null;
+  secret?: boolean;
+};
+
+export type BlockIODiscriminatedOneOfSubSchema = {
+  oneOf: BlockIOObjectSubSchema[];
+  discriminator: {
+    propertyName: string;
+    mapping: Record<string, BlockIOObjectSubSchema>;
+  };
+  default?: Record<string, any>;
+  secret?: boolean;
+};
+
 /* Mirror of backend/data/graph.py:Node */
 export type Node = {
   id: string;