diff --git a/packages/cli/package.json b/packages/cli/package.json
index 39540625b..7c81f4251 100644
--- a/packages/cli/package.json
+++ b/packages/cli/package.json
@@ -35,6 +35,7 @@
         "@zenstackhq/common-helpers": "workspace:*",
         "@zenstackhq/language": "workspace:*",
         "@zenstackhq/sdk": "workspace:*",
+        "chokidar": "^5.0.0",
         "colors": "1.4.0",
         "commander": "^8.3.0",
         "execa": "^9.6.0",
diff --git a/packages/cli/src/actions/generate.ts b/packages/cli/src/actions/generate.ts
index c41a99ea4..16e3826c7 100644
--- a/packages/cli/src/actions/generate.ts
+++ b/packages/cli/src/actions/generate.ts
@@ -1,5 +1,6 @@
-import { invariant } from '@zenstackhq/common-helpers';
-import { isPlugin, LiteralExpr, Plugin, type Model } from '@zenstackhq/language/ast';
+import { invariant, singleDebounce } from '@zenstackhq/common-helpers';
+import { ZModelLanguageMetaData } from '@zenstackhq/language';
+import { type AbstractDeclaration, isPlugin, LiteralExpr, Plugin, type Model } from '@zenstackhq/language/ast';
 import { getLiteral, getLiteralArray } from '@zenstackhq/language/utils';
 import { type CliPlugin } from '@zenstackhq/sdk';
 import colors from 'colors';
@@ -7,6 +8,7 @@ import { createJiti } from 'jiti';
 import fs from 'node:fs';
 import path from 'node:path';
 import { pathToFileURL } from 'node:url';
+import { watch } from 'chokidar';
 import ora, { type Ora } from 'ora';
 import { CliError } from '../cli-error';
 import * as corePlugins from '../plugins';
@@ -16,6 +18,7 @@ type Options = {
     schema?: string;
     output?: string;
     silent: boolean;
+    watch: boolean;
     lite: boolean;
     liteOnly: boolean;
 };
@@ -24,6 +27,96 @@ type Options = {
  * CLI action for generating code from schema
  */
 export async function run(options: Options) {
+    const model = await pureGenerate(options, false);
+
+    if (options.watch) {
+        const logsEnabled = !options.silent;
+
+        if (logsEnabled) {
+            console.log(colors.green(`\nEnabled watch mode!`));
+        }
+
+        const schemaExtensions = ZModelLanguageMetaData.fileExtensions;
+
+        // Get real models file path (cuz its merged into single document -> we need use cst nodes)
+        const getRootModelWatchPaths = (model: Model) => new Set<string>(
+            (
+                model.declarations.filter(
+                    (v) =>
+                        v.$cstNode?.parent?.element.$type === 'Model' &&
+                        !!v.$cstNode.parent.element.$document?.uri?.fsPath,
+                ) as AbstractDeclaration[]
+            ).map((v) => v.$cstNode!.parent!.element.$document!.uri!.fsPath),
+        );
+
+        const watchedPaths = getRootModelWatchPaths(model);
+
+        if (logsEnabled) {
+            const logPaths = [...watchedPaths].map((at) => `- ${at}`).join('\n');
+            console.log(`Watched file paths:\n${logPaths}`);
+        }
+
+        const watcher = watch([...watchedPaths], {
+            alwaysStat: false,
+            ignoreInitial: true,
+            ignorePermissionErrors: true,
+            ignored: (at) => !schemaExtensions.some((ext) => at.endsWith(ext)),
+        });
+
+        // prevent save multiple files and run multiple times
+        const reGenerateSchema = singleDebounce(async () => {
+            if (logsEnabled) {
+                console.log('Got changes, run generation!');
+            }
+
+            try {
+                const newModel = await pureGenerate(options, true);
+                const allModelsPaths = getRootModelWatchPaths(newModel);
+                const newModelPaths = [...allModelsPaths].filter((at) => !watchedPaths.has(at));
+                const removeModelPaths = [...watchedPaths].filter((at) => !allModelsPaths.has(at));
+
+                if (newModelPaths.length) {
+                    if (logsEnabled) {
+                        const logPaths = newModelPaths.map((at) => `- ${at}`).join('\n');
+                        console.log(`Added file(s) to watch:\n${logPaths}`);
+                    }
+
+                    newModelPaths.forEach((at) => watchedPaths.add(at));
+                    watcher.add(newModelPaths);
+                }
+
+                if (removeModelPaths.length) {
+                    if (logsEnabled) {
+                        const logPaths = removeModelPaths.map((at) => `- ${at}`).join('\n');
+                        console.log(`Removed file(s) from watch:\n${logPaths}`);
+                    }
+
+                    removeModelPaths.forEach((at) => watchedPaths.delete(at));
+                    watcher.unwatch(removeModelPaths);
+                }
+            } catch (e) {
+                console.error(e);
+            }
+        }, 500, true);
+
+        watcher.on('unlink', (pathAt) => {
+            if (logsEnabled) {
+                console.log(`Removed file from watch: ${pathAt}`);
+            }
+
+            watchedPaths.delete(pathAt);
+            watcher.unwatch(pathAt);
+
+            reGenerateSchema();
+        });
+
+        watcher.on('change', () => {
+            reGenerateSchema();
+        });
+    }
+}
+
+async function pureGenerate(options: Options, fromWatch: boolean) {
     const start = Date.now();
 
     const schemaFile = getSchemaFile(options.schema);
@@ -35,7 +128,9 @@ export async function run(options: Options) {
 
     if (!options.silent) {
         console.log(colors.green(`Generation completed successfully in ${Date.now() - start}ms.\n`));
-        console.log(`You can now create a ZenStack client with it.
+
+        if (!fromWatch) {
+            console.log(`You can now create a ZenStack client with it.
 
 \`\`\`ts
 import { ZenStackClient } from '@zenstackhq/orm';
@@ -47,7 +142,10 @@ const client = new ZenStackClient(schema, {
 \`\`\`
 
 Check documentation: https://zenstack.dev/docs/`);
+        }
     }
+
+    return model;
 }
 
 function getOutputPath(options: Options, schemaFile: string) {
diff --git a/packages/cli/src/index.ts b/packages/cli/src/index.ts
index c2307fa1d..0d663044c 100644
--- a/packages/cli/src/index.ts
+++ b/packages/cli/src/index.ts
@@ -68,6 +68,7 @@ function createProgram() {
         .addOption(schemaOption)
         .addOption(noVersionCheckOption)
         .addOption(new Option('-o, --output <path>', 'default output directory for code generation'))
+        .addOption(new Option('-w, --watch', 'enable watch mode').default(false))
         .addOption(new Option('--lite', 'also generate a lite version of schema without attributes').default(false))
         .addOption(new Option('--lite-only', 'only generate lite version of schema without attributes').default(false))
         .addOption(new Option('--silent', 'suppress all output except errors').default(false))
@@ -220,6 +221,11 @@ async function main() {
         }
     }
 
+    if (program.args.includes('generate') && (program.args.includes('-w') || program.args.includes('--watch'))) {
+        // A "hack" way to prevent the process from terminating because we don't want to stop it.
+        return;
+    }
+
     if (telemetry.isTracking) {
         // give telemetry a chance to send events before exit
         setTimeout(() => {
diff --git a/packages/common-helpers/src/index.ts b/packages/common-helpers/src/index.ts
index 07c4fff56..146609fb7 100644
--- a/packages/common-helpers/src/index.ts
+++ b/packages/common-helpers/src/index.ts
@@ -4,6 +4,7 @@ export * from './is-plain-object';
 export * from './lower-case-first';
 export * from './param-case';
 export * from './safe-json-stringify';
+export * from './single-debounce';
 export * from './sleep';
 export * from './tiny-invariant';
 export * from './upper-case-first';
diff --git a/packages/common-helpers/src/single-debounce.ts b/packages/common-helpers/src/single-debounce.ts
new file mode 100644
index 000000000..f16091e5c
--- /dev/null
+++ b/packages/common-helpers/src/single-debounce.ts
@@ -0,0 +1,34 @@
+export function singleDebounce(cb: () => void | PromiseLike<void>, debounceMc: number, reRunOnInProgressCall: boolean = false) {
+    let timeout: ReturnType<typeof setTimeout> | undefined;
+    let inProgress = false;
+    let pendingInProgress = false;
+
+    const run = async () => {
+        if (inProgress) {
+            if (reRunOnInProgressCall) {
+                pendingInProgress = true;
+            }
+
+            return;
+        }
+
+        inProgress = true;
+        pendingInProgress = false;
+
+        try {
+            await cb();
+        } finally {
+            inProgress = false;
+
+            if (pendingInProgress) {
+                await run();
+            }
+        }
+    };
+
+    return () => {
+        clearTimeout(timeout);
+
+        timeout = setTimeout(run, debounceMc);
+    }
+}
diff --git a/packages/orm/src/client/crud/validator/index.ts b/packages/orm/src/client/crud/validator/index.ts
index dc1a6f836..8b16a9de1 100644
--- a/packages/orm/src/client/crud/validator/index.ts
+++ b/packages/orm/src/client/crud/validator/index.ts
@@ -382,9 +382,13 @@ export class InputValidator<Schema extends SchemaDef> {
         // zod doesn't preserve object field order after parsing, here we use a
         // validation-only custom schema and use the original data if parsing
         // is successful
-        const finalSchema = z.custom((v) => {
-            return schema.safeParse(v).success;
+        const finalSchema = z.any().superRefine((value, ctx) => {
+            const parseResult = schema.safeParse(value);
+            if (!parseResult.success) {
+                parseResult.error.issues.forEach((issue) => ctx.addIssue(issue as any));
+            }
         });
+
         this.setSchemaCache(key!, finalSchema);
         return finalSchema;
     }
@@ -495,7 +499,7 @@ export class InputValidator<Schema extends SchemaDef> {
         }
 
         // expression builder
-        fields['$expr'] = z.custom((v) => typeof v === 'function').optional();
+        fields['$expr'] = z.custom((v) => typeof v === 'function', { error: '"$expr" must be a function' }).optional();
 
         // logical operators
         fields['AND'] = this.orArray(
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index ba5f04d5d..df889c6bd 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -183,6 +183,9 @@ importers:
       '@zenstackhq/sdk':
         specifier: workspace:*
         version: link:../sdk
+      chokidar:
+        specifier: ^5.0.0
+        version: 5.0.0
       colors:
         specifier: 1.4.0
         version: 1.4.0
diff --git a/tests/e2e/orm/client-api/typed-json-fields.test.ts b/tests/e2e/orm/client-api/typed-json-fields.test.ts
index a213c1d82..f5a8945c1 100644
--- a/tests/e2e/orm/client-api/typed-json-fields.test.ts
+++ b/tests/e2e/orm/client-api/typed-json-fields.test.ts
@@ -121,7 +121,7 @@ model User {
                     },
                 },
             }),
-        ).rejects.toThrow(/invalid/i);
+        ).rejects.toThrow('data.identity.providers[0].id');
     });
 
     it('works with find', async () => {
diff --git a/tests/regression/test/issue-558.test.ts b/tests/regression/test/issue-558.test.ts
new file mode 100644
index 000000000..4a76e31f4
--- /dev/null
+++ b/tests/regression/test/issue-558.test.ts
@@ -0,0 +1,19 @@
+import { createTestClient } from '@zenstackhq/testtools';
+import { describe, expect, it } from 'vitest';
+
+describe('Regression for issue #558', () => {
+    it('verifies issue 558', async () => {
+        const db = await createTestClient(`
+type Foo {
+    x Int
+}
+
+model Model {
+    id String @id @default(cuid())
+    foo Foo @json
+}
+        `);
+
+        await expect(db.model.create({ data: { foo: { x: 'hello' } } })).rejects.toThrow('data.foo.x');
+    });
+});