From 62b94228342980646a233d75f3e0f7e712038856 Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Thu, 10 Apr 2025 03:38:49 +0000
Subject: [PATCH] feat(infra): op call with signal (#11567)

---
 packages/common/infra/src/op/client.ts | 24 +++++++++++++++++++-----
 packages/common/infra/src/op/types.ts  |  5 +++++
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/packages/common/infra/src/op/client.ts b/packages/common/infra/src/op/client.ts
index 806b681e1a..8034225a2f 100644
--- a/packages/common/infra/src/op/client.ts
+++ b/packages/common/infra/src/op/client.ts
@@ -11,7 +11,13 @@ import {
   type SubscribeMessage,
   type UnsubscribeMessage,
 } from './message';
-import type { OpInput, OpNames, OpOutput, OpSchema } from './types';
+import type {
+  OpInput,
+  OpInputWithSignal,
+  OpNames,
+  OpOutput,
+  OpSchema,
+} from './types';
 
 export interface CancelablePromise<T> extends Promise<T> {
   cancel(): void;
@@ -107,10 +113,14 @@ export class OpClient<Ops extends OpSchema> extends AutoMessageHandler {
 
   call<Op extends OpNames<Ops>>(
     op: Op,
-    ...args: OpInput<Ops, Op>
+    ...args: OpInputWithSignal<Ops, Op>
   ): CancelablePromise<OpOutput<Ops, Op>> {
     const promiseWithResolvers = Promise.withResolvers<any>();
-    const payload = args[0];
+    const abortSignal =
+      args[args.length - 1] instanceof AbortSignal
+        ? (args.pop() as AbortSignal)
+        : undefined;
+    const payload = args.pop();
 
     const msg = {
       type: 'call',
@@ -121,7 +131,7 @@ export class OpClient<Ops extends OpSchema> extends AutoMessageHandler {
 
     const promise = promiseWithResolvers.promise as CancelablePromise<any>;
 
-    const raise = (reason: string) => {
+    const raise = (reason: any) => {
       const pending = this.pendingCalls.get(msg.id);
       if (!pending) {
         return;
@@ -130,11 +140,15 @@ export class OpClient<Ops extends OpSchema> extends AutoMessageHandler {
         type: 'cancel',
         id: msg.id,
       } satisfies CancelMessage);
-      promiseWithResolvers.reject(new Error(reason));
+      promiseWithResolvers.reject(reason);
       clearTimeout(pending.timeout);
       this.pendingCalls.delete(msg.id);
     };
 
+    abortSignal?.addEventListener('abort', () => {
+      raise(abortSignal.reason);
+    });
+
     promise.cancel = () => {
       raise('canceled');
     };
diff --git a/packages/common/infra/src/op/types.ts b/packages/common/infra/src/op/types.ts
index 7de14e03e2..cae15e68bd 100644
--- a/packages/common/infra/src/op/types.ts
+++ b/packages/common/infra/src/op/types.ts
@@ -26,6 +26,11 @@ export type OpInput<
       : never
   : never;
 
+export type OpInputWithSignal<Ops extends OpSchema, Type extends OpNames<Ops>> =
+  OpInput<Ops, Type> extends [infer In]
+    ? [In, AbortSignal | undefined] | [In]
+    : [AbortSignal | undefined] | [];
+
 export type OpOutput<
   Ops extends OpSchema,
   Type extends OpNames<Ops>,