From 174ccc8fb60c7ea28819d09beb84e8565f29e888 Mon Sep 17 00:00:00 2001
From: Geoffry Song <goffrie@gmail.com>
Date: Tue, 3 Dec 2024 14:13:29 -0800
Subject: [PATCH] Allow handler to return a promise when using
 PropertyValidators (#31895)

GitOrigin-RevId: f210a4fb93772b0fa2c4dbca3561622e5ea2d157
---
 src/server/registration.test.ts | 30 ++++++++++++++++++++++++++++++
 src/server/registration.ts      |  2 +-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/server/registration.test.ts b/src/server/registration.test.ts
index f08bb0c..42b28a8 100644
--- a/src/server/registration.test.ts
+++ b/src/server/registration.test.ts
@@ -430,6 +430,36 @@ describe("argument and return value validators can be objects or validators", ()
       },
     }),
 
+    queryAsync: query({
+      args: v.object({
+        arg: v.string(),
+      }),
+      returns: { arg: v.string() },
+      handler: async (_, { arg }) => {
+        return { arg: arg };
+      },
+    }),
+
+    mutationAsync: mutation({
+      args: v.object({
+        arg: v.string(),
+      }),
+      returns: { arg: v.string() },
+      handler: async (_, { arg }) => {
+        return { arg: arg };
+      },
+    }),
+
+    actionAsync: action({
+      args: v.object({
+        arg: v.string(),
+      }),
+      returns: { arg: v.string() },
+      handler: async (_, { arg }) => {
+        return { arg: arg };
+      },
+    }),
+
     // This is syntx that we no longer want to support when typechecking because they result in undefined behavior.
     mutationNoOptionalValidators: mutation({
       // @ts-expect-error Optional validators are not supported at the top level
diff --git a/src/server/registration.ts b/src/server/registration.ts
index 072a1ea..edec289 100644
--- a/src/server/registration.ts
+++ b/src/server/registration.ts
@@ -583,7 +583,7 @@ export type ReturnValueForOptionalValidator<
 > = [ReturnsValidator] extends [Validator<any, any, any>]
   ? ValidatorTypeToReturnType<Infer<ReturnsValidator>>
   : [ReturnsValidator] extends [PropertyValidators]
-    ? ObjectType<ReturnsValidator>
+    ? ValidatorTypeToReturnType<ObjectType<ReturnsValidator>>
     : any;
 
 export type ArgsArrayForOptionalValidator<