From 7ba354e8db8b8bd38f329d35d6ab9285850d42d8 Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Fri, 15 Aug 2025 20:56:15 +0800
Subject: [PATCH] fix: find input args validation

---
 packages/runtime/src/client/crud/operations/base.ts |  2 ++
 packages/runtime/src/client/crud/validator.ts       | 10 +++++-----
 packages/runtime/src/client/plugin.ts               |  4 ++--
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/packages/runtime/src/client/crud/operations/base.ts b/packages/runtime/src/client/crud/operations/base.ts
index 24874d69..f11ad80c 100644
--- a/packages/runtime/src/client/crud/operations/base.ts
+++ b/packages/runtime/src/client/crud/operations/base.ts
@@ -68,6 +68,8 @@ export type CrudOperation =
     | 'aggregate'
     | 'groupBy';
 
+export type AllCrudOperation = CrudOperation | 'findUniqueOrThrow' | 'findFirstOrThrow';
+
 export type FromRelationContext<Schema extends SchemaDef> = {
     model: GetModels<Schema>;
     field: string;
diff --git a/packages/runtime/src/client/crud/validator.ts b/packages/runtime/src/client/crud/validator.ts
index 7ef9d677..be83d102 100644
--- a/packages/runtime/src/client/crud/validator.ts
+++ b/packages/runtime/src/client/crud/validator.ts
@@ -44,7 +44,7 @@ export class InputValidator<Schema extends SchemaDef> {
         return this.validate<FindArgs<Schema, GetModels<Schema>, true>, Parameters<typeof this.makeFindSchema>[1]>(
             model,
             'find',
-            { unique, collection: true },
+            { unique },
             (model, options) => this.makeFindSchema(model, options),
             args,
         );
@@ -196,7 +196,7 @@ export class InputValidator<Schema extends SchemaDef> {
 
     // #region Find
 
-    private makeFindSchema(model: string, options: { unique: boolean; collection: boolean }) {
+    private makeFindSchema(model: string, options: { unique: boolean }) {
         const fields: Record<string, z.ZodSchema> = {};
         const where = this.makeWhereSchema(model, options.unique);
         if (options.unique) {
@@ -208,13 +208,13 @@ export class InputValidator<Schema extends SchemaDef> {
         fields['select'] = this.makeSelectSchema(model).optional();
         fields['include'] = this.makeIncludeSchema(model).optional();
         fields['omit'] = this.makeOmitSchema(model).optional();
-        fields['distinct'] = this.makeDistinctSchema(model).optional();
-        fields['cursor'] = this.makeCursorSchema(model).optional();
 
-        if (options.collection) {
+        if (!options.unique) {
             fields['skip'] = this.makeSkipSchema().optional();
             fields['take'] = this.makeTakeSchema().optional();
             fields['orderBy'] = this.orArray(this.makeOrderBySchema(model, true, false), true).optional();
+            fields['cursor'] = this.makeCursorSchema(model).optional();
+            fields['distinct'] = this.makeDistinctSchema(model).optional();
         }
 
         let result: ZodType = z.strictObject(fields);
diff --git a/packages/runtime/src/client/plugin.ts b/packages/runtime/src/client/plugin.ts
index 7f087a5d..99ee2d92 100644
--- a/packages/runtime/src/client/plugin.ts
+++ b/packages/runtime/src/client/plugin.ts
@@ -2,7 +2,7 @@ import type { OperationNode, QueryResult, RootOperationNode, UnknownRow } from '
 import type { ClientContract, ToKysely } from '.';
 import type { GetModels, SchemaDef } from '../schema';
 import type { MaybePromise } from '../utils/type-utils';
-import type { CrudOperation } from './crud/operations/base';
+import type { AllCrudOperation } from './crud/operations/base';
 
 /**
  * ZenStack runtime plugin.
@@ -61,7 +61,7 @@ type OnQueryHookContext<Schema extends SchemaDef> = {
     /**
      * The operation that is being performed.
      */
-    operation: CrudOperation;
+    operation: AllCrudOperation;
 
     /**
      * The query arguments.