diff --git a/packages/orm/src/client/contract.ts b/packages/orm/src/client/contract.ts
index d5fe86ad..01fec352 100644
--- a/packages/orm/src/client/contract.ts
+++ b/packages/orm/src/client/contract.ts
@@ -2,6 +2,7 @@ import type Decimal from 'decimal.js';
 import {
     type FieldIsArray,
     type GetModels,
+    type GetTypeDefs,
     type IsDelegateModel,
     type ProcedureDef,
     type RelationFields,
@@ -32,6 +33,7 @@ import type {
     SelectSubset,
     SimplifiedModelResult,
     Subset,
+    TypeDefResult,
     UpdateArgs,
     UpdateManyAndReturnArgs,
     UpdateManyArgs,
@@ -854,10 +856,10 @@ type AuthModelType<Schema extends SchemaDef, Model extends GetModels<Schema>> =
 };
 
 export type AuthType<Schema extends SchemaDef> =
-    string extends GetModels<Schema>
-        ? Record<string, unknown>
-        : Schema['authType'] extends GetModels<Schema>
-          ? AuthModelType<Schema, Schema['authType']>
-          : never;
+    Schema['authType'] extends GetModels<Schema>
+        ? AuthModelType<Schema, Schema['authType']>
+        : Schema['authType'] extends GetTypeDefs<Schema>
+          ? TypeDefResult<Schema, Schema['authType'], true>
+          : Record<string, unknown>;
 
 //#endregion
diff --git a/packages/orm/src/client/crud-types.ts b/packages/orm/src/client/crud-types.ts
index 21e75967..55e3edfd 100644
--- a/packages/orm/src/client/crud-types.ts
+++ b/packages/orm/src/client/crud-types.ts
@@ -38,6 +38,7 @@ import type {
     NullableIf,
     Optional,
     OrArray,
+    PartialIf,
     Simplify,
     ValueOfPotentialTuple,
     WrapType,
@@ -246,17 +247,31 @@ export type SimplifiedModelResult<
     Array = false,
 > = Simplify<ModelResult<Schema, Model, Options, Args, Optional, Array>>;
 
-export type TypeDefResult<Schema extends SchemaDef, TypeDef extends GetTypeDefs<Schema>> = Optional<
-    {
-        [Key in GetTypeDefFields<Schema, TypeDef>]: MapTypeDefFieldType<Schema, TypeDef, Key>;
-    },
-    // optionality
-    keyof {
-        [Key in GetTypeDefFields<Schema, TypeDef> as TypeDefFieldIsOptional<Schema, TypeDef, Key> extends true
-            ? Key
-            : never]: true;
-    }
->;
+export type TypeDefResult<
+    Schema extends SchemaDef,
+    TypeDef extends GetTypeDefs<Schema>,
+    Partial extends boolean = false,
+> = PartialIf<
+    Optional<
+        {
+            [Key in GetTypeDefFields<Schema, TypeDef>]: MapFieldDefType<
+                Schema,
+                GetTypeDefField<Schema, TypeDef, Key>,
+                Partial
+            >;
+        },
+        // optionality
+        Partial extends true
+            ? never
+            : keyof {
+                  [Key in GetTypeDefFields<Schema, TypeDef> as TypeDefFieldIsOptional<Schema, TypeDef, Key> extends true
+                      ? Key
+                      : never]: true;
+              }
+    >,
+    Partial
+> &
+    Record<string, unknown>;
 
 export type BatchResult = { count: number };
 
@@ -617,17 +632,15 @@ type MapModelFieldType<
     Field extends GetModelFields<Schema, Model>,
 > = MapFieldDefType<Schema, GetModelField<Schema, Model, Field>>;
 
-type MapTypeDefFieldType<
+type MapFieldDefType<
     Schema extends SchemaDef,
-    TypeDef extends GetTypeDefs<Schema>,
-    Field extends GetTypeDefFields<Schema, TypeDef>,
-> = MapFieldDefType<Schema, GetTypeDefField<Schema, TypeDef, Field>>;
-
-type MapFieldDefType<Schema extends SchemaDef, T extends Pick<FieldDef, 'type' | 'optional' | 'array'>> = WrapType<
+    T extends Pick<FieldDef, 'type' | 'optional' | 'array'>,
+    Partial extends boolean = false,
+> = WrapType<
     T['type'] extends GetEnums<Schema>
         ? keyof GetEnum<Schema, T['type']>
         : T['type'] extends GetTypeDefs<Schema>
-          ? TypeDefResult<Schema, T['type']> & Record<string, unknown>
+          ? TypeDefResult<Schema, T['type'], Partial> & Record<string, unknown>
           : MapBaseType<T['type']>,
     T['optional'],
     T['array']
diff --git a/packages/orm/src/utils/type-utils.ts b/packages/orm/src/utils/type-utils.ts
index 47b4e095..26828cf6 100644
--- a/packages/orm/src/utils/type-utils.ts
+++ b/packages/orm/src/utils/type-utils.ts
@@ -2,8 +2,12 @@ import type Decimal from 'decimal.js';
 
 export type Optional<T extends object, K extends keyof T = keyof T> = Omit<T, K> & Partial<Pick<T, K>>;
 
+export type PartialIf<T, Condition extends boolean> = Condition extends true ? Partial<T> : T;
+
 export type NullableIf<T, Condition extends boolean> = Condition extends true ? T | null : T;
 
+export type ArrayIf<T, Condition extends boolean> = Condition extends true ? T[] : T;
+
 export type PartialRecord<K extends string | number | symbol, T> = Partial<Record<K, T>>;
 
 type _Preserve = Date | Function | Decimal | Uint8Array | JsonObject | JsonValue;
diff --git a/packages/schema/src/schema.ts b/packages/schema/src/schema.ts
index 95abd005..abc98ace 100644
--- a/packages/schema/src/schema.ts
+++ b/packages/schema/src/schema.ts
@@ -15,7 +15,7 @@ export type SchemaDef = {
     typeDefs?: Record<string, TypeDefDef>;
     plugins: Record<string, unknown>;
     procedures?: Record<string, ProcedureDef>;
-    authType?: GetModels<SchemaDef>;
+    authType?: GetModels<SchemaDef> | GetTypeDefs<SchemaDef>;
 };
 
 export type ModelDef = {
@@ -241,6 +241,12 @@ export type TypeDefFieldIsOptional<
     Field extends GetTypeDefFields<Schema, TypeDef>,
 > = GetTypeDefField<Schema, TypeDef, Field>['optional'] extends true ? true : false;
 
+export type TypeDefFieldIsArray<
+    Schema extends SchemaDef,
+    TypeDef extends GetTypeDefs<Schema>,
+    Field extends GetTypeDefFields<Schema, TypeDef>,
+> = GetTypeDefField<Schema, TypeDef, Field>['array'] extends true ? true : false;
+
 export type FieldIsRelation<
     Schema extends SchemaDef,
     Model extends GetModels<Schema>,
diff --git a/tests/e2e/orm/client-api/auth-typing.test.ts b/tests/e2e/orm/client-api/auth-typing.test.ts
new file mode 100644
index 00000000..ce510227
--- /dev/null
+++ b/tests/e2e/orm/client-api/auth-typing.test.ts
@@ -0,0 +1,18 @@
+import { createTestClient } from '@zenstackhq/testtools';
+import { describe, it } from 'vitest';
+import { schema } from '../schemas/auth-type/schema';
+
+describe('Custom auth typing tests', () => {
+    it('works with custom auth typing', async () => {
+        const db = await createTestClient(schema);
+        db.$setAuth({
+            id: 1,
+            role: 'ADMIN',
+            permissions: [
+                {
+                    actionCode: 'MANAGE',
+                },
+            ],
+        });
+    });
+});
diff --git a/tests/e2e/orm/schemas/auth-type/input.ts b/tests/e2e/orm/schemas/auth-type/input.ts
new file mode 100644
index 00000000..9564b7f2
--- /dev/null
+++ b/tests/e2e/orm/schemas/auth-type/input.ts
@@ -0,0 +1,30 @@
+//////////////////////////////////////////////////////////////////////////////////////////////
+// DO NOT MODIFY THIS FILE                                                                  //
+// This file is automatically generated by ZenStack CLI and should not be manually updated. //
+//////////////////////////////////////////////////////////////////////////////////////////////
+
+/* eslint-disable */
+
+import { type SchemaType as $Schema } from "./schema";
+import type { FindManyArgs as $FindManyArgs, FindUniqueArgs as $FindUniqueArgs, FindFirstArgs as $FindFirstArgs, CreateArgs as $CreateArgs, CreateManyArgs as $CreateManyArgs, CreateManyAndReturnArgs as $CreateManyAndReturnArgs, UpdateArgs as $UpdateArgs, UpdateManyArgs as $UpdateManyArgs, UpdateManyAndReturnArgs as $UpdateManyAndReturnArgs, UpsertArgs as $UpsertArgs, DeleteArgs as $DeleteArgs, DeleteManyArgs as $DeleteManyArgs, CountArgs as $CountArgs, AggregateArgs as $AggregateArgs, GroupByArgs as $GroupByArgs, WhereInput as $WhereInput, SelectInput as $SelectInput, IncludeInput as $IncludeInput, OmitInput as $OmitInput, ClientOptions as $ClientOptions } from "@zenstackhq/orm";
+import type { SimplifiedModelResult as $SimplifiedModelResult, SelectIncludeOmit as $SelectIncludeOmit } from "@zenstackhq/orm";
+export type FooFindManyArgs = $FindManyArgs<$Schema, "Foo">;
+export type FooFindUniqueArgs = $FindUniqueArgs<$Schema, "Foo">;
+export type FooFindFirstArgs = $FindFirstArgs<$Schema, "Foo">;
+export type FooCreateArgs = $CreateArgs<$Schema, "Foo">;
+export type FooCreateManyArgs = $CreateManyArgs<$Schema, "Foo">;
+export type FooCreateManyAndReturnArgs = $CreateManyAndReturnArgs<$Schema, "Foo">;
+export type FooUpdateArgs = $UpdateArgs<$Schema, "Foo">;
+export type FooUpdateManyArgs = $UpdateManyArgs<$Schema, "Foo">;
+export type FooUpdateManyAndReturnArgs = $UpdateManyAndReturnArgs<$Schema, "Foo">;
+export type FooUpsertArgs = $UpsertArgs<$Schema, "Foo">;
+export type FooDeleteArgs = $DeleteArgs<$Schema, "Foo">;
+export type FooDeleteManyArgs = $DeleteManyArgs<$Schema, "Foo">;
+export type FooCountArgs = $CountArgs<$Schema, "Foo">;
+export type FooAggregateArgs = $AggregateArgs<$Schema, "Foo">;
+export type FooGroupByArgs = $GroupByArgs<$Schema, "Foo">;
+export type FooWhereInput = $WhereInput<$Schema, "Foo">;
+export type FooSelect = $SelectInput<$Schema, "Foo">;
+export type FooInclude = $IncludeInput<$Schema, "Foo">;
+export type FooOmit = $OmitInput<$Schema, "Foo">;
+export type FooGetPayload<Args extends $SelectIncludeOmit<$Schema, "Foo", true>, Options extends $ClientOptions<$Schema> = $ClientOptions<$Schema>> = $SimplifiedModelResult<$Schema, "Foo", Options, Args>;
diff --git a/tests/e2e/orm/schemas/auth-type/models.ts b/tests/e2e/orm/schemas/auth-type/models.ts
new file mode 100644
index 00000000..75911fc7
--- /dev/null
+++ b/tests/e2e/orm/schemas/auth-type/models.ts
@@ -0,0 +1,12 @@
+//////////////////////////////////////////////////////////////////////////////////////////////
+// DO NOT MODIFY THIS FILE                                                                  //
+// This file is automatically generated by ZenStack CLI and should not be manually updated. //
+//////////////////////////////////////////////////////////////////////////////////////////////
+
+/* eslint-disable */
+
+import { type SchemaType as $Schema } from "./schema";
+import { type ModelResult as $ModelResult, type TypeDefResult as $TypeDefResult } from "@zenstackhq/orm";
+export type Foo = $ModelResult<$Schema, "Foo">;
+export type Permission = $TypeDefResult<$Schema, "Permission">;
+export type Auth = $TypeDefResult<$Schema, "Auth">;
diff --git a/tests/e2e/orm/schemas/auth-type/schema.ts b/tests/e2e/orm/schemas/auth-type/schema.ts
new file mode 100644
index 00000000..5904569b
--- /dev/null
+++ b/tests/e2e/orm/schemas/auth-type/schema.ts
@@ -0,0 +1,74 @@
+//////////////////////////////////////////////////////////////////////////////////////////////
+// DO NOT MODIFY THIS FILE                                                                  //
+// This file is automatically generated by ZenStack CLI and should not be manually updated. //
+//////////////////////////////////////////////////////////////////////////////////////////////
+
+/* eslint-disable */
+
+import { type SchemaDef } from "@zenstackhq/orm/schema";
+const _schema = {
+    provider: {
+        type: "sqlite"
+    },
+    models: {
+        Foo: {
+            name: "Foo",
+            fields: {
+                id: {
+                    name: "id",
+                    type: "Int",
+                    id: true,
+                    attributes: [{ name: "@id" }]
+                }
+            },
+            idFields: ["id"],
+            uniqueFields: {
+                id: { type: "Int" }
+            }
+        }
+    },
+    typeDefs: {
+        Permission: {
+            name: "Permission",
+            fields: {
+                actionCode: {
+                    name: "actionCode",
+                    type: "String"
+                }
+            }
+        },
+        Auth: {
+            name: "Auth",
+            fields: {
+                id: {
+                    name: "id",
+                    type: "Int",
+                    attributes: [{ name: "@id" }]
+                },
+                name: {
+                    name: "name",
+                    type: "String"
+                },
+                permissions: {
+                    name: "permissions",
+                    type: "Permission",
+                    array: true
+                },
+                role: {
+                    name: "role",
+                    type: "String"
+                }
+            },
+            attributes: [
+                { name: "@@auth" }
+            ]
+        }
+    },
+    authType: "Auth",
+    plugins: {}
+} as const satisfies SchemaDef;
+type Schema = typeof _schema & {
+    __brand?: "schema";
+};
+export const schema: Schema = _schema;
+export type SchemaType = Schema;
diff --git a/tests/e2e/orm/schemas/auth-type/schema.zmodel b/tests/e2e/orm/schemas/auth-type/schema.zmodel
new file mode 100644
index 00000000..9ded3aae
--- /dev/null
+++ b/tests/e2e/orm/schemas/auth-type/schema.zmodel
@@ -0,0 +1,21 @@
+datasource db {
+    provider = "sqlite"
+    url      = "file:./dev.db"
+}
+
+type Permission {
+    actionCode String
+}
+
+type Auth {
+    id          Int          @id
+    name        String
+    permissions Permission[]
+    role        String
+
+    @@auth
+}
+
+model Foo {
+    id Int @id
+}
\ No newline at end of file