fix: handle server-sent termination in RPC

aws · Oct 22, 2023 · 8c43067 · 8c43067
1 parent 86f1134
commit 8c43067
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/lib/eventstream_rpc.ts b/lib/eventstream_rpc.ts
@@ -785,6 +785,13 @@ class OperationBase extends EventEmitter {
      * @return the operation's underlying event stream binding object
      */
     getStream() : eventstream.ClientStream { return this.stream; }
+
+    /**
+     * Set this operation state to be "Ended" so that closing the operation will not send a terminate message.
+     */
+    setStateEnded() {
+        this.state = OperationState.Ended;
+    }
 }
 
 /**
@@ -837,6 +844,11 @@ export class RequestResponseOperation<RequestType, ResponseType> extends EventEm
                 await this.operation.activate(requestMessage);
 
                 let message : eventstream.Message = await responsePromise;
+
+                // If the server terminated the stream, then set the operation to be ended immediately
+                if ((message.flags ?? 0) & eventstream.MessageFlags.TerminateStream) {
+                    this.operation.setStateEnded();
+                }
                 let response : ResponseType = deserializeResponse(this.serviceModel, this.operationConfig.name, message);
 
                 resolve(response);
@@ -926,6 +938,10 @@ export class StreamingOperation<RequestType, ResponseType, OutboundMessageType,
                 await this.operation.activate(requestMessage);
 
                 let message : eventstream.Message = await responsePromise;
+                // If the server terminated the stream, then set the operation to be ended immediately
+                if ((message.flags ?? 0) & eventstream.MessageFlags.TerminateStream) {
+                    this.operation.setStateEnded();
+                }
                 let response : ResponseType = deserializeResponse(this.serviceModel, this.operationConfig.name, message);
 
                 resolve(response);