diff --git a/.changeset/polite-crabs-rush.md b/.changeset/polite-crabs-rush.md
new file mode 100644
index 000000000..dd96cee72
--- /dev/null
+++ b/.changeset/polite-crabs-rush.md
@@ -0,0 +1,5 @@
+---
+"@aws-amplify/data-schema": minor
+---
+
+Allow CustomType and RefType in arguments for custom operations
diff --git a/packages/data-schema/__tests__/CustomOperations.test-d.ts b/packages/data-schema/__tests__/CustomOperations.test-d.ts
index 02ee314ae..f81306d9a 100644
--- a/packages/data-schema/__tests__/CustomOperations.test-d.ts
+++ b/packages/data-schema/__tests__/CustomOperations.test-d.ts
@@ -16,6 +16,7 @@ import type {
 import { configure } from '../src/ModelSchema';
 import { Nullable } from '../src/ModelField';
 import { defineFunctionStub } from './utils';
+import type { CustomOperation } from '../src/CustomOperation';
 
 describe('custom operations return types', () => {
   describe('when .ref() a basic custom type', () => {
@@ -855,3 +856,153 @@ describe('.for() modifier', () => {
     a.subscription().for(a.ref('Model'));
   });
 });
+
+describe('.arguments() modifier', () => {
+  // Test to verify that CustomType can be used as an argument in custom operations
+  it('accepts CustomType in arguments', () => {
+    type ExpectedType = CustomOperation<
+      any,
+      'arguments' | 'for',
+      'queryCustomOperation'
+    >;
+
+    const operation = a.query().arguments({
+      customArg: a.customType({
+        field1: a.string(),
+        field2: a.integer(),
+      }),
+    });
+
+    type test = Expect<Equal<ExpectedType, typeof operation>>;
+  });
+
+  // Test to verify that RefType can be used as an argument in custom operations
+  it('accepts RefType in arguments', () => {
+    type ExpectedType = CustomOperation<
+      any,
+      'arguments' | 'for',
+      'queryCustomOperation'
+    >;
+
+    const operation = a.query().arguments({
+      refArg: a.ref('SomeType'),
+    });
+
+    type test = Expect<Equal<ExpectedType, typeof operation>>;
+  });
+
+  it('handles deeply nested custom types', () => {
+    const schema = a.schema({
+      DeepNested: a.customType({
+        level1: a.customType({
+          level2: a.customType({
+            level3: a.string(),
+          }),
+        }),
+      }),
+      deepQuery: a
+        .query()
+        .arguments({
+          input: a.ref('DeepNested'),
+        })
+        .returns(a.string()),
+    });
+
+    type Schema = ClientSchema<typeof schema>;
+
+    type ExpectedArgs = {
+      input?: {
+        level1?: {
+          level2?: {
+            level3?: string | null;
+          } | null;
+        } | null;
+      } | null;
+    };
+
+    type ActualArgs = Schema['deepQuery']['args'];
+
+    type Test = Expect<Equal<ActualArgs, ExpectedArgs>>;
+  });
+
+  it('handles mixed custom types and refs', () => {
+    const schema = a.schema({
+      RefType: a.customType({
+        field: a.string(),
+      }),
+      MixedType: a.customType({
+        nested: a.customType({
+          refField: a.ref('RefType'),
+          customField: a.customType({
+            deepField: a.integer(),
+          }),
+        }),
+      }),
+      mixedQuery: a
+        .query()
+        .arguments({
+          input: a.ref('MixedType'),
+        })
+        .returns(a.string()),
+    });
+
+    type Schema = ClientSchema<typeof schema>;
+
+    type ExpectedArgs = {
+      input?: {
+        nested?: {
+          refField?: {
+            field?: string | null;
+          } | null;
+          customField?: {
+            deepField?: number | null;
+          } | null;
+        } | null;
+      } | null;
+    };
+
+    type ActualArgs = Schema['mixedQuery']['args'];
+
+    type Test = Expect<Equal<ActualArgs, ExpectedArgs>>;
+  });
+
+  it('handles RefType with multi-layered custom types in nested structures', () => {
+    const schema = a.schema({
+      NestedCustomType: a.customType({
+        nestedField: a.string(),
+      }),
+      RefType: a.customType({
+        field: a.string(),
+        nestedCustom: a.ref('NestedCustomType'),
+      }),
+      OuterType: a.customType({
+        refField: a.ref('RefType'),
+        otherField: a.integer(),
+      }),
+      complexQuery: a
+        .query()
+        .arguments({
+          input: a.ref('OuterType'),
+        })
+        .returns(a.string()),
+    });
+
+    type Schema = ClientSchema<typeof schema>;
+
+    type ExpectedArgs = {
+      input?: {
+        refField?: {
+          field?: string | null;
+          nestedCustom?: {
+            nestedField?: string | null;
+          } | null;
+        } | null;
+        otherField?: number | null;
+      } | null;
+    };
+
+    type ActualArgs = Schema['complexQuery']['args'];
+
+    type Test = Expect<Equal<ActualArgs, ExpectedArgs>>;
+  });
+});
diff --git a/packages/data-schema/__tests__/CustomOperations.test.ts b/packages/data-schema/__tests__/CustomOperations.test.ts
index ca9820776..b342a6a20 100644
--- a/packages/data-schema/__tests__/CustomOperations.test.ts
+++ b/packages/data-schema/__tests__/CustomOperations.test.ts
@@ -1089,6 +1089,195 @@ describe('CustomOperation transform', () => {
     });
   });
 
+  describe('custom operations with custom types and refs', () => {
+    test('Schema with custom query using custom type argument', () => {
+      const s = a
+        .schema({
+          CustomArgType: a.customType({
+            field: a.string(),
+          }),
+          queryWithCustomTypeArg: a
+            .query()
+            .arguments({ customArg: a.ref('CustomArgType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom query using ref argument', () => {
+      const s = a
+        .schema({
+          RefArgType: a.customType({
+            field: a.string(),
+          }),
+          queryWithRefArg: a
+            .query()
+            .arguments({ refArg: a.ref('RefArgType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom mutation using custom type argument', () => {
+      const s = a
+        .schema({
+          CustomArgType: a.customType({
+            field: a.string(),
+          }),
+          mutateWithCustomTypeArg: a
+            .mutation()
+            .arguments({ customArg: a.ref('CustomArgType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom mutation using ref argument', () => {
+      const s = a
+        .schema({
+          RefArgType: a.customType({
+            field: a.string(),
+          }),
+          mutationWithRefArg: a
+            .mutation()
+            .arguments({ refArg: a.ref('RefArgType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom mutation using enum ref in nested custom type', () => {
+      const s = a
+        .schema({
+          NestedCustomType: a.customType({
+            field: a.string(),
+            enumField: a.ref('SomeEnum'),
+          }),
+          SomeEnum: a.enum(['VALUE1', 'VALUE2']),
+          mutateWithNestedEnumRef: a
+            .mutation()
+            .arguments({ nestedArg: a.ref('NestedCustomType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom query using enum ref as argument', () => {
+      const s = a
+        .schema({
+          SomeEnum: a.enum(['OPTION1', 'OPTION2']),
+          queryWithEnumRefArg: a
+            .query()
+            .arguments({ enumArg: a.ref('SomeEnum') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom operation using enum directly in arguments', () => {
+      const s = a
+        .schema({
+          directEnumQuery: a
+            .query()
+            .arguments({
+              status: a.enum(['PENDING', 'APPROVED', 'REJECTED']),
+            })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+      expect(result).toContain('enum DirectEnumQueryStatus {');
+      expect(result).toContain(
+        'directEnumQuery(status: DirectEnumQueryStatus): String',
+      );
+    });
+
+    test('Schema with custom operation using enum in inline custom type arguments', () => {
+      const s = a
+        .schema({
+          inlineEnumQuery: a
+            .query()
+            .arguments({
+              input: a.customType({
+                name: a.string(),
+                status: a.enum(['ACTIVE', 'INACTIVE']),
+              }),
+            })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+      expect(result).toContain('enum InlineEnumQueryInputStatus {');
+      expect(result).toContain('input InlineEnumQueryInput {');
+      expect(result).toContain(
+        'inlineEnumQuery(input: InlineEnumQueryInput): String',
+      );
+    });
+
+    test('Schema with custom operation using enum in referenced custom type arguments', () => {
+      const s = a
+        .schema({
+          StatusInput: a.customType({
+            name: a.string(),
+            status: a.enum(['NEW', 'IN_PROGRESS', 'COMPLETED']),
+          }),
+          refEnumQuery: a
+            .query()
+            .arguments({
+              input: a.ref('StatusInput'),
+            })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+      expect(result).toContain('enum StatusInputStatus {');
+      expect(result).toContain('input StatusInput {');
+      expect(result).toContain('refEnumQuery(input: StatusInput): String');
+    });
+  });
+
   const fakeSecret = () => ({}) as any;
 
   const datasourceConfigMySQL = {
diff --git a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
index 98d9201fc..0e4696ae3 100644
--- a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
@@ -1,5 +1,162 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom mutation using custom type argument 1`] = `
+"input MutateWithCustomTypeArgCustomArgInput
+{
+  field: String
+}
+
+type CustomArgType 
+{
+  field: String
+}
+
+type Mutation {
+  mutateWithCustomTypeArg(customArg: MutateWithCustomTypeArgCustomArgInput): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom mutation using enum ref in nested custom type 1`] = `
+"enum SomeEnum {
+  VALUE1
+  VALUE2
+}
+
+input MutateWithNestedEnumRefNestedArgInput
+{
+  field: String
+  enumField: SomeEnum
+}
+
+type NestedCustomType 
+{
+  field: String
+  enumField: SomeEnum
+}
+
+type Mutation {
+  mutateWithNestedEnumRef(nestedArg: MutateWithNestedEnumRefNestedArgInput): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom mutation using ref argument 1`] = `
+"input MutationWithRefArgRefArgInput
+{
+  field: String
+}
+
+type RefArgType 
+{
+  field: String
+}
+
+type Mutation {
+  mutationWithRefArg(refArg: MutationWithRefArgRefArgInput): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom operation using enum directly in arguments 1`] = `
+"enum DirectEnumQueryStatus {
+  PENDING
+  APPROVED
+  REJECTED
+}
+
+type Query {
+  directEnumQuery(status: DirectEnumQueryStatus): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom operation using enum in inline custom type arguments 1`] = `
+"input InlineEnumQueryInputInput
+{
+  name: String
+  status: InlineEnumQueryInputInputStatus
+}
+
+enum InlineEnumQueryInputInputStatus {
+  ACTIVE
+  INACTIVE
+}
+
+type Query {
+  inlineEnumQuery(input: InlineEnumQueryInputInput): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom operation using enum in referenced custom type arguments 1`] = `
+"input RefEnumQueryInputInput
+{
+  name: String
+  status: RefEnumQueryInputInputStatus
+}
+
+enum RefEnumQueryInputInputStatus {
+  NEW
+  IN_PROGRESS
+  COMPLETED
+}
+
+type StatusInput 
+{
+  name: String
+  status: StatusInputStatus
+}
+
+enum StatusInputStatus {
+  NEW
+  IN_PROGRESS
+  COMPLETED
+}
+
+type Query {
+  refEnumQuery(input: RefEnumQueryInputInput): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom query using custom type argument 1`] = `
+"input QueryWithCustomTypeArgCustomArgInput
+{
+  field: String
+}
+
+type CustomArgType 
+{
+  field: String
+}
+
+type Query {
+  queryWithCustomTypeArg(customArg: QueryWithCustomTypeArgCustomArgInput): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom query using enum ref as argument 1`] = `
+"enum SomeEnum {
+  OPTION1
+  OPTION2
+}
+
+type Query {
+  queryWithEnumRefArg(enumArg: SomeEnum): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
+exports[`CustomOperation transform custom operations with custom types and refs Schema with custom query using ref argument 1`] = `
+"input QueryWithRefArgRefArgInput
+{
+  field: String
+}
+
+type RefArgType 
+{
+  field: String
+}
+
+type Query {
+  queryWithRefArg(refArg: QueryWithRefArgRefArgInput): String @function(name: "myFunc") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
 exports[`CustomOperation transform dynamo schema Custom Mutation w required arg and enum 1`] = `
 "type Post @model @auth(rules: [{allow: private}])
 {
@@ -45,7 +202,7 @@ exports[`CustomOperation transform dynamo schema Custom mutation w inline boolea
 `;
 
 exports[`CustomOperation transform dynamo schema Custom mutation w inline custom return type 1`] = `
-"type LikePostReturnType 
+"type LikePostReturnType
 {
   stringField: String
   intField: Int
@@ -109,7 +266,7 @@ type Query {
 `;
 
 exports[`CustomOperation transform dynamo schema Custom query w inline custom return type 1`] = `
-"type GetPostDetailsReturnType 
+"type GetPostDetailsReturnType
 {
   stringField: String
   intField: Int
diff --git a/packages/data-schema/src/ClientSchema/Core/ClientCustomOperations.ts b/packages/data-schema/src/ClientSchema/Core/ClientCustomOperations.ts
index 6f7d3dfac..f96d87e18 100644
--- a/packages/data-schema/src/ClientSchema/Core/ClientCustomOperations.ts
+++ b/packages/data-schema/src/ClientSchema/Core/ClientCustomOperations.ts
@@ -9,8 +9,8 @@ import type { AppSyncResolverHandler } from 'aws-lambda';
 import type { CustomType } from '../../CustomType';
 import type { FieldTypesOfCustomType } from '../../MappedTypes/ResolveSchema';
 import type { ResolveRef } from '../utilities/ResolveRef';
-import type { EnumType } from '../../EnumType';
 import { ClientSchemaProperty } from './ClientSchemaProperty';
+import type { ResolveFields } from '../utilities';
 
 type CustomOperationSubType<Op extends CustomOperationParamShape> =
   `custom${Op['typeName']}`;
@@ -38,7 +38,7 @@ export interface ClientCustomOperation<
    * ```
    */
   functionHandler: AppSyncResolverHandler<
-    CustomOpArguments<Op>,
+    CustomOpArguments<Op, RefBag>,
     // If the final handler is an async function, the Schema['fieldname']['functionhandler']
     // should have a return type of `void`. This only applies to `functionHandler` and not
     // `returnType` because `returnType` determines the type returned by the mutation / query
@@ -60,7 +60,7 @@ export interface ClientCustomOperation<
    * }
    * ```
    */
-  args: CustomOpArguments<Op>;
+  args: CustomOpArguments<Op, RefBag>;
 
   /**
    * The return type expected by a lambda function handler.
@@ -84,19 +84,18 @@ export interface ClientCustomOperation<
 
 /**
  * Digs out custom operation arguments, mapped to the intended graphql types.
+ * using the existing ResolveFields utility type. This handles:
+ * - Basic scalar fields
+ * - Enum types
+ * - Custom types (including nested structures)
+ * - Reference types
  */
-type CustomOpArguments<Shape extends CustomOperationParamShape> =
-  Shape['arguments'] extends null
-    ? never
-    : ResolveFieldRequirements<{
-        [FieldName in keyof Shape['arguments']]: Shape['arguments'][FieldName] extends BaseModelField<
-          infer R
-        >
-          ? R
-          : Shape['arguments'][FieldName] extends EnumType<infer Values>
-            ? Values[number] | null
-            : never;
-      }>;
+type CustomOpArguments<
+  Shape extends CustomOperationParamShape,
+  RefBag extends Record<string, any> = any,
+> = Shape['arguments'] extends null
+  ? never
+  : ResolveFields<RefBag, Shape['arguments']>;
 
 /**
  * Removes `null | undefined` from the return type if the operation is a subscription,
diff --git a/packages/data-schema/src/ClientSchema/utilities/ResolveField.ts b/packages/data-schema/src/ClientSchema/utilities/ResolveField.ts
index 8da96ff0a..502126bd3 100644
--- a/packages/data-schema/src/ClientSchema/utilities/ResolveField.ts
+++ b/packages/data-schema/src/ClientSchema/utilities/ResolveField.ts
@@ -18,7 +18,7 @@ import { LazyLoader } from '../../runtime';
  * The first type parameter (`Bag`) should always just be the top-level `ClientSchema` that
  * references and related model definitions can be resolved against.
  */
-export type ResolveFields<Bag extends Record<string, any>, T> = ShallowPretty<
+export type ResolveFields<Bag extends Record<string, any>, T> = Expand<
   {
     [K in keyof T as IsRequired<T[K]> extends true
       ? K
@@ -34,9 +34,7 @@ export type ResolveFields<Bag extends Record<string, any>, T> = ShallowPretty<
 // down the line *as-needed*. Performing this *here* is somehow essential to getting 2 unit
 // tests to pass, but hurts performance significantly. E.g., p50/operations/p50-prod-CRUDL.bench.ts
 // goes from `783705` to `1046408`.
-type ShallowPretty<T> = {
-  [K in keyof T]: T[K];
-};
+type Expand<T> = T extends infer O ? { [K in keyof O]: O[K] } : never;
 
 export type ResolveIndividualField<Bag extends Record<string, any>, T> =
   T extends BaseModelField<infer FieldShape>
diff --git a/packages/data-schema/src/CustomOperation.ts b/packages/data-schema/src/CustomOperation.ts
index 0e5fdb5ec..4f79fe524 100644
--- a/packages/data-schema/src/CustomOperation.ts
+++ b/packages/data-schema/src/CustomOperation.ts
@@ -29,7 +29,7 @@ type CustomOperationBrand =
   | typeof subscriptionBrand
   | typeof generationBrand;
 
-type CustomArguments = Record<string, BaseModelField | EnumType>;
+type CustomArguments = Record<string, BaseModelField | EnumType | CustomType<any> | RefType<any, any>>;
 type SubscriptionSource = RefType<any, any>;
 type InternalSubscriptionSource = InternalRef;
 type CustomReturnType = RefType<any> | CustomType<any>;
diff --git a/packages/data-schema/src/SchemaProcessor.ts b/packages/data-schema/src/SchemaProcessor.ts
index abc3ad883..8857b44cb 100644
--- a/packages/data-schema/src/SchemaProcessor.ts
+++ b/packages/data-schema/src/SchemaProcessor.ts
@@ -344,6 +344,8 @@ function customOperationToGql(
 ): {
   gqlField: string;
   implicitTypes: [string, any][];
+  inputTypes: string[];
+  returnTypes: string[];
   customTypeAuthRules: CustomTypeAuthRules;
   lambdaFunctionDefinition: LambdaFunctionDefinition;
   customSqlDataSourceStrategy: CustomSqlDataSourceStrategy | undefined;
@@ -358,6 +360,8 @@ function customOperationToGql(
 
   let callSignature: string = typeName;
   const implicitTypes: [string, any][] = [];
+  const inputTypes: string[] = [];
+  let returnTypes: string[] = [];
 
   // When Custom Operations are defined with a Custom Type return type,
   // the Custom Type inherits the operation's auth rules
@@ -393,6 +397,8 @@ function customOperationToGql(
           typeName: returnType.data.link,
           authRules: authorization,
         };
+      } else if (type === 'Enum') {
+        return refFieldToGql(returnType.data);
       }
 
       return refFieldToGql(returnType?.data);
@@ -404,7 +410,10 @@ function customOperationToGql(
           authRules: authorization,
         };
 
-        implicitTypes.push([returnTypeName, returnType]);
+        implicitTypes.push([
+          returnTypeName,
+          { ...returnType, generateInputType: false },
+        ]);
       }
       return returnTypeName;
     } else if (isEnumType(returnType)) {
@@ -441,6 +450,16 @@ function customOperationToGql(
       refererTypeName: typeName,
     });
   }
+  // After resolving returnTypeName
+  if (isCustomType(returnType)) {
+    returnTypes = generateInputTypes(
+      [[returnTypeName, { ...returnType, isgenerateInputTypeInput: false }]],
+      false,
+      getRefType,
+      authorization,
+    );
+  }
+  const dedupedInputTypes = new Set<string>(inputTypes);
 
   if (Object.keys(fieldArgs).length > 0) {
     const { gqlFields, implicitTypes: implied } = processFields(
@@ -448,9 +467,22 @@ function customOperationToGql(
       fieldArgs,
       {},
       {},
+      getRefType,
+      undefined,
+      undefined,
+      {},
+      databaseType === 'sql' ? 'postgresql' : 'dynamodb',
+      true,
     );
     callSignature += `(${gqlFields.join(', ')})`;
     implicitTypes.push(...implied);
+
+    const newTypes = generateInputTypes(implied, true, getRefType);
+    for (const t of newTypes) {
+      if (!dedupedInputTypes.has(t)) {
+        dedupedInputTypes.add(t);
+      }
+    }
   }
 
   const handler = handlers && handlers[0];
@@ -547,6 +579,8 @@ function customOperationToGql(
   return {
     gqlField,
     implicitTypes: implicitTypes,
+    inputTypes,
+    returnTypes,
     customTypeAuthRules,
     lambdaFunctionDefinition,
     customSqlDataSourceStrategy,
@@ -974,15 +1008,19 @@ function processFields(
   fields: Record<string, any>,
   impliedFields: Record<string, any>,
   fieldLevelAuthRules: Record<string, string | null>,
+  getRefType: ReturnType<typeof getRefTypeForSchema>,
+
   identifier?: readonly string[],
   partitionKey?: string,
   secondaryIndexes: TransformedSecondaryIndexes = {},
   databaseEngine: DatasourceEngine = 'dynamodb',
+  generateInputType: boolean = false,
 ) {
   const gqlFields: string[] = [];
   // stores nested, field-level type definitions (custom types and enums)
   // the need to be hoisted to top-level schema types and processed accordingly
   const implicitTypes: [string, any][] = [];
+  const gqlComponents: string[] = [];
 
   validateImpliedFields(fields, impliedFields);
   validateDBGeneration(fields, databaseEngine);
@@ -1001,50 +1039,115 @@ function processFields(
       if (fieldName === partitionKey) {
         gqlFields.push(
           `${fieldName}: ${scalarFieldToGql(
-            fieldDef.data,
+            (fieldDef as any).data,
             identifier,
             secondaryIndexes[fieldName],
           )}${fieldAuth}`,
         );
       } else if (isRefField(fieldDef)) {
-        gqlFields.push(
-          `${fieldName}: ${refFieldToGql(fieldDef.data, secondaryIndexes[fieldName])}${fieldAuth}`,
-        );
+        const refTypeInfo = getRefType(fieldDef.data.link, typeName);
+        if (refTypeInfo.type === 'Enum') {
+          gqlFields.push(
+            `${fieldName}: ${refFieldToGql(fieldDef.data, secondaryIndexes[fieldName])}${fieldAuth}`,
+          );
+        } else if (refTypeInfo.type === 'CustomType') {
+          if (generateInputType) {
+            const inputTypeName = `${capitalize(typeName)}${capitalize(fieldName)}Input`;
+            gqlFields.push(`${fieldName}: ${inputTypeName}${fieldAuth}`);
+
+            const { implicitTypes: nestedImplicitTypes } = processFields(
+              inputTypeName,
+              refTypeInfo.def.data.fields,
+              {},
+              {},
+              getRefType,
+              undefined,
+              undefined,
+              {},
+              'dynamodb',
+              true,
+            );
+
+            implicitTypes.push([
+              inputTypeName,
+              {
+                data: {
+                  type: 'customType',
+                  fields: refTypeInfo.def.data.fields,
+                },
+              },
+            ]);
+
+            implicitTypes.push(...nestedImplicitTypes);
+          } else {
+            gqlFields.push(
+              `${fieldName}: ${refFieldToGql(fieldDef.data, secondaryIndexes[fieldName])}${fieldAuth}`,
+            );
+          }
+        } else {
+          throw new Error(
+            `Field '${fieldName}' in type '${typeName}' references '${fieldDef.data.link}' which is neither a CustomType nor an Enum.`,
+          );
+        }
       } else if (isEnumType(fieldDef)) {
         // The inline enum type name should be `<TypeName><FieldName>` to avoid
         // enum type name conflicts
         const enumName = `${capitalize(typeName)}${capitalize(fieldName)}`;
-
         implicitTypes.push([enumName, fieldDef]);
-
+        const enumValues = (fieldDef as any).values;
+        if (Array.isArray(enumValues)) {
+          gqlComponents.push(
+            `enum ${enumName} {\n  ${enumValues.join('\n  ')}\n}`,
+          );
+        }
         gqlFields.push(
           `${fieldName}: ${enumFieldToGql(enumName, secondaryIndexes[fieldName])}`,
         );
       } else if (isCustomType(fieldDef)) {
         // The inline CustomType name should be `<TypeName><FieldName>` to avoid
         // CustomType name conflicts
-        const customTypeName = `${capitalize(typeName)}${capitalize(
-          fieldName,
-        )}`;
-
+        const customTypeName = `${capitalize(typeName)}${capitalize(fieldName)}${generateInputType ? 'Input' : ''}`;
         implicitTypes.push([customTypeName, fieldDef]);
+        // Recursively process the fields of the nested custom type
+        const {
+          gqlFields: _nestedFields,
+          implicitTypes: nestedImplicitTypes,
+          gqlComponents: nestedComponents,
+        } = processFields(
+          customTypeName,
+          fieldDef.data.fields,
+          {},
+          {},
+          getRefType,
+          undefined,
+          undefined,
+          {},
+          databaseEngine,
+          generateInputType,
+        );
 
-        gqlFields.push(`${fieldName}: ${customTypeName}`);
+        // Add nested components (including enums) to gqlComponents
+        gqlComponents.push(...nestedComponents);
+        // Add nested implicit types to the main implicitTypes array
+        implicitTypes.push(...nestedImplicitTypes);
+        gqlFields.push(`${fieldName}: ${customTypeName}${fieldAuth}`);
       } else {
         gqlFields.push(
           `${fieldName}: ${scalarFieldToGql(
-            (fieldDef as any).data,
+            fieldDef.data,
             undefined,
             secondaryIndexes[fieldName],
           )}${fieldAuth}`,
         );
       }
     } else {
-      throw new Error(`Unexpected field definition: ${fieldDef}`);
+      throw new Error(
+        `Unexpected field definition for ${typeName}.${fieldName}: ${JSON.stringify(fieldDef)}`,
+      );
     }
   }
 
-  return { gqlFields, implicitTypes };
+  return { gqlFields, implicitTypes, gqlComponents };
 }
 
 type TransformedSecondaryIndexes = {
@@ -1316,6 +1419,75 @@ const mergeCustomTypeAuthRules = (
   }
 };
 
+function generateInputTypes(
+  implicitTypes: [string, any][],
+  generateInputType: boolean,
+  getRefType: ReturnType<typeof getRefTypeForSchema>,
+  authRules?: Authorization<any, any, any>[],
+  isInlineType = false,
+  definedTypes: Set<string> = new Set(),
+): string[] {
+  const generatedTypes = new Set<string>();
+
+  implicitTypes.forEach(([typeName, typeDef]) => {
+    if (!definedTypes.has(typeName)) {
+      if (isCustomType(typeDef)) {
+        const { gqlFields, implicitTypes: nestedTypes } = processFields(
+          typeName,
+          typeDef.data.fields,
+          {},
+          {},
+          getRefType,
+          undefined,
+          undefined,
+          {},
+          'dynamodb',
+          generateInputType,
+        );
+        const authString =
+          !isInlineType && authRules
+            ? mapToNativeAppSyncAuthDirectives(authRules, false).authString
+            : '';
+        const typeKeyword = generateInputType ? 'input' : 'type';
+        const customType = `${typeKeyword} ${typeName}${authString ? ` ${authString}` : ''}\n{\n  ${gqlFields.join('\n  ')}\n}`;
+        generatedTypes.add(customType);
+        definedTypes.add(typeName);
+
+        // Process nested types
+        if (nestedTypes.length > 0) {
+          const nestedGeneratedTypes = generateInputTypes(
+            nestedTypes,
+            generateInputType,
+            getRefType,
+            authRules,
+            false,
+            definedTypes,
+          );
+          nestedGeneratedTypes.forEach((type) => {
+            generatedTypes.add(type);
+          });
+        }
+      } else if (typeDef.type === 'enum') {
+        if (!definedTypes.has(typeName)) {
+          generatedTypes.add(
+            `enum ${typeName} {\n  ${typeDef.values.join('\n  ')}\n}`,
+          );
+          definedTypes.add(typeName);
+        }
+      } else if (typeDef?.data?.type === 'ref') {
+        getRefType(typeDef.data.link, typeName);
+      } else if (typeDef.type === 'scalar') {
+        generatedTypes.add(`scalar ${typeName}`);
+        definedTypes.add(typeName);
+      } else {
+        console.warn(`Unexpected type definition for ${typeName}:`, typeDef);
+      }
+    }
+  });
+
+  return Array.from(generatedTypes);
+}
+
 const schemaPreprocessor = (
   schema: InternalSchema,
 ): {
@@ -1325,12 +1497,13 @@ const schemaPreprocessor = (
   lambdaFunctions: LambdaFunctionDefinition;
   customSqlDataSourceStrategies?: CustomSqlDataSourceStrategy[];
 } => {
-  const gqlModels: string[] = [];
-
+  const gqlComponents: string[] = [];
   const customQueries = [];
   const customMutations = [];
   const customSubscriptions = [];
   let shouldAddConversationTypes = false;
+  const definedTypes = new Set<string>();
+  const nestedTypeDefinitions: string[] = [];
 
   // Dict of auth rules to be applied to custom types
   // Inherited from the auth configured on the custom operations that return these custom types
@@ -1382,17 +1555,16 @@ const schemaPreprocessor = (
         : schemaAuth;
 
     if (!isInternalModel(typeDef)) {
-      if (isEnumType(typeDef)) {
+      if (isEnumType(typeDef) && !definedTypes.has(typeName)) {
         if (typeDef.values.some((value) => /\s/.test(value))) {
           throw new Error(
             `Values of the enum type ${typeName} should not contain any whitespace.`,
           );
         }
-        const enumType = `enum ${typeName} {\n  ${typeDef.values.join(
-          '\n  ',
-        )}\n}`;
-        gqlModels.push(enumType);
-      } else if (isCustomType(typeDef)) {
+        const enumType = `enum ${typeName} {\n  ${typeDef.values.join('\n  ')}\n}`;
+        gqlComponents.push(enumType);
+        definedTypes.add(typeName);
+      } else if (isCustomType(typeDef) && !definedTypes.has(typeName)) {
         const fields = typeDef.data.fields;
 
         validateRefUseCases(typeName, 'customType', fields, getRefType);
@@ -1420,23 +1592,65 @@ const schemaPreprocessor = (
           authFields,
         );
 
-        const { gqlFields, implicitTypes } = processFields(
+        const {
+          gqlFields,
+          implicitTypes,
+          gqlComponents: _nestedComponents,
+        } = processFields(
           typeName,
           fields,
           authFields,
           fieldLevelAuthRules,
-          undefined,
+          getRefType,
           undefined,
           undefined,
           databaseEngine,
         );
-
-        topLevelTypes.push(...implicitTypes);
-
         const joined = gqlFields.join('\n  ');
 
         const model = `type ${typeName} ${customAuth}\n{\n  ${joined}\n}`;
-        gqlModels.push(model);
+        gqlComponents.push(model);
+        definedTypes.add(typeName);
+
+        /**
+         * Process implicit types found during schema generation.
+         * This section handles two cases:
+         * 1. Enum types - Creates GraphQL enum definitions
+         * 2. Custom types - Creates GraphQL type definitions with inherited auth rules
+         *
+         * Uses definedTypes set to prevent duplicate type definitions.
+         */
+
+        implicitTypes.forEach(([nestedTypeName, nestedTypeDef]) => {
+          if (!definedTypes.has(nestedTypeName)) {
+            if (isEnumType(nestedTypeDef)) {
+              nestedTypeDefinitions.push(
+                `enum ${nestedTypeName} {\n  ${nestedTypeDef.values.join('\n  ')}\n}`,
+              );
+            } else if (isCustomType(nestedTypeDef)) {
+              const nestedAuth = customTypeInheritedAuthRules[nestedTypeName]
+                ? mapToNativeAppSyncAuthDirectives(
+                    customTypeInheritedAuthRules[nestedTypeName],
+                    false,
+                  ).authString
+                : '';
+              const { gqlFields: nestedFields } = processFields(
+                nestedTypeName,
+                nestedTypeDef.data.fields,
+                {},
+                {},
+                getRefType,
+                undefined,
+                undefined,
+                {},
+                databaseEngine,
+              );
+              const nestedModel = `type ${nestedTypeName} ${nestedAuth}\n{\n  ${nestedFields.join('\n  ')}\n}`;
+              nestedTypeDefinitions.push(nestedModel);
+            }
+            definedTypes.add(nestedTypeName);
+          }
+        });
       } else if (isCustomOperation(typeDef)) {
         // TODO: add generation route logic.
 
@@ -1445,6 +1659,8 @@ const schemaPreprocessor = (
         const {
           gqlField,
           implicitTypes,
+          inputTypes: operationInputTypes,
+          returnTypes: operationReturnTypes,
           customTypeAuthRules,
           jsFunctionForField,
           lambdaFunctionDefinition,
@@ -1456,8 +1672,48 @@ const schemaPreprocessor = (
           databaseType,
           getRefType,
         );
-
-        topLevelTypes.push(...implicitTypes);
+        operationInputTypes.forEach((type) => {
+          gqlComponents.push(type);
+        });
+        gqlComponents.push(...operationReturnTypes);
+
+        /**
+         * Processes implicit types to generate GraphQL definitions.
+         *
+         * - Enums are converted to 'enum' definitions and added to the schema.
+         * - Custom types are conditionally treated as input types if they are
+         *   not part of the operation's return types.
+         * - Input types are generated and added to the inputTypes array when required.
+         *
+         * This ensures that all necessary type definitions, including enums and input types,
+         * are correctly generated and available in the schema output.
+         */
+
+        implicitTypes.forEach(([typeName, typeDef]) => {
+          if (!definedTypes.has(typeName)) {
+            if (isEnumType(typeDef)) {
+              gqlComponents.push(
+                `enum ${typeName} {\n  ${typeDef.values.join('\n  ')}\n}`,
+              );
+              definedTypes.add(typeName);
+            } else if (isCustomType(typeDef) || typeDef?.data?.type === 'ref') {
+              const isReturnType = operationReturnTypes.some((returnType) =>
+                returnType.includes(typeName),
+              );
+              if (!isReturnType) {
+                const generatedTypes = generateInputTypes(
+                  [[typeName, typeDef]],
+                  true,
+                  getRefType,
+                  undefined,
+                  false,
+                  definedTypes,
+                );
+                gqlComponents.push(...generatedTypes);
+              }
+            }
+          }
+        });
 
         mergeCustomTypeAuthRules(
           customTypeInheritedAuthRules,
@@ -1537,12 +1793,19 @@ const schemaPreprocessor = (
         fields,
         authFields,
         fieldLevelAuthRules,
+        getRefType,
         identifier,
         partitionKey,
         undefined,
         databaseEngine,
       );
-
+      const existingTypeNames = new Set<string>(topLevelTypes.map(([n]) => n));
+      for (const [name, type] of implicitTypes) {
+        if (!existingTypeNames.has(name)) {
+          topLevelTypes.push([name, type]);
+          existingTypeNames.add(name);
+        }
+      }
       topLevelTypes.push(...implicitTypes);
 
       const joined = gqlFields.join('\n  ');
@@ -1556,7 +1819,8 @@ const schemaPreprocessor = (
       // passing (timestamps: null) to @model to suppress this behavior as a short
       // term solution.
       const model = `type ${typeName} @model(timestamps: null) ${authString}${refersToString}\n{\n  ${joined}\n}`;
-      gqlModels.push(model);
+      gqlComponents.push(model);
+      definedTypes.add(typeName);
     } else {
       const fields = typeDef.data.fields as Record<string, BaseModelField>;
 
@@ -1602,6 +1866,7 @@ const schemaPreprocessor = (
         fields,
         authFields,
         fieldLevelAuthRules,
+        getRefType,
         identifier,
         partitionKey,
         transformedSecondaryIndexes,
@@ -1618,7 +1883,7 @@ const schemaPreprocessor = (
       const modelDirective = modelAttrs ? `@model(${modelAttrs})` : '@model';
 
       const model = `type ${typeName} ${modelDirective} ${authString}\n{\n  ${joined}\n}`;
-      gqlModels.push(model);
+      gqlComponents.push(model);
     }
   }
 
@@ -1627,13 +1892,14 @@ const schemaPreprocessor = (
     mutations: customMutations,
     subscriptions: customSubscriptions,
   };
+  gqlComponents.push(...nestedTypeDefinitions);
+
+  gqlComponents.push(...generateCustomOperationTypes(customOperations));
 
-  gqlModels.push(...generateCustomOperationTypes(customOperations));
   if (shouldAddConversationTypes) {
-    gqlModels.push(CONVERSATION_SCHEMA_GRAPHQL_TYPES);
+    gqlComponents.push(CONVERSATION_SCHEMA_GRAPHQL_TYPES);
   }
-
-  const processedSchema = gqlModels.join('\n\n');
+  const processedSchema = gqlComponents.join('\n\n');
 
   return {
     schema: processedSchema,
@@ -1932,6 +2198,8 @@ function transformCustomOperations(
   const {
     gqlField,
     implicitTypes,
+    inputTypes,
+    returnTypes,
     customTypeAuthRules,
     lambdaFunctionDefinition,
     customSqlDataSourceStrategy,
@@ -1947,6 +2215,8 @@ function transformCustomOperations(
   return {
     gqlField,
     implicitTypes,
+    inputTypes,
+    returnTypes,
     customTypeAuthRules,
     jsFunctionForField,
     lambdaFunctionDefinition,
@@ -1981,7 +2251,14 @@ function extractNestedCustomTypeNames(
   topLevelTypes: [string, any][],
   getRefType: ReturnType<typeof getRefTypeForSchema>,
 ): string[] {
-  if (!customTypeAuthRules) {
+  if (!customTypeAuthRules || !topLevelTypes || topLevelTypes.length === 0) {
+    return [];
+  }
+  const foundType = topLevelTypes.find(
+    ([topLevelTypeName]) => customTypeAuthRules.typeName === topLevelTypeName,
+  );
+
+  if (!foundType) {
     return [];
   }
 
diff --git a/packages/data-schema/src/runtime/bridge-types.ts b/packages/data-schema/src/runtime/bridge-types.ts
index 93185146c..ffe73e22e 100644
--- a/packages/data-schema/src/runtime/bridge-types.ts
+++ b/packages/data-schema/src/runtime/bridge-types.ts
@@ -16,6 +16,8 @@ import { Observable } from 'rxjs';
 import { CustomHeaders, ModelSortDirection } from './client';
 
 import { AiAction, AiCategory } from './internals/ai/getCustomUserAgentDetails';
+import { CustomType } from '../CustomType';
+import { RefType } from '../RefType';
 
 export declare namespace AmplifyServer {
   export interface ContextToken {
@@ -165,7 +167,7 @@ export type CustomOperationArguments = Record<string, CustomOperationArgument>;
 
 export interface CustomOperationArgument {
   name: string;
-  type: InputFieldType;
+  type: InputFieldType | CustomType<any> | RefType<any, any>;
   isArray: boolean;
   isRequired: boolean;
   isArrayNullable?: boolean;
@@ -240,7 +242,7 @@ export type FieldType =
   | ModelFieldType
   | NonModelFieldType;
 
-export type InputFieldType = ScalarType | EnumType | InputType;
+export type InputFieldType = ScalarType | EnumType | InputType | CustomType<any> | RefType<any, any>;
 
 export type FieldAttribute = ModelAttribute;
 
diff --git a/packages/integration-tests/__tests__/defined-behavior/2-expected-use/custom-operations.ts b/packages/integration-tests/__tests__/defined-behavior/2-expected-use/custom-operations.ts
index d23ec3cf4..f86c03fdb 100644
--- a/packages/integration-tests/__tests__/defined-behavior/2-expected-use/custom-operations.ts
+++ b/packages/integration-tests/__tests__/defined-behavior/2-expected-use/custom-operations.ts
@@ -123,12 +123,202 @@ describe('custom operations', () => {
         a.handler.function(dummyHandler).async(),
       ])
       .authorization((allow) => [allow.publicApiKey()]),
+    CustomArgType: a.customType({
+      message: a.string(),
+      count: a.integer(),
+    }),
+    NestedObjectType: a.customType({
+      innerField1: a.boolean(),
+      innerField2: a.string(),
+    }),
+
+    NestedFieldType: a.customType({
+      nestedObject1: a.ref('NestedObjectType'),
+    }),
+    queryWithCustomTypeArg: a
+      .query()
+      .arguments({
+        customArg: a.ref('CustomArgType'),
+      })
+      .returns(a.string())
+      .handler(a.handler.function(dummyHandler))
+      .authorization((allow) => [allow.publicApiKey()]),
+    mutateWithCustomTypeArg: a
+      .mutation()
+      .arguments({
+        customArg: a.ref('CustomArgType'),
+      })
+      .returns(a.string())
+      .handler(a.handler.function(dummyHandler))
+      .authorization((allow) => [allow.publicApiKey()]),
+    mutationWithNestedCustomType: a
+      .mutation()
+      .arguments({
+        nestedField: a.ref('NestedFieldType'),
+      })
+      .returns(a.string())
+      .handler(a.handler.function(dummyHandler))
+      .authorization((allow) => [allow.publicApiKey()]),
+    queryWithRefArg: a
+      .query()
+      .arguments({
+        refArg: a.ref('EchoResult'),
+      })
+      .returns(a.string())
+      .handler(a.handler.function(dummyHandler))
+      .authorization((allow) => [allow.publicApiKey()]),
+    mutationWithRefArg: a
+      .mutation()
+      .arguments({
+        refArg: a.ref('EchoResult'),
+      })
+      .returns(a.string())
+      .handler(a.handler.function(dummyHandler))
+      .authorization((allow) => [allow.publicApiKey()]),
+    ComplexCustomArgType: a.customType({
+      field1: a.string(),
+      field2: a.integer(),
+    }),
+    complexQueryOperation: a
+      .query()
+      .arguments({
+        scalarArg: a.string(),
+        customArg: a.ref('ComplexCustomArgType'),
+        refArg: a.ref('EchoResult'),
+      })
+      .returns(a.string())
+      .handler(a.handler.function(dummyHandler))
+      .authorization((allow) => [allow.publicApiKey()]),
+    complexMutation: a
+      .mutation()
+      .arguments({
+        scalarArg: a.string(),
+        customArg: a.ref('ComplexCustomArgType'),
+        refArg: a.ref('EchoResult'),
+      })
+      .returns(a.string())
+      .handler(a.handler.function(dummyHandler))
+      .authorization((allow) => [allow.publicApiKey()]),
   });
 
   type Schema = ClientSchema<typeof schema>;
 
+  type ExpectedQueryWithCustomTypeArg = {
+    customArg?: {
+      message?: string | null;
+      count?: number | null;
+    } | null;
+  };
+  type ActualQuertWithCustomTypeArg = Schema['queryWithCustomTypeArg']['args'];
+  type TestEchoWithCustomTypeArg = Expect<
+    Equal<ActualQuertWithCustomTypeArg, ExpectedQueryWithCustomTypeArg>
+  >;
+
+  type ExpectedMutateWithCustomTypeArg = {
+    customArg?: {
+      message?: string | null;
+      count?: number | null;
+    } | null;
+  };
+  type ActualMutateWithCustomTypeArg =
+    Schema['mutateWithCustomTypeArg']['args'];
+  type TestMutateWithCustomTypeArg = Expect<
+    Equal<ActualMutateWithCustomTypeArg, ExpectedMutateWithCustomTypeArg>
+  >;
+
+  type ExpectedNestedCustomTypeArgs = {
+    nestedField?: {
+      nestedObject1?: {
+        innerField1?: boolean | null;
+        innerField2?: string | null;
+      } | null;
+    } | null;
+  };
+  type ActualNestedCustomTypeArgs =
+    Schema['mutationWithNestedCustomType']['args'];
+  type TestNestedCustomTypeArgs = Expect<
+    Equal<ActualNestedCustomTypeArgs, ExpectedNestedCustomTypeArgs>
+  >;
+
+  type ExpectedQueryWithRefArg = {
+    refArg?: {
+      result?: string | null;
+    } | null;
+  };
+  type ActualQueryWithRefArg = Schema['queryWithRefArg']['args'];
+  type TestQueryWithRefArg = Expect<
+    Equal<ActualQueryWithRefArg, ExpectedQueryWithRefArg>
+  >;
+
+  type ExpectedMutationWithRefArg = {
+    refArg?: {
+      result?: string | null;
+    } | null;
+  };
+  type ActualMutationWithRefArg = Schema['mutationWithRefArg']['args'];
+  type TestMutationWithRefArg = Expect<
+    Equal<ActualMutationWithRefArg, ExpectedMutationWithRefArg>
+  >;
+
+  type ExpectedComplexQueryArgs = {
+    scalarArg?: string | null;
+    customArg?: {
+      field1?: string | null;
+      field2?: number | null;
+    } | null;
+    refArg?: {
+      result?: string | null;
+    } | null;
+  };
+  type ActualComplexArgs = Schema['complexQueryOperation']['args'];
+  type TestComplexArgs = Expect<
+    Equal<ActualComplexArgs, ExpectedComplexQueryArgs>
+  >;
+
+  type ExpectedComplexMutationArgs = {
+    scalarArg?: string | null;
+    customArg?: {
+      field1?: string | null;
+      field2?: number | null;
+    } | null;
+    refArg?: {
+      result?: string | null;
+    } | null;
+  };
+  type ActualComplexMutationArgs = Schema['complexMutation']['args'];
+  type TestComplexMutationArgs = Expect<
+    Equal<ActualComplexMutationArgs, ExpectedComplexMutationArgs>
+  >;
   // #endregion
 
+  test('schema.transform() includes custom types, ref types, and operations', () => {
+    const transformedSchema = schema.transform();
+    const expectedTypes = ['CustomArgType', 'EchoResult', 'Query', 'Mutation'];
+    const expectedOperations = [
+      'queryWithCustomTypeArg(customArg: QueryWithCustomTypeArgCustomArgInput): String',
+      'queryWithRefArg(refArg: QueryWithRefArgRefArgInput): String',
+      'mutateWithCustomTypeArg(customArg: MutateWithCustomTypeArgCustomArgInput): String',
+      'mutationWithRefArg(refArg: MutationWithRefArgRefArgInput): String',
+    ];
+    const expectedInputTypes = [
+      'input QueryWithCustomTypeArgCustomArgInput',
+      'input QueryWithRefArgRefArgInput',
+      'input MutateWithCustomTypeArgCustomArgInput',
+      'input MutationWithRefArgRefArgInput',
+    ];
+
+    expectedTypes.forEach((type) => {
+      expect(transformedSchema.schema).toContain(`type ${type}`);
+    });
+
+    expectedOperations.forEach((operation) => {
+      expect(transformedSchema.schema).toContain(operation);
+    });
+
+    expectedInputTypes.forEach((inputType) => {
+      expect(transformedSchema.schema).toContain(inputType);
+    });
+  });
   test('primitive type result', async () => {
     const { spy, generateClient } = mockedGenerateClient([
       {
@@ -141,7 +331,6 @@ describe('custom operations', () => {
     const config = await buildAmplifyConfig(schema);
     Amplify.configure(config);
     const client = generateClient<Schema>();
-
     // #region covers ffefd700b1e323c9
     const { data } = await client.queries.echo({ value: 'something' });
     // #endregion