Expose invocation id of calls

Author: igalshilman

packages/restate-sdk/src/common_api.ts

@@ -18,6 +18,9 @@ export type {
   Rand,
   GenericCall,
   GenericSend,
+  InvocationId,
+  InvocationHandle,
+  InvocationPromise,
 } from "./context.js";
 export { CombineablePromise } from "./context.js";
 

packages/restate-sdk/src/context.ts

@@ -537,7 +537,7 @@ export interface Context extends RestateContext {
 
   genericCall<REQ = Uint8Array, RES = Uint8Array>(
     call: GenericCall<REQ, RES>
-  ): CombineablePromise<RES>;
+  ): InvocationPromise<RES>;
 
   genericSend<REQ = Uint8Array>(call: GenericSend<REQ>): void;
 
@@ -633,6 +633,14 @@ export type CombineablePromise<T> = Promise<T> & {
   orTimeout(millis: number): CombineablePromise<T>;
 };
 
+export type InvocationId = string & { __brand: "InvocationId" };
+
+export type InvocationHandle = {
+  invocationId(): Promise<InvocationId>;
+};
+
+export type InvocationPromise<T> = CombineablePromise<T> & InvocationHandle;
+
 export const CombineablePromise = {
   /**
    * Creates a Promise that is resolved with an array of results when all of the provided Promises

packages/restate-sdk/src/context_impl.ts

@@ -17,6 +17,8 @@ import type {
   DurablePromise,
   GenericCall,
   GenericSend,
+  InvocationPromise,
+  InvocationId,
   ObjectContext,
   Rand,
   Request,
@@ -63,8 +65,10 @@ import {
   extractContext,
   pendingPromise,
   PromisesExecutor,
+  RestateInvocationPromise,
   RestateCombinatorPromise,
   RestatePendingPromise,
+  InvocationPendingPromise,
   RestateSinglePromise,
 } from "./promises.js";
 import { InputPump, OutputPump } from "./io.js";
@@ -168,13 +172,14 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
   //
   public genericCall<REQ = Uint8Array, RES = Uint8Array>(
     call: GenericCall<REQ, RES>
-  ): CombineablePromise<RES> {
+  ): InvocationPromise<RES> {
     const requestSerde: Serde<REQ> =
       call.inputSerde ?? (serde.binary as Serde<REQ>);
     const responseSerde: Serde<RES> =
       call.outputSerde ?? (serde.binary as Serde<RES>);
 
-    return this.processCompletableEntry((vm) => {
+    try {
+      const vm = this.coreVm;
       const parameter = requestSerde.serialize(call.parameter);
       const call_handles = vm.sys_call(
         call.service,
@@ -188,13 +193,25 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
           : []
       );
 
-      // TODO eventually we return this promise back to the user
-      // const invocationIdPromise = this.createInvocationIdPromise(
-      //   call_handles.invocation_id_completion_id
-      // );
+      const invocationIdHandle = call_handles.invocation_id_completion_id;
+      const invocationHandle = () => {
+        // create a promise, on demand that resolves to an invocation id if needed.
+        return this.createInvocationIdPromise(invocationIdHandle);
+      };
+
+      const callHandle = call_handles.call_completion_id;
 
-      return call_handles.call_completion_id;
-    }, completeUsing(SuccessWithSerde(responseSerde), Failure));
+      return new RestateInvocationPromise(
+        this,
+        callHandle,
+        completeUsing(SuccessWithSerde(responseSerde), Failure),
+        invocationHandle
+      );
+    } catch (e) {
+      this.handleInvocationEndError(e);
+      // We return a pending promise to avoid the caller to see the error.
+      return new InvocationPendingPromise(this);
+    }
   }
 
   public genericSend<REQ = Uint8Array>(send: GenericSend<REQ>) {
@@ -228,7 +245,7 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
 
   private createInvocationIdPromise(
     handle: number
-  ): RestateSinglePromise<string> {
+  ): RestateSinglePromise<InvocationId> {
     return new RestateSinglePromise(this, handle, completeUsing(InvocationId));
   }
 
@@ -696,17 +713,25 @@ const VoidAsUndefined: Completer = (value, prom) => {
   return false;
 };
 
-function SuccessWithSerde<T>(serde?: Serde<T>): Completer {
+function SuccessWithSerde<T>(
+  serde?: Serde<T>,
+  transform?: <U>(success: T) => U
+): Completer {
   return (value, prom) => {
-    if (typeof value === "object" && "Success" in value) {
-      if (!serde) {
-        prom.resolve(defaultSerde<T>().deserialize(value.Success));
-      } else {
-        prom.resolve(serde.deserialize(value.Success));
-      }
-      return true;
+    if (typeof value !== "object" || !("Success" in value)) {
+      return false;
     }
-    return false;
+    let val: T;
+    if (serde) {
+      val = serde.deserialize(value.Success);
+    } else {
+      val = defaultSerde<T>().deserialize(value.Success);
+    }
+    if (transform) {
+      val = transform(val);
+    }
+    prom.resolve(val);
+    return true;
   };
 }
 

packages/restate-sdk/src/promises.ts

@@ -11,7 +11,11 @@
 
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import type { CombineablePromise } from "./context.js";
+import type {
+  CombineablePromise,
+  InvocationId,
+  InvocationPromise,
+} from "./context.js";
 import type * as vm from "./endpoint/handlers/vm/sdk_shared_core_wasm_bindings.js";
 import { CancelledError, TimeoutError } from "./types/errors.js";
 import { CompletablePromise } from "./utils/completable_promise.js";
@@ -191,6 +195,24 @@ export class RestateSinglePromise<T> extends AbstractRestatePromise<T> {
   readonly [Symbol.toStringTag] = "RestateSinglePromise";
 }
 
+export class RestateInvocationPromise<T>
+  extends RestateSinglePromise<T>
+  implements InvocationPromise<T>
+{
+  constructor(
+    ctx: ContextImpl,
+    handle: number,
+    completer: (value: AsyncResultValue, prom: CompletablePromise<T>) => void,
+    private readonly invocationIdSupplier: () => Promise<InvocationId>
+  ) {
+    super(ctx, handle, completer);
+  }
+
+  invocationId(): Promise<InvocationId> {
+    return this.invocationIdSupplier();
+  }
+}
+
 export class RestateCombinatorPromise extends AbstractRestatePromise<any> {
   private state: PromiseState = PromiseState.NOT_COMPLETED;
   private readonly combinatorPromise: Promise<any>;
@@ -278,6 +300,19 @@ export class RestatePendingPromise<T> implements RestatePromise<T> {
   readonly [Symbol.toStringTag] = "RestatePendingPromise";
 }
 
+export class InvocationPendingPromise<T>
+  extends RestatePendingPromise<T>
+  implements InvocationPromise<T>
+{
+  constructor(ctx: ContextImpl) {
+    super(ctx);
+  }
+
+  invocationId(): Promise<InvocationId> {
+    return pendingPromise();
+  }
+}
+
 /**
  * Promises executor, gluing VM with I/O and Promises given to user space.
  */

packages/restate-sdk/src/types/rpc.ts

@@ -13,10 +13,10 @@
 /* eslint-disable @typescript-eslint/no-namespace */
 /* eslint-disable @typescript-eslint/ban-types */
 import type {
-  CombineablePromise,
   Context,
   GenericCall,
   GenericSend,
+  InvocationPromise,
   ObjectContext,
   ObjectSharedContext,
   WorkflowContext,
@@ -218,7 +218,7 @@ export type Client<M> = {
   ) => PromiseLike<infer O>
     ? (
         ...args: [...P, ...[opts?: Opts<InferArg<P>, O>]]
-      ) => CombineablePromise<O>
+      ) => InvocationPromise<O>
     : never;
 };