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 bac164008..65bbd936e 100644
--- a/packages/data-schema/__tests__/CustomOperations.test-d.ts
+++ b/packages/data-schema/__tests__/CustomOperations.test-d.ts
@@ -16,6 +16,9 @@ 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', () => {
@@ -858,3 +861,22 @@ 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', () => {
+    const operation: CustomOperation<any, "arguments" | "for", "queryCustomOperation"> = a.query().arguments({
+      customArg: a.customType({
+        field1: a.string(),
+        field2: a.integer()
+      })
+    });
+  });
+
+  // Test to verify that RefType can be used as an argument in custom operations
+  it('accepts RefType in arguments', () => {
+    const operation:CustomOperation<any, "arguments" | "for", "queryCustomOperation">  = a.query().arguments({
+      refArg: a.ref('SomeType')
+    });
+  });
+});
\ No newline at end of file
diff --git a/packages/data-schema/__tests__/CustomOperations.test.ts b/packages/data-schema/__tests__/CustomOperations.test.ts
index dc6bd7b96..a841ed336 100644
--- a/packages/data-schema/__tests__/CustomOperations.test.ts
+++ b/packages/data-schema/__tests__/CustomOperations.test.ts
@@ -907,7 +907,7 @@ describe('CustomOperation transform', () => {
               .query()
               .arguments({})
               .handler(a.handler.function(fn1).async())
-              .authorization((allow) => allow.authenticated())
+              .authorization((allow) => allow.authenticated()),
           });
 
           const { schema, lambdaFunctions } = s.transform();
@@ -915,7 +915,7 @@ describe('CustomOperation transform', () => {
           expect(lambdaFunctions).toMatchObject({
             FnGetPostDetails: fn1,
           });
-        })
+        });
 
         test('defineFunction sync - async', () => {
           const fn1 = defineFunctionStub({});
@@ -927,7 +927,7 @@ describe('CustomOperation transform', () => {
                 a.handler.function(fn1),
                 a.handler.function(fn1).async(),
               ])
-              .authorization((allow) => allow.authenticated())
+              .authorization((allow) => allow.authenticated()),
           });
 
           const { schema, lambdaFunctions } = s.transform();
@@ -935,7 +935,7 @@ describe('CustomOperation transform', () => {
           expect(lambdaFunctions).toMatchObject({
             FnGetPostDetails: fn1,
           });
-        })
+        });
 
         test('defineFunction sync - async with returns generates type errors', () => {
           const fn1 = defineFunctionStub({});
@@ -949,9 +949,9 @@ describe('CustomOperation transform', () => {
               ])
               .authorization((allow) => allow.authenticated())
               // @ts-expect-error
-              .returns({ })
+              .returns({}),
           });
-        })
+        });
 
         test('defineFunction async - async', () => {
           const fn1 = defineFunctionStub({});
@@ -965,7 +965,7 @@ describe('CustomOperation transform', () => {
                 a.handler.function(fn1).async(),
                 a.handler.function(fn2).async(),
               ])
-              .authorization((allow) => allow.authenticated())
+              .authorization((allow) => allow.authenticated()),
           });
 
           const { schema, lambdaFunctions } = s.transform();
@@ -974,7 +974,7 @@ describe('CustomOperation transform', () => {
             FnGetPostDetails: fn1,
             FnGetPostDetails2: fn2,
           });
-        })
+        });
 
         test('defineFunction async - sync', () => {
           const fn1 = defineFunctionStub({});
@@ -987,12 +987,12 @@ describe('CustomOperation transform', () => {
                 a.handler.function(fn1),
               ])
               .returns(a.customType({}))
-              .authorization((allow) => allow.authenticated())
+              .authorization((allow) => allow.authenticated()),
           });
 
           const { schema, lambdaFunctions } = s.transform();
           expect(schema).toMatchSnapshot();
-        })
+        });
 
         test('pipeline / mix', () => {
           const fn1 = defineFunctionStub({});
@@ -1341,15 +1341,14 @@ describe('custom operations + custom type auth inheritance', () => {
 
   test('implicit custom type inherits auth rules from referencing op', () => {
     const s = a.schema({
+      MyQueryReturnType: a.customType({
+        fieldA: a.string(),
+        fieldB: a.integer(),
+      }),
       myQuery: a
         .query()
         .handler(a.handler.function('myFn'))
-        .returns(
-          a.customType({
-            fieldA: a.string(),
-            fieldB: a.integer(),
-          }),
-        )
+        .returns(a.ref('MyQueryReturnType'))
         .authorization((allow) => allow.publicApiKey()),
     });
 
@@ -1363,23 +1362,22 @@ describe('custom operations + custom type auth inheritance', () => {
 
   test('nested custom types inherit auth rules from top-level referencing op', () => {
     const s = a.schema({
+      MyQueryReturnType: a.customType({
+        fieldA: a.string(),
+        fieldB: a.integer(),
+        nestedCustomType: a.customType({
+          nestedA: a.string(),
+          nestedB: a.string(),
+          grandChild: a.customType({
+            grandA: a.string(),
+            grandB: a.string(),
+          }),
+        }),
+      }),
       myQuery: a
         .query()
         .handler(a.handler.function('myFn'))
-        .returns(
-          a.customType({
-            fieldA: a.string(),
-            fieldB: a.integer(),
-            nestedCustomType: a.customType({
-              nestedA: a.string(),
-              nestedB: a.string(),
-              grandChild: a.customType({
-                grandA: a.string(),
-                grandB: a.string(),
-              }),
-            }),
-          }),
-        )
+        .returns(a.ref('MyQueryReturnType'))
         .authorization((allow) => allow.publicApiKey()),
     });
 
@@ -1401,6 +1399,28 @@ describe('custom operations + custom type auth inheritance', () => {
     );
   });
 
+  test('inline custom type inherits auth rules from referencing op', () => {
+    const s = a.schema({
+      myQuery: a
+        .query()
+        .handler(a.handler.function('myFn'))
+        .returns(
+          a.customType({
+            fieldA: a.string(),
+            fieldB: a.integer(),
+          }),
+        )
+        .authorization((allow) => allow.publicApiKey()),
+    });
+
+    const result = s.transform().schema;
+
+    expect(result).toMatchSnapshot();
+    expect(result).toEqual(
+      expect.stringContaining('type MyQueryReturnType @aws_api_key\n{'),
+    );
+  });
+
   test('top-level custom type with nested top-level custom types inherits combined auth rules from referencing ops', () => {
     const s = a.schema({
       myQuery: a
diff --git a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
index 059b9149b..8f6bb1318 100644
--- a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
@@ -242,20 +242,20 @@ type Query {
 `;
 
 exports[`custom operations Add entities to SQL schema add custom type, enum, and custom query to generated SQL schema 1`] = `
-"type post @model(timestamps: null) @auth(rules: [{allow: private}])
-{
-  title: String!
-  description: String
-  author: String
-}
-
-enum PostStatus {
+"enum PostStatus {
   draft
   pending
   approved
   published
 }
 
+type post @model(timestamps: null) @auth(rules: [{allow: private}])
+{
+  title: String!
+  description: String
+  author: String
+}
+
 type PostMeta @aws_cognito_user_pools
 {
   viewCount: Int
@@ -595,7 +595,11 @@ exports[`schema auth rules global public auth - multiple models 1`] = `
   "functionSlots": [],
   "jsFunctions": [],
   "lambdaFunctions": {},
-  "schema": "type A @model @auth(rules: [{allow: public, provider: apiKey}])
+  "schema": "enum DTired {
+  ?
+}
+
+type A @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   field: String
 }
@@ -621,10 +625,6 @@ type D @model @auth(rules: [{allow: public, provider: apiKey}])
   tired: DTired
   cId: ID
   c: C @belongsTo(references: ["cId"])
-}
-
-enum DTired {
-  ?
 }",
 }
 `;
diff --git a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
index 98d9201fc..a81507f10 100644
--- a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
@@ -1,16 +1,16 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
 exports[`CustomOperation transform dynamo schema Custom Mutation w required arg and enum 1`] = `
-"type Post @model @auth(rules: [{allow: private}])
-{
-  title: String
-}
-
-enum LikePostReactionType {
+"enum LikePostReactionType {
   :shipit:
   :risitas:
 }
 
+type Post @model @auth(rules: [{allow: private}])
+{
+  title: String
+}
+
 type Mutation {
   likePost(postId: String!, reactionType: LikePostReactionType): Post 
 }"
@@ -45,7 +45,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 +109,7 @@ type Query {
 `;
 
 exports[`CustomOperation transform dynamo schema Custom query w inline custom return type 1`] = `
-"type GetPostDetailsReturnType 
+"type GetPostDetailsReturnType
 {
   stringField: String
   intField: Int
@@ -545,6 +545,18 @@ type Query {
 }"
 `;
 
+exports[`custom operations + custom type auth inheritance inline custom type inherits auth rules from referencing op 1`] = `
+"type MyQueryReturnType @aws_api_key
+{
+  fieldA: String
+  fieldB: Int
+}
+
+type Query {
+  myQuery: MyQueryReturnType @function(name: "myFn") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
 exports[`custom operations + custom type auth inheritance nested custom types inherit auth rules from top-level referencing op 1`] = `
 "type MyQueryReturnType @aws_api_key
 {
diff --git a/packages/data-schema/__tests__/__snapshots__/CustomType.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/CustomType.test.ts.snap
index 3b1a10820..f7b601e28 100644
--- a/packages/data-schema/__tests__/__snapshots__/CustomType.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/CustomType.test.ts.snap
@@ -40,16 +40,16 @@ type Location
 `;
 
 exports[`CustomType transform Explicit CustomType nests explicit CustomType 1`] = `
-"type Post @model @auth(rules: [{allow: public, provider: apiKey}])
-{
-  meta: Meta
-}
-
-enum PostStatus {
+"enum PostStatus {
   unpublished
   published
 }
 
+type Post @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  meta: Meta
+}
+
 type Meta 
 {
   status: PostStatus
@@ -63,16 +63,16 @@ type AltMeta
 `;
 
 exports[`CustomType transform Explicit CustomType nests explicit enum 1`] = `
-"type Post @model @auth(rules: [{allow: public, provider: apiKey}])
-{
-  meta: Meta
-}
-
-enum PostStatus {
+"enum PostStatus {
   unpublished
   published
 }
 
+type Post @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  meta: Meta
+}
+
 type Meta 
 {
   status: PostStatus
@@ -81,7 +81,12 @@ type Meta
 `;
 
 exports[`CustomType transform Explicit CustomType nests implicit CustomType 1`] = `
-"type Post @model @auth(rules: [{allow: public, provider: apiKey}])
+"enum MetaStatus {
+  unpublished
+  published
+}
+
+type Post @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   meta: Meta
 }
@@ -92,11 +97,6 @@ type Meta
   nestedMeta: MetaNestedMeta
 }
 
-enum MetaStatus {
-  unpublished
-  published
-}
-
 type MetaNestedMeta 
 {
   field1: String
@@ -104,7 +104,12 @@ type MetaNestedMeta
 `;
 
 exports[`CustomType transform Explicit CustomType nests implicit enum type 1`] = `
-"type Post @model @auth(rules: [{allow: public, provider: apiKey}])
+"enum MetaStatus {
+  unpublished
+  published
+}
+
+type Post @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   meta: Meta
 }
@@ -113,11 +118,6 @@ type Meta
 {
   status: MetaStatus
   publishedDate: AWSDate
-}
-
-enum MetaStatus {
-  unpublished
-  published
 }"
 `;
 
@@ -135,7 +135,12 @@ type PostLocation
 `;
 
 exports[`CustomType transform Implicit CustomType nests explicit CustomType 1`] = `
-"type Post @model @auth(rules: [{allow: public, provider: apiKey}])
+"enum PostMetaStatus {
+  unpublished
+  published
+}
+
+type Post @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   meta: PostMeta
 }
@@ -149,25 +154,20 @@ type PostMeta
 {
   status: PostMetaStatus
   nestedMeta: AltMeta
-}
-
-enum PostMetaStatus {
-  unpublished
-  published
 }"
 `;
 
 exports[`CustomType transform Implicit CustomType nests explicit enum 1`] = `
-"type Post @model @auth(rules: [{allow: public, provider: apiKey}])
-{
-  meta: PostMeta
-}
-
-enum PostStatus {
+"enum PostStatus {
   unpublished
   published
 }
 
+type Post @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  meta: PostMeta
+}
+
 type PostMeta 
 {
   status: PostStatus
@@ -176,7 +176,12 @@ type PostMeta
 `;
 
 exports[`CustomType transform Implicit CustomType nests implicit CustomType 1`] = `
-"type Post @model @auth(rules: [{allow: public, provider: apiKey}])
+"enum PostMetaStatus {
+  unpublished
+  published
+}
+
+type Post @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   meta: PostMeta
 }
@@ -187,11 +192,6 @@ type PostMeta
   nestedMeta: PostMetaNestedMeta
 }
 
-enum PostMetaStatus {
-  unpublished
-  published
-}
-
 type PostMetaNestedMeta 
 {
   field1: String
@@ -199,7 +199,12 @@ type PostMetaNestedMeta
 `;
 
 exports[`CustomType transform Implicit CustomType nests implicit enum type 1`] = `
-"type Post @model @auth(rules: [{allow: public, provider: apiKey}])
+"enum PostMetaStatus {
+  unpublished
+  published
+}
+
+type Post @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   meta: PostMeta
 }
@@ -208,10 +213,5 @@ type PostMeta
 {
   status: PostMetaStatus
   publishedDate: AWSDate
-}
-
-enum PostMetaStatus {
-  unpublished
-  published
 }"
 `;
diff --git a/packages/data-schema/__tests__/__snapshots__/EnumType.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/EnumType.test.ts.snap
index 524a00326..2492d578d 100644
--- a/packages/data-schema/__tests__/__snapshots__/EnumType.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/EnumType.test.ts.snap
@@ -1,53 +1,53 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
 exports[`EnumType transform Explicit Enum - auth 1`] = `
-"type File @model @auth(rules: [{allow: public, provider: apiKey}])
-{
-  accessLevel: AccessLevel @auth(rules: [{allow: owner, ownerField: "owner"}])
-}
-
-enum AccessLevel {
+"enum AccessLevel {
   public
   protected
   private
+}
+
+type File @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  accessLevel: AccessLevel @auth(rules: [{allow: owner, ownerField: "owner"}])
 }"
 `;
 
 exports[`EnumType transform Explicit Enum - required 1`] = `
-"type File @model @auth(rules: [{allow: public, provider: apiKey}])
-{
-  accessLevel: AccessLevel!
-}
-
-enum AccessLevel {
+"enum AccessLevel {
   public
   protected
   private
+}
+
+type File @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  accessLevel: AccessLevel!
 }"
 `;
 
 exports[`EnumType transform Explicit Enum 1`] = `
-"type File @model @auth(rules: [{allow: public, provider: apiKey}])
-{
-  accessLevel: AccessLevel
-}
-
-enum AccessLevel {
+"enum AccessLevel {
   public
   protected
   private
+}
+
+type File @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  accessLevel: AccessLevel
 }"
 `;
 
 exports[`EnumType transform Implicit Enum 1`] = `
-"type File @model @auth(rules: [{allow: public, provider: apiKey}])
-{
-  accessLevel: FileAccessLevel
-}
-
-enum FileAccessLevel {
+"enum FileAccessLevel {
   public
   protected
   private
+}
+
+type File @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  accessLevel: FileAccessLevel
 }"
 `;
diff --git a/packages/data-schema/__tests__/__snapshots__/ModelIndex.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/ModelIndex.test.ts.snap
index 913272736..112cf923f 100644
--- a/packages/data-schema/__tests__/__snapshots__/ModelIndex.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/ModelIndex.test.ts.snap
@@ -1,61 +1,61 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
 exports[`secondary index schema generation generates correct schema for using a.enum() as the partition key 1`] = `
-"type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
+"enum TodoStatus {
+  open
+  in_progress
+  completed
+}
+
+type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   title: String!
   content: String
   status: TodoStatus @index(sortKeyFields: ["title"], queryField: "listTodoByStatusAndTitle")
-}
+}"
+`;
 
-enum TodoStatus {
+exports[`secondary index schema generation generates correct schema for using a.enum() as the sort key 1`] = `
+"enum TodoStatus {
   open
   in_progress
   completed
-}"
-`;
+}
 
-exports[`secondary index schema generation generates correct schema for using a.enum() as the sort key 1`] = `
-"type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
+type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   title: String! @index(sortKeyFields: ["status"], queryField: "listTodoByTitleAndStatus")
   content: String
   status: TodoStatus
-}
+}"
+`;
 
-enum TodoStatus {
+exports[`secondary index schema generation generates correct schema for using a.ref() (refer to an enum) as the partition key 1`] = `
+"enum TodoStatus {
   open
   in_progress
   completed
-}"
-`;
+}
 
-exports[`secondary index schema generation generates correct schema for using a.ref() (refer to an enum) as the partition key 1`] = `
-"type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
+type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   title: String!
   content: String
   status: TodoStatus @index(sortKeyFields: ["title"], queryField: "listTodoByStatusAndTitle")
-}
+}"
+`;
 
-enum TodoStatus {
+exports[`secondary index schema generation generates correct schema for using a.ref() (refer to an enum) as the sort key 1`] = `
+"enum TodoStatus {
   open
   in_progress
   completed
-}"
-`;
+}
 
-exports[`secondary index schema generation generates correct schema for using a.ref() (refer to an enum) as the sort key 1`] = `
-"type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
+type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   title: String! @index(sortKeyFields: ["status"], queryField: "listTodoByTitleAndStatus")
   content: String
   status: TodoStatus
-}
-
-enum TodoStatus {
-  open
-  in_progress
-  completed
 }"
 `;
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/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 88e386a1f..966dbdf59 100644
--- a/packages/data-schema/src/SchemaProcessor.ts
+++ b/packages/data-schema/src/SchemaProcessor.ts
@@ -155,9 +155,7 @@ function isRefField(
   return isRefFieldDef((field as any)?.data);
 }
 
-function canGenerateFieldType(
-  fieldType: ModelFieldType
-): boolean {
+function canGenerateFieldType(fieldType: ModelFieldType): boolean {
   return fieldType === 'Int';
 }
 
@@ -175,7 +173,6 @@ function scalarFieldToGql(
   } = fieldDef;
   let field: string = fieldType;
 
-
   if (identifier !== undefined) {
     field += '!';
     if (identifier.length > 1) {
@@ -347,6 +344,8 @@ function customOperationToGql(
 ): {
   gqlField: string;
   implicitTypes: [string, any][];
+  inputTypes: string[];
+  returnTypes: string[];
   customTypeAuthRules: CustomTypeAuthRules;
   lambdaFunctionDefinition: LambdaFunctionDefinition;
   customSqlDataSourceStrategy: CustomSqlDataSourceStrategy | undefined;
@@ -361,6 +360,9 @@ 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
@@ -407,7 +409,10 @@ function customOperationToGql(
           authRules: authorization,
         };
 
-        implicitTypes.push([returnTypeName, returnType]);
+        implicitTypes.push([
+          returnTypeName,
+          { ...returnType, generateInputType: false },
+        ]);
       }
       return returnTypeName;
     } else if (isEnumType(returnType)) {
@@ -444,6 +449,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(
@@ -451,9 +466,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];
@@ -550,6 +578,8 @@ function customOperationToGql(
   return {
     gqlField,
     implicitTypes: implicitTypes,
+    inputTypes,
+    returnTypes,
     customTypeAuthRules,
     lambdaFunctionDefinition,
     customSqlDataSourceStrategy,
@@ -929,23 +959,33 @@ function processFieldLevelAuthRules(
   return fieldLevelAuthRules;
 }
 
-function validateDBGeneration(fields: Record<string, any>, databaseEngine: DatasourceEngine) {
+function validateDBGeneration(
+  fields: Record<string, any>,
+  databaseEngine: DatasourceEngine,
+) {
   for (const [fieldName, fieldDef] of Object.entries(fields)) {
     const _default = fieldDef.data?.default;
     const fieldType = fieldDef.data?.fieldType;
     const isGenerated = _default === __generated;
 
     if (isGenerated && databaseEngine !== 'postgresql') {
-      throw new Error(`Invalid field definition for ${fieldName}. DB-generated fields are only supported with PostgreSQL data sources.`);
+      throw new Error(
+        `Invalid field definition for ${fieldName}. DB-generated fields are only supported with PostgreSQL data sources.`,
+      );
     }
 
     if (isGenerated && !canGenerateFieldType(fieldType)) {
-      throw new Error(`Incompatible field type. Field type ${fieldType} in field ${fieldName} cannot be configured as a DB-generated field.`);
+      throw new Error(
+        `Incompatible field type. Field type ${fieldType} in field ${fieldName} cannot be configured as a DB-generated field.`,
+      );
     }
   }
 }
 
-function validateNullableIdentifiers(fields: Record<string, any>, identifier?: readonly string[]){
+function validateNullableIdentifiers(
+  fields: Record<string, any>,
+  identifier?: readonly string[],
+) {
   for (const [fieldName, fieldDef] of Object.entries(fields)) {
     const fieldType = fieldDef.data?.fieldType;
     const required = fieldDef.data?.required;
@@ -954,7 +994,9 @@ function validateNullableIdentifiers(fields: Record<string, any>, identifier?: r
 
     if (identifier !== undefined && identifier.includes(fieldName)) {
       if (!required && fieldType !== 'ID' && !isGenerated) {
-        throw new Error(`Invalid identifier definition. Field ${fieldName} cannot be used in the identifier. Identifiers must reference required or DB-generated fields)`);
+        throw new Error(
+          `Invalid identifier definition. Field ${fieldName} cannot be used in the identifier. Identifiers must reference required or DB-generated fields)`,
+        );
       }
     }
   }
@@ -965,10 +1007,13 @@ 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)
@@ -977,7 +1022,7 @@ function processFields(
 
   validateImpliedFields(fields, impliedFields);
   validateDBGeneration(fields, databaseEngine);
-  validateNullableIdentifiers(fields, identifier)
+  validateNullableIdentifiers(fields, identifier);
 
   for (const [fieldName, fieldDef] of Object.entries(fields)) {
     const fieldAuth = fieldLevelAuthRules[fieldName]
@@ -998,40 +1043,74 @@ function processFields(
           )}${fieldAuth}`,
         );
       } else if (isRefField(fieldDef)) {
-        gqlFields.push(
-          `${fieldName}: ${refFieldToGql(fieldDef.data, secondaryIndexes[fieldName])}${fieldAuth}`,
-        );
+        if (generateInputType) {
+          const inputTypeName = `${capitalize(typeName)}${capitalize(fieldName)}Input`;
+          gqlFields.push(`${fieldName}: ${inputTypeName}${fieldAuth}`);
+
+          const refTypeInfo = getRefType(fieldDef.data.link, typeName);
+
+          if (refTypeInfo.type === 'CustomType') {
+            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 {
+            throw new Error(
+              `Field '${fieldName}' in type '${typeName}' references '${fieldDef.data.link}' which is not a CustomType. Check schema definitions.`,
+            );
+          }
+        } else {
+          gqlFields.push(
+            `${fieldName}: ${refFieldToGql(fieldDef.data, secondaryIndexes[fieldName])}${fieldAuth}`,
+          );
+        }
       } 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]);
-
         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]);
-
-        gqlFields.push(`${fieldName}: ${customTypeName}`);
+        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)}`,
+      );
     }
   }
 
@@ -1307,6 +1386,50 @@ const mergeCustomTypeAuthRules = (
   }
 };
 
+function generateInputTypes(
+  implicitTypes: [string, any][],
+  generateInputType: boolean,
+  getRefType: ReturnType<typeof getRefTypeForSchema>,
+  authRules?: Authorization<any, any, any>[],
+  isInlineType = false,
+): string[] {
+  const generatedTypes = new Set<string>();
+
+  implicitTypes.forEach(([typeName, typeDef]) => {
+    if (isCustomType(typeDef)) {
+      const { gqlFields } = 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);
+    } else if (typeDef.type === 'enum') {
+      const enumDefinition = `enum ${typeName} {\n  ${typeDef.values.join('\n  ')}\n}`;
+      generatedTypes.add(enumDefinition);
+    } else if (typeDef.type === 'scalar') {
+      const scalarDefinition = `scalar ${typeName}`;
+      generatedTypes.add(scalarDefinition);
+    } else {
+      console.warn(`Unexpected type definition for ${typeName}:`, typeDef);
+    }
+  });
+
+  return Array.from(generatedTypes);
+}
+
 const schemaPreprocessor = (
   schema: InternalSchema,
 ): {
@@ -1316,7 +1439,9 @@ const schemaPreprocessor = (
   lambdaFunctions: LambdaFunctionDefinition;
   customSqlDataSourceStrategies?: CustomSqlDataSourceStrategy[];
 } => {
+  const enumTypes = new Set<string>();
   const gqlModels: string[] = [];
+  const inputTypes: string[] = [];
 
   const customQueries = [];
   const customMutations = [];
@@ -1334,11 +1459,8 @@ const schemaPreprocessor = (
   const lambdaFunctions: LambdaFunctionDefinition = {};
   const customSqlDataSourceStrategies: CustomSqlDataSourceStrategy[] = [];
 
-  const databaseEngine = schema.data.configuration.database.engine
-  const databaseType =
-    databaseEngine === 'dynamodb'
-      ? 'dynamodb'
-      : 'sql';
+  const databaseEngine = schema.data.configuration.database.engine;
+  const databaseType = databaseEngine === 'dynamodb' ? 'dynamodb' : 'sql';
 
   const staticSchema = databaseType === 'sql';
 
@@ -1382,10 +1504,8 @@ const schemaPreprocessor = (
             `Values of the enum type ${typeName} should not contain any whitespace.`,
           );
         }
-        const enumType = `enum ${typeName} {\n  ${typeDef.values.join(
-          '\n  ',
-        )}\n}`;
-        gqlModels.push(enumType);
+        const enumType = `enum ${typeName} {\n  ${typeDef.values.join('\n  ')}\n}`;
+        enumTypes.add(enumType);
       } else if (isCustomType(typeDef)) {
         const fields = typeDef.data.fields;
 
@@ -1419,14 +1539,20 @@ const schemaPreprocessor = (
           fields,
           authFields,
           fieldLevelAuthRules,
+          getRefType,
           undefined,
           undefined,
-          undefined,
-          databaseEngine
+          databaseEngine,
         );
-
-        topLevelTypes.push(...implicitTypes);
-
+        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);
+          }
+        }
         const joined = gqlFields.join('\n  ');
 
         const model = `type ${typeName} ${customAuth}\n{\n  ${joined}\n}`;
@@ -1439,6 +1565,8 @@ const schemaPreprocessor = (
         const {
           gqlField,
           implicitTypes,
+          inputTypes: operationInputTypes,
+          returnTypes: operationReturnTypes,
           customTypeAuthRules,
           jsFunctionForField,
           lambdaFunctionDefinition,
@@ -1450,8 +1578,39 @@ const schemaPreprocessor = (
           databaseType,
           getRefType,
         );
-
-        topLevelTypes.push(...implicitTypes);
+        inputTypes.push(...operationInputTypes);
+        gqlModels.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 (isEnumType(typeDef)) {
+            const enumTypeDefinition = `enum ${typeName} {\n  ${typeDef.values.join('\n  ')}\n}`;
+            enumTypes.add(enumTypeDefinition);
+          } else {
+            const shouldGenerateInputType = !operationReturnTypes.some(
+              (returnType) => returnType.includes(typeName),
+            );
+            const generatedTypeDefinition = generateInputTypes(
+              [[typeName, typeDef]],
+              shouldGenerateInputType,
+              getRefType,
+            )[0];
+            if (shouldGenerateInputType) {
+              inputTypes.push(generatedTypeDefinition);
+            }
+          }
+        });
 
         mergeCustomTypeAuthRules(
           customTypeInheritedAuthRules,
@@ -1531,12 +1690,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  ');
@@ -1596,6 +1762,7 @@ const schemaPreprocessor = (
         fields,
         authFields,
         fieldLevelAuthRules,
+        getRefType,
         identifier,
         partitionKey,
         transformedSecondaryIndexes,
@@ -1622,12 +1789,21 @@ const schemaPreprocessor = (
     subscriptions: customSubscriptions,
   };
 
-  gqlModels.push(...generateCustomOperationTypes(customOperations));
+  const customOperationTypes = generateCustomOperationTypes(customOperations);
+
+  const schemaComponents = [
+    ...Array.from(enumTypes),
+    ...gqlModels,
+    ...customOperationTypes,
+  ];
+
   if (shouldAddConversationTypes) {
-    gqlModels.push(CONVERSATION_SCHEMA_GRAPHQL_TYPES);
+    schemaComponents.push(CONVERSATION_SCHEMA_GRAPHQL_TYPES);
   }
 
-  const processedSchema = gqlModels.join('\n\n');
+  schemaComponents.push(...inputTypes);
+
+  const processedSchema = schemaComponents.join('\n\n');
 
   return {
     schema: processedSchema,
@@ -1926,6 +2102,8 @@ function transformCustomOperations(
   const {
     gqlField,
     implicitTypes,
+    inputTypes,
+    returnTypes,
     customTypeAuthRules,
     lambdaFunctionDefinition,
     customSqlDataSourceStrategy,
@@ -1941,6 +2119,8 @@ function transformCustomOperations(
   return {
     gqlField,
     implicitTypes,
+    inputTypes,
+    returnTypes,
     customTypeAuthRules,
     jsFunctionForField,
     lambdaFunctionDefinition,
@@ -1975,7 +2155,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/1-patterns/__snapshots__/client-schema.ts.snap b/packages/integration-tests/__tests__/defined-behavior/1-patterns/__snapshots__/client-schema.ts.snap
index f35e2fdc1..14cdf30bc 100644
--- a/packages/integration-tests/__tests__/defined-behavior/1-patterns/__snapshots__/client-schema.ts.snap
+++ b/packages/integration-tests/__tests__/defined-behavior/1-patterns/__snapshots__/client-schema.ts.snap
@@ -1,7 +1,13 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
 exports[`ClientSchema overview sample schema 1`] = `
-"type Customer @model @auth(rules: [{allow: public, provider: apiKey}])
+"enum CustomerEngagementStage {
+  PROSPECT
+  INTERESTED
+  PURCHASED
+}
+
+type Customer @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   customerId: ID! @primaryKey
   name: String
@@ -22,11 +28,5 @@ type CustomerLocation
 {
   lat: Float!
   long: Float!
-}
-
-enum CustomerEngagementStage {
-  PROSPECT
-  INTERESTED
-  PURCHASED
 }"
 `;
diff --git a/packages/integration-tests/__tests__/defined-behavior/2-expected-use/__snapshots__/custom-operations.ts.snap b/packages/integration-tests/__tests__/defined-behavior/2-expected-use/__snapshots__/custom-operations.ts.snap
index d9625d55f..270c95500 100644
--- a/packages/integration-tests/__tests__/defined-behavior/2-expected-use/__snapshots__/custom-operations.ts.snap
+++ b/packages/integration-tests/__tests__/defined-behavior/2-expected-use/__snapshots__/custom-operations.ts.snap
@@ -40,6 +40,86 @@ exports[`custom operations async sync 1`] = `
 ]
 `;
 
+exports[`custom operations client operations with custom types and refs client can call custom mutation with custom type argument 1`] = `
+[
+  [
+    {
+      "authMode": undefined,
+      "authToken": undefined,
+      "query": "
+    mutation($customArg: MutateWithCustomTypeArgCustomArgInput) {
+      mutateWithCustomTypeArg(customArg: $customArg) 
+    }
+  ",
+      "variables": {
+        "customArg": {},
+      },
+    },
+    {},
+  ],
+]
+`;
+
+exports[`custom operations client operations with custom types and refs client can call custom mutation with ref argument 1`] = `
+[
+  [
+    {
+      "authMode": undefined,
+      "authToken": undefined,
+      "query": "
+    mutation($refArg: MutationWithRefArgRefArgInput) {
+      mutationWithRefArg(refArg: $refArg) 
+    }
+  ",
+      "variables": {
+        "refArg": {},
+      },
+    },
+    {},
+  ],
+]
+`;
+
+exports[`custom operations client operations with custom types and refs client can call custom query with custom type argument 1`] = `
+[
+  [
+    {
+      "authMode": undefined,
+      "authToken": undefined,
+      "query": "
+    query($customArg: QueryWithCustomTypeArgCustomArgInput) {
+      queryWithCustomTypeArg(customArg: $customArg) 
+    }
+  ",
+      "variables": {
+        "customArg": {},
+      },
+    },
+    {},
+  ],
+]
+`;
+
+exports[`custom operations client operations with custom types and refs client can call custom query with ref argument 1`] = `
+[
+  [
+    {
+      "authMode": undefined,
+      "authToken": undefined,
+      "query": "
+    query($refArg: QueryWithRefArgRefArgInput) {
+      queryWithRefArg(refArg: $refArg) 
+    }
+  ",
+      "variables": {
+        "refArg": {},
+      },
+    },
+    {},
+  ],
+]
+`;
+
 exports[`custom operations custom type array result 1`] = `
 [
   [
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..4a00b7f05 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
@@ -540,4 +729,85 @@ describe('custom operations', () => {
       const { data } = await client.queries.echoEnum({ status: 'BAD VALUE' });
     });
   });
+  describe('client operations with custom types and refs', () => {
+    test('client can call custom query with custom type argument', async () => {
+      const { spy, generateClient } = mockedGenerateClient([
+        {
+          data: {
+            queryWithCustomTypeArg: 'Custom type query result',
+          },
+        },
+      ]);
+
+      const config = await buildAmplifyConfig(schema);
+      Amplify.configure(config);
+      const client = generateClient<Schema>();
+
+      const { data } = await client.queries.queryWithCustomTypeArg({
+        customArg: {},
+      });
+      expect(data).toEqual('Custom type query result');
+      expect(optionsAndHeaders(spy)).toMatchSnapshot();
+    });
+
+    test('client can call custom query with ref argument', async () => {
+      const { spy, generateClient } = mockedGenerateClient([
+        {
+          data: {
+            queryWithRefArg: 'Ref type query result',
+          },
+        },
+      ]);
+
+      const config = await buildAmplifyConfig(schema);
+      Amplify.configure(config);
+      const client = generateClient<Schema>();
+
+      const { data } = await client.queries.queryWithRefArg({
+        refArg: {},
+      });
+      expect(data).toEqual('Ref type query result');
+      expect(optionsAndHeaders(spy)).toMatchSnapshot();
+    });
+
+    test('client can call custom mutation with custom type argument', async () => {
+      const { spy, generateClient } = mockedGenerateClient([
+        {
+          data: {
+            mutateWithCustomTypeArg: 'Custom type mutation result',
+          },
+        },
+      ]);
+
+      const config = await buildAmplifyConfig(schema);
+      Amplify.configure(config);
+      const client = generateClient<Schema>();
+
+      const { data } = await client.mutations.mutateWithCustomTypeArg({
+        customArg: {},
+      });
+      expect(data).toEqual('Custom type mutation result');
+      expect(optionsAndHeaders(spy)).toMatchSnapshot();
+    });
+
+    test('client can call custom mutation with ref argument', async () => {
+      const { spy, generateClient } = mockedGenerateClient([
+        {
+          data: {
+            mutationWithRefArg: 'Ref type mutation result',
+          },
+        },
+      ]);
+
+      const config = await buildAmplifyConfig(schema);
+      Amplify.configure(config);
+      const client = generateClient<Schema>();
+
+      const { data } = await client.mutations.mutationWithRefArg({
+        refArg: {},
+      });
+      expect(data).toEqual('Ref type mutation result');
+      expect(optionsAndHeaders(spy)).toMatchSnapshot();
+    });
+  });
 });