From 534fcc5d0ab1f938d998a3410ec7a8fea0e3ee9a Mon Sep 17 00:00:00 2001
From: Mike Sawka <mike@commandline.dev>
Date: Mon, 19 Aug 2024 15:44:30 -0700
Subject: [PATCH] client support for rpc cancel (#249)

---
 frontend/app/store/wshrpc.ts          | 12 ++++++-
 pkg/tsgen/tsgen.go                    |  7 +++++
 pkg/wshrpc/wshclient/wshclientutil.go | 45 ++++++++++++++-------------
 pkg/wshrpc/wshrpctypes.go             |  2 ++
 pkg/wshutil/wshrpc.go                 |  8 +++++
 5 files changed, 52 insertions(+), 22 deletions(-)

diff --git a/frontend/app/store/wshrpc.ts b/frontend/app/store/wshrpc.ts
index ea172933..bdcac511 100644
--- a/frontend/app/store/wshrpc.ts
+++ b/frontend/app/store/wshrpc.ts
@@ -51,7 +51,11 @@ async function* rpcResponseGenerator(
                     return;
                 }
                 const shouldTerminate = yield msg.data;
-                if (shouldTerminate || !msg.cont) {
+                if (shouldTerminate) {
+                    sendRpcCancel(reqid);
+                    return;
+                }
+                if (!msg.cont) {
                     return;
                 }
             }
@@ -65,6 +69,12 @@ async function* rpcResponseGenerator(
     }
 }
 
+function sendRpcCancel(reqid: string) {
+    const rpcMsg: RpcMessage = { reqid: reqid, cancel: true };
+    const wsMsg: WSRpcCommand = { wscommand: "rpc", message: rpcMsg };
+    globalWS.pushMessage(wsMsg);
+}
+
 function sendRpcCommand(msg: RpcMessage): AsyncGenerator<RpcMessage, void, boolean> {
     const wsMsg: WSRpcCommand = { wscommand: "rpc", message: msg };
     globalWS.pushMessage(wsMsg);
diff --git a/pkg/tsgen/tsgen.go b/pkg/tsgen/tsgen.go
index 9f4be100..3a0efe88 100644
--- a/pkg/tsgen/tsgen.go
+++ b/pkg/tsgen/tsgen.go
@@ -79,6 +79,13 @@ func generateTSMethodTypes(method reflect.Method, tsTypesMap map[reflect.Type]st
 }
 
 func getTSFieldName(field reflect.StructField) string {
+	tsFieldTag := field.Tag.Get("tsfield")
+	if tsFieldTag != "" {
+		if tsFieldTag == "-" {
+			return ""
+		}
+		return tsFieldTag
+	}
 	jsonTag := field.Tag.Get("json")
 	if jsonTag != "" {
 		parts := strings.Split(jsonTag, ",")
diff --git a/pkg/wshrpc/wshclient/wshclientutil.go b/pkg/wshrpc/wshclient/wshclientutil.go
index a5066a34..033a9e7f 100644
--- a/pkg/wshrpc/wshclient/wshclientutil.go
+++ b/pkg/wshrpc/wshclient/wshclientutil.go
@@ -57,27 +57,30 @@ func sendRpcRequestResponseStreamHelper[T any](w *wshutil.WshRpc, command string
 	if err != nil {
 		rtnErr(respChan, err)
 		return respChan
-	} else {
-		go func() {
-			defer close(respChan)
-			for {
-				if reqHandler.ResponseDone() {
-					break
-				}
-				resp, err := reqHandler.NextResponse()
-				if err != nil {
-					respChan <- wshrpc.RespOrErrorUnion[T]{Error: err}
-					break
-				}
-				var respData T
-				err = utilfn.ReUnmarshal(&respData, resp)
-				if err != nil {
-					respChan <- wshrpc.RespOrErrorUnion[T]{Error: err}
-					break
-				}
-				respChan <- wshrpc.RespOrErrorUnion[T]{Response: respData}
-			}
-		}()
 	}
+	opts.StreamCancelFn = func() {
+		// TODO coordinate the cancel with the for loop below
+		reqHandler.SendCancel()
+	}
+	go func() {
+		defer close(respChan)
+		for {
+			if reqHandler.ResponseDone() {
+				break
+			}
+			resp, err := reqHandler.NextResponse()
+			if err != nil {
+				respChan <- wshrpc.RespOrErrorUnion[T]{Error: err}
+				break
+			}
+			var respData T
+			err = utilfn.ReUnmarshal(&respData, resp)
+			if err != nil {
+				respChan <- wshrpc.RespOrErrorUnion[T]{Error: err}
+				break
+			}
+			respChan <- wshrpc.RespOrErrorUnion[T]{Response: respData}
+		}
+	}()
 	return respChan
 }
diff --git a/pkg/wshrpc/wshrpctypes.go b/pkg/wshrpc/wshrpctypes.go
index 489a9307..0ebd71f1 100644
--- a/pkg/wshrpc/wshrpctypes.go
+++ b/pkg/wshrpc/wshrpctypes.go
@@ -105,6 +105,8 @@ type RpcOpts struct {
 	Timeout    int    `json:"timeout,omitempty"`
 	NoResponse bool   `json:"noresponse,omitempty"`
 	Route      string `json:"route,omitempty"`
+
+	StreamCancelFn func() `json:"-"` // this is an *output* parameter, set by the handler
 }
 
 type RpcContext struct {
diff --git a/pkg/wshutil/wshrpc.go b/pkg/wshutil/wshrpc.go
index f15b67f3..fc86e844 100644
--- a/pkg/wshutil/wshrpc.go
+++ b/pkg/wshutil/wshrpc.go
@@ -76,6 +76,14 @@ func GetRpcSourceFromContext(ctx context.Context) string {
 	return rtn.(*RpcResponseHandler).GetSource()
 }
 
+func GetIsCanceledFromContext(ctx context.Context) bool {
+	rtn := ctx.Value(wshRpcRespHandlerContextKey{})
+	if rtn == nil {
+		return false
+	}
+	return rtn.(*RpcResponseHandler).IsCanceled()
+}
+
 func GetRpcResponseHandlerFromContext(ctx context.Context) *RpcResponseHandler {
 	rtn := ctx.Value(wshRpcRespHandlerContextKey{})
 	if rtn == nil {