1 files changed, 47 insertions, 34 deletions

diff --git a/js/dispatch.ts b/js/dispatch.ts
index 0c6e70709..36f97363f 100644
--- a/js/dispatch.ts
+++ b/js/dispatch.ts
@@ -5,36 +5,30 @@ import * as msg from "gen/cli/msg_generated";
 import * as errors from "./errors";
 import * as util from "./util";
 import {
-  nextPromiseId,
   recordFromBufMinimal,
   handleAsyncMsgFromRustMinimal
 } from "./dispatch_minimal";
 
 const promiseTable = new Map<number, util.Resolvable<msg.Base>>();
 
-interface FlatbufferRecord {
-  promiseId: number;
-  base: msg.Base;
-}
-
-function flatbufferRecordFromBuf(buf: Uint8Array): FlatbufferRecord {
+function flatbufferRecordFromBuf(buf: Uint8Array): msg.Base {
   const bb = new flatbuffers.ByteBuffer(buf);
   const base = msg.Base.getRootAsBase(bb);
-  return {
-    promiseId: base.cmdId(),
-    base
-  };
+  return base;
 }
 
-export function handleAsyncMsgFromRust(ui8: Uint8Array): void {
+export function handleAsyncMsgFromRust(
+  promiseId: number,
+  ui8: Uint8Array
+): void {
   const buf32 = new Int32Array(ui8.buffer, ui8.byteOffset, ui8.byteLength / 4);
   const recordMin = recordFromBufMinimal(buf32);
   if (recordMin) {
     // Fast and new
-    handleAsyncMsgFromRustMinimal(ui8, recordMin);
+    handleAsyncMsgFromRustMinimal(promiseId, ui8, recordMin);
   } else {
     // Legacy
-    let { promiseId, base } = flatbufferRecordFromBuf(ui8);
+    let base = flatbufferRecordFromBuf(ui8);
     const promise = promiseTable.get(promiseId);
     util.assert(promise != null, `Expecting promise in table. ${promiseId}`);
     promiseTable.delete(promiseId);
@@ -56,14 +50,26 @@ function sendInternal(
   innerType: msg.Any,
   inner: flatbuffers.Offset,
   zeroCopy: undefined | ArrayBufferView,
-  sync = true
-): [number, null | Uint8Array] {
-  const cmdId = nextPromiseId();
+  isSync: true
+): Uint8Array | null;
+function sendInternal(
+  builder: flatbuffers.Builder,
+  innerType: msg.Any,
+  inner: flatbuffers.Offset,
+  zeroCopy: undefined | ArrayBufferView,
+  isSync: false
+): Promise<msg.Base>;
+function sendInternal(
+  builder: flatbuffers.Builder,
+  innerType: msg.Any,
+  inner: flatbuffers.Offset,
+  zeroCopy: undefined | ArrayBufferView,
+  isSync: boolean
+): Promise<msg.Base> | Uint8Array | null {
   msg.Base.startBase(builder);
+  msg.Base.addSync(builder, isSync);
   msg.Base.addInner(builder, inner);
   msg.Base.addInnerType(builder, innerType);
-  msg.Base.addSync(builder, sync);
-  msg.Base.addCmdId(builder, cmdId);
   builder.finish(msg.Base.endBase(builder));
 
   const control = builder.asUint8Array();
@@ -74,7 +80,25 @@ function sendInternal(
   );
 
   builder.inUse = false;
-  return [cmdId, response];
+
+  if (typeof response === "number") {
+    const promise = util.createResolvable<msg.Base>();
+    promiseTable.set(response, promise);
+    util.assert(!isSync);
+    return promise;
+  } else {
+    if (!isSync) {
+      util.assert(response !== null);
+      const base = flatbufferRecordFromBuf(response as Uint8Array);
+      const err = errors.maybeError(base);
+      if (err != null) {
+        return Promise.reject(err);
+      } else {
+        return Promise.resolve(base);
+      }
+    }
+    return response;
+  }
 }
 
 // @internal
@@ -84,16 +108,7 @@ export function sendAsync(
   inner: flatbuffers.Offset,
   data?: ArrayBufferView
 ): Promise<msg.Base> {
-  const [cmdId, response] = sendInternal(
-    builder,
-    innerType,
-    inner,
-    data,
-    false
-  );
-  util.assert(response == null); // null indicates async.
-  const promise = util.createResolvable<msg.Base>();
-  promiseTable.set(cmdId, promise);
+  const promise = sendInternal(builder, innerType, inner, data, false);
   return promise;
 }
 
@@ -104,10 +119,8 @@ export function sendSync(
   inner: flatbuffers.Offset,
   data?: ArrayBufferView
 ): null | msg.Base {
-  const [cmdId, response] = sendInternal(builder, innerType, inner, data, true);
-  util.assert(cmdId >= 0);
-  util.assert(response != null); // null indicates async.
-  if (response!.length === 0) {
+  const response = sendInternal(builder, innerType, inner, data, true);
+  if (response == null || response.length === 0) {
     return null;
   } else {
     const bb = new flatbuffers.ByteBuffer(response!);