@restatedev/restate-sdk 1.3.2 → 1.4.0

package/dist/cjs/src/context_impl.js

@@ -11,23 +11,18 @@
  */
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.ContextImpl = void 0;
-const protocol_pb_js_1 = require("./generated/proto/protocol_pb.js");
-const protocol_js_1 = require("./types/protocol.js");
 const errors_js_1 = require("./types/errors.js");
-const protobuf_1 = require("@bufbuild/protobuf");
 const rpc_js_1 = require("./types/rpc.js");
 const restate_sdk_core_1 = require("@restatedev/restate-sdk-core");
 const rand_js_1 = require("./utils/rand.js");
-const promise_combinator_tracker_js_1 = require("./promise_combinator_tracker.js");
-const node_buffer_1 = require("node:buffer");
 class ContextImpl {
+    coreVm;
+    input;
     console;
     handlerKind;
-    keyedContextKey;
-    stateMachine;
-    // This is used to guard users against calling ctx.sideEffect without awaiting it.
-    // See https://github.com/restatedev/sdk-typescript/issues/197 for more details.
-    executingRun = false;
+    invocationEndPromise;
+    inputReader;
+    outputWriter;
     invocationRequest;
     rand;
     date = {
@@ -38,137 +33,122 @@ class ContextImpl {
             return this.run(() => new Date().toJSON());
         },
     };
-    constructor(id, console, handlerKind, keyedContextKey, invocationValue, invocationHeaders, attemptHeaders, extraArgs, 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    stateMachine) {
+    currentRead;
+    constructor(coreVm, input, console, handlerKind, attemptHeaders, extraArgs, invocationEndPromise, inputReader, outputWriter) {
+        this.coreVm = coreVm;
+        this.input = input;
         this.console = console;
         this.handlerKind = handlerKind;
-        this.keyedContextKey = keyedContextKey;
-        this.stateMachine = stateMachine;
+        this.invocationEndPromise = invocationEndPromise;
+        this.inputReader = inputReader;
+        this.outputWriter = outputWriter;
         this.invocationRequest = {
-            id,
-            headers: invocationHeaders,
-            attemptHeaders,
-            body: invocationValue,
+            id: input.invocation_id,
+            headers: input.headers.reduce((headers, { key, value }) => {
+                headers.set(key, value);
+                return headers;
+            }, new Map()),
+            attemptHeaders: Object.entries(attemptHeaders).reduce((headers, [key, value]) => {
+                if (value !== undefined) {
+                    headers.set(key, value instanceof Array ? value[0] : value);
+                }
+                return headers;
+            }, new Map()),
+            body: input.input,
             extraArgs,
         };
-        this.rand = new rand_js_1.RandImpl(id, this.checkState.bind(this));
-    }
-    promise(name, serde) {
-        return new DurablePromiseImpl(this, name, serde);
+        this.rand = new rand_js_1.RandImpl(input.invocation_id, () => {
+            if (coreVm.is_inside_run()) {
+                throw new Error("Cannot generate random numbers within a run closure. Use the random object outside the run closure.");
+            }
+        });
     }
     get key() {
         switch (this.handlerKind) {
             case rpc_js_1.HandlerKind.EXCLUSIVE:
             case rpc_js_1.HandlerKind.SHARED:
             case rpc_js_1.HandlerKind.WORKFLOW: {
-                if (this.keyedContextKey === undefined) {
-                    throw new errors_js_1.TerminalError("unexpected missing key");
-                }
-                return this.keyedContextKey;
+                return this.input.key;
             }
-            case rpc_js_1.HandlerKind.SERVICE:
-                throw new errors_js_1.TerminalError("unexpected missing key");
             default:
-                throw new errors_js_1.TerminalError("unknown handler type");
+                throw new errors_js_1.TerminalError("this handler type doesn't support key()");
         }
     }
     request() {
         return this.invocationRequest;
     }
-    // DON'T make this function async!!! see sideEffect comment for details.
     get(name, serde) {
-        // Check if this is a valid action
-        this.checkState("get state");
-        // Create the message and let the state machine process it
-        const msg = new protocol_pb_js_1.GetStateEntryMessage({
-            key: new TextEncoder().encode(name),
-        });
-        const completed = this.stateMachine.localStateStore.tryCompleteGet(name, msg);
-        const getState = async () => {
-            const result = await this.stateMachine.handleUserCodeMessage(protocol_js_1.GET_STATE_ENTRY_MESSAGE_TYPE, msg, completed);
-            // If the GetState message did not have a value or empty,
-            // then we went to the runtime to get the value.
-            // When we get the response, we set it in the localStateStore,
-            // to answer subsequent requests
-            if (!completed) {
-                this.stateMachine.localStateStore.add(name, result);
-            }
-            if (!(result instanceof Uint8Array)) {
+        return this.processCompletableEntry((vm) => vm.sys_get_state(name), (asyncResultValue) => {
+            if (asyncResultValue === "Empty") {
+                // Empty
                 return null;
             }
-            return (serde ?? (0, rpc_js_1.defaultSerde)()).deserialize(result);
-        };
-        return getState();
+            else if ("Success" in asyncResultValue) {
+                return (serde ?? (0, rpc_js_1.defaultSerde)()).deserialize(asyncResultValue.Success);
+            }
+            else if ("Failure" in asyncResultValue) {
+                throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                    errorCode: asyncResultValue.Failure.code,
+                });
+            }
+            throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
+        });
     }
-    // DON'T make this function async!!! see sideEffect comment for details.
     stateKeys() {
-        // Check if this is a valid action
-        this.checkState("state keys");
-        // Create the message and let the state machine process it
-        const msg = new protocol_pb_js_1.GetStateKeysEntryMessage({});
-        const completed = this.stateMachine.localStateStore.tryCompletedGetStateKeys(msg);
-        const getStateKeys = async () => {
-            const result = await this.stateMachine.handleUserCodeMessage(protocol_js_1.GET_STATE_KEYS_ENTRY_MESSAGE_TYPE, msg, completed);
-            return result.keys.map((b) => new TextDecoder().decode(b));
-        };
-        return getStateKeys();
+        return this.processCompletableEntry((vm) => vm.sys_get_state_keys(), (asyncResultValue) => {
+            if (typeof asyncResultValue === "object" &&
+                "StateKeys" in asyncResultValue) {
+                return asyncResultValue.StateKeys;
+            }
+            else if (typeof asyncResultValue === "object" &&
+                "Failure" in asyncResultValue) {
+                throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                    errorCode: asyncResultValue.Failure.code,
+                });
+            }
+            throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
+        });
     }
     set(name, value, serde) {
-        this.checkState("set state");
-        const bytes = (serde ?? (0, rpc_js_1.defaultSerde)()).serialize(value);
-        const msg = this.stateMachine.localStateStore.set(name, bytes);
-        this.stateMachine
-            .handleUserCodeMessage(protocol_js_1.SET_STATE_ENTRY_MESSAGE_TYPE, msg)
-            .catch((e) => this.stateMachine.handleDanglingPromiseError(e));
+        this.processNonCompletableEntry((vm) => vm.sys_set_state(name, (serde ?? (0, rpc_js_1.defaultSerde)()).serialize(value)));
     }
     clear(name) {
-        this.checkState("clear state");
-        const msg = this.stateMachine.localStateStore.clear(name);
-        this.stateMachine
-            .handleUserCodeMessage(protocol_js_1.CLEAR_STATE_ENTRY_MESSAGE_TYPE, msg)
-            .catch((e) => this.stateMachine.handleDanglingPromiseError(e));
+        this.processNonCompletableEntry((vm) => vm.sys_clear_state(name));
     }
     clearAll() {
-        this.checkState("clear all state");
-        const msg = this.stateMachine.localStateStore.clearAll();
-        this.stateMachine
-            .handleUserCodeMessage(protocol_js_1.CLEAR_ALL_STATE_ENTRY_MESSAGE_TYPE, msg)
-            .catch((e) => this.stateMachine.handleDanglingPromiseError(e));
+        this.processNonCompletableEntry((vm) => vm.sys_clear_all_state());
     }
     // --- Calls, background calls, etc
     //
     genericCall(call) {
         const requestSerde = call.inputSerde ?? restate_sdk_core_1.serde.binary;
         const responseSerde = call.outputSerde ?? restate_sdk_core_1.serde.binary;
-        const parameter = requestSerde.serialize(call.parameter);
-        const msg = new protocol_pb_js_1.CallEntryMessage({
-            serviceName: call.service,
-            handlerName: call.method,
-            parameter,
-            key: call.key,
+        return this.processCompletableEntry((vm) => {
+            const parameter = requestSerde.serialize(call.parameter);
+            return vm.sys_call(call.service, call.method, parameter, call.key);
+        }, (asyncResultValue) => {
+            if (typeof asyncResultValue === "object" &&
+                "Success" in asyncResultValue) {
+                return responseSerde.deserialize(asyncResultValue.Success);
+            }
+            else if (typeof asyncResultValue === "object" &&
+                "Failure" in asyncResultValue) {
+                throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                    errorCode: asyncResultValue.Failure.code,
+                });
+            }
+            throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
         });
-        const rawRequest = this.stateMachine.handleUserCodeMessage(protocol_js_1.INVOKE_ENTRY_MESSAGE_TYPE, msg);
-        const decoded = rawRequest.transform((res) => responseSerde.deserialize(res));
-        return this.markCombineablePromise(decoded);
     }
     genericSend(send) {
-        const requestSerde = send.inputSerde ?? restate_sdk_core_1.serde.binary;
-        const parameter = requestSerde.serialize(send.parameter);
-        const actualDelay = send.delay || 0;
-        const jsInvokeTime = actualDelay > 0 ? Date.now() + actualDelay : protobuf_1.protoInt64.zero;
-        const invokeTime = protobuf_1.protoInt64.parse(jsInvokeTime);
-        const msg = new protocol_pb_js_1.OneWayCallEntryMessage({
-            serviceName: send.service,
-            handlerName: send.method,
-            parameter,
-            invokeTime,
-            key: send.key,
-        });
-        this.stateMachine
-            .handleUserCodeMessage(protocol_js_1.BACKGROUND_INVOKE_ENTRY_MESSAGE_TYPE, msg)
-            .catch((e) => {
-            this.stateMachine.handleDanglingPromiseError(e);
+        this.processNonCompletableEntry((vm) => {
+            const requestSerde = send.inputSerde ?? restate_sdk_core_1.serde.binary;
+            const parameter = requestSerde.serialize(send.parameter);
+            let delay;
+            if (send.delay !== undefined) {
+                delay = BigInt(send.delay);
+            }
+            vm.sys_send(send.service, send.method, parameter, send.key, delay);
         });
     }
     serviceClient({ name }) {
@@ -194,225 +174,406 @@ class ContextImpl {
     // and not in the promise context. To understand the semantic difference, make this function async and run the
     // UnawaitedSideEffectShouldFailSubsequentContextCall test.
     run(nameOrAction, actionSecondParameter, options) {
-        this.checkState("run");
         const { name, action } = unpack(nameOrAction, actionSecondParameter);
-        this.executingRun = true;
         const serde = options?.serde ?? (0, rpc_js_1.defaultSerde)();
-        const executeRun = async () => {
-            // in replay mode, we directly return the value from the log
-            if (this.stateMachine.nextEntryWillBeReplayed()) {
-                const emptyMsg = new protocol_pb_js_1.RunEntryMessage({});
-                return this.stateMachine
-                    .handleUserCodeMessage(protocol_js_1.SIDE_EFFECT_ENTRY_MESSAGE_TYPE, emptyMsg)
-                    .transform((result) => {
-                    if (!result || result instanceof protocol_pb_js_1.Empty) {
-                        return undefined;
-                    }
-                    return serde.deserialize(result);
-                });
+        try {
+            const runEnterResult = this.coreVm.sys_run_enter(name || "");
+            // Check if the run was already executed
+            if (typeof runEnterResult === "object" &&
+                "ExecutedWithSuccess" in runEnterResult) {
+                return Promise.resolve(serde.deserialize(runEnterResult.ExecutedWithSuccess));
+            }
+            else if (typeof runEnterResult === "object" &&
+                "ExecutedWithFailure" in runEnterResult) {
+                return Promise.reject(new errors_js_1.TerminalError(runEnterResult.ExecutedWithFailure.message, {
+                    errorCode: runEnterResult.ExecutedWithFailure.code,
+                }));
             }
-            let sideEffectResult;
+        }
+        catch (e) {
+            this.handleInvocationEndError(e);
+            return pendingPromise();
+        }
+        // We wrap the rest of the execution in this closure to create a future
+        const doRun = async () => {
+            const startTime = Date.now();
+            let res;
+            let err;
             try {
-                sideEffectResult = await action();
+                res = await action();
             }
             catch (e) {
-                if (!(e instanceof errors_js_1.TerminalError)) {
-                    ///non terminal errors are retirable.
-                    // we do not commit the error itself into the journal, but rather let restate know about this
-                    // so that restate can retry this invocation later.
-                    // Before we can propagate this error to the user, we must let the state machine know that this attempt
-                    // is finished with an error, and it should not append anything else to the journal from now on.
-                    const error = (0, errors_js_1.ensureError)(e);
-                    const additionalContext = {
-                        relatedEntryName: name,
-                        relatedEntryType: protocol_js_1.SIDE_EFFECT_ENTRY_MESSAGE_TYPE,
-                    };
-                    await this.stateMachine.sendErrorAndFinish(error, additionalContext);
-                    throw e;
+                err = (0, errors_js_1.ensureError)(e);
+            }
+            const attemptDuration = Date.now() - startTime;
+            // Record the result/failure, get back the handle for the ack.
+            let handle;
+            try {
+                if (err !== undefined) {
+                    if (err instanceof errors_js_1.TerminalError) {
+                        // Record failure, go ahead
+                        handle = this.coreVm.sys_run_exit_failure({
+                            code: err.code,
+                            message: err.message,
+                        });
+                    }
+                    else {
+                        if (options?.retryIntervalFactor === undefined &&
+                            options?.initialRetryIntervalMillis === undefined &&
+                            options?.maxRetryAttempts === undefined &&
+                            options?.maxRetryDurationMillis === undefined &&
+                            options?.maxRetryIntervalMillis === undefined) {
+                            // If no retry option was set, simply throw the error.
+                            // This will lead to the invoker applying its retry, without the SDK overriding it.
+                            throw err;
+                        }
+                        handle = this.coreVm.sys_run_exit_failure_transient(err.message, err.cause?.toString(), BigInt(attemptDuration), {
+                            factor: options?.retryIntervalFactor || 2.0,
+                            initial_interval: options?.initialRetryIntervalMillis || 50,
+                            max_attempts: options?.maxRetryAttempts,
+                            max_duration: options?.maxRetryDurationMillis,
+                            max_interval: options?.maxRetryIntervalMillis || 10 * 1000,
+                        });
+                    }
+                }
+                else {
+                    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+                    // @ts-expect-error
+                    handle = this.coreVm.sys_run_exit_success(serde.serialize(res));
                 }
-                // we commit a terminal error from the side effect to the journal, and re-throw it into
-                // the function. that way, any catching by the user and reacting to it will be
-                // deterministic on replay
-                const error = (0, errors_js_1.ensureError)(e);
-                const failure = (0, errors_js_1.errorToFailure)(error);
-                const sideEffectMsg = new protocol_pb_js_1.RunEntryMessage({
-                    name,
-                    result: { case: "failure", value: failure },
-                });
-                // this may throw an error from the SDK/runtime/connection side, in case the
-                // failure message cannot be committed to the journal. That error would then
-                // be returned from this function (replace the original error)
-                // that is acceptable, because in such a situation (failure to append to journal),
-                // the state machine closes anyways and no further operations will succeed and the
-                // the execution aborts
-                await this.stateMachine.handleUserCodeMessage(protocol_js_1.SIDE_EFFECT_ENTRY_MESSAGE_TYPE, sideEffectMsg, false, true);
-                throw e;
             }
-            // we have this code outside the above try/catch block, to ensure that any error arising
-            // from here is not incorrectly attributed to the side-effect
-            const sideEffectMsg = sideEffectResult !== undefined
-                ? new protocol_pb_js_1.RunEntryMessage({
-                    name,
-                    result: {
-                        case: "value",
-                        value: serde.serialize(sideEffectResult),
-                    },
-                })
-                : new protocol_pb_js_1.RunEntryMessage({
-                    name,
-                });
-            // if an error arises from committing the side effect result, then this error will
-            // be thrown here (reject the returned promise) and the function will see that error,
-            // even if the side-effect function completed correctly
-            // that is acceptable, because in such a situation (failure to append to journal),
-            // the state machine closes anyways and reports an execution failure, meaning no further
-            // operations will succeed and the the execution will be retried.
-            // If the side-effect result did in fact not make it to the journal, then the side-effect
-            // re-executes, and if it made it to the journal after all (error happend inly during
-            // ack-back), then retries will use the journaled result.
-            // So all good in any case, due to the beauty of "the runtime log is the ground thruth" approach.
-            await this.stateMachine.handleUserCodeMessage(protocol_js_1.SIDE_EFFECT_ENTRY_MESSAGE_TYPE, sideEffectMsg, false, true);
-            return sideEffectResult;
+            catch (e) {
+                this.handleInvocationEndError(e);
+                return pendingPromise();
+            }
+            // Got the handle, wait for the result now (which we get once we get the ack)
+            return await this.pollAsyncResult(handle, (asyncResultValue) => {
+                if (typeof asyncResultValue === "object" &&
+                    "Success" in asyncResultValue) {
+                    return serde.deserialize(asyncResultValue.Success);
+                }
+                else if (typeof asyncResultValue === "object" &&
+                    "Failure" in asyncResultValue) {
+                    throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                        errorCode: asyncResultValue.Failure.code,
+                    });
+                }
+                throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
+            });
         };
-        return executeRun().finally(() => {
-            this.executingRun = false;
-        });
+        return doRun();
     }
     sleep(millis) {
-        this.checkState("sleep");
-        return this.markCombineablePromise(this.sleepInternal(millis));
-    }
-    sleepInternal(millis) {
-        return this.stateMachine.handleUserCodeMessage(protocol_js_1.SLEEP_ENTRY_MESSAGE_TYPE, new protocol_pb_js_1.SleepEntryMessage({
-            wakeUpTime: protobuf_1.protoInt64.parse(Date.now() + millis),
-        }));
-    }
-    // -- Awakeables
-    awakeable(serde) {
-        this.checkState("awakeable");
-        const msg = new protocol_pb_js_1.AwakeableEntryMessage();
-        const promise = this.stateMachine
-            .handleUserCodeMessage(protocol_js_1.AWAKEABLE_ENTRY_MESSAGE_TYPE, msg)
-            .transform((result) => {
-            if (!(result instanceof Uint8Array)) {
-                // This should either be a filled buffer or an empty buffer but never anything else.
-                throw errors_js_1.RetryableError.internal("Awakeable was not resolved with a buffer payload");
-            }
-            if (!serde) {
-                return (0, rpc_js_1.defaultSerde)().deserialize(result);
-            }
-            if (result.length == 0) {
+        return this.processCompletableEntry((vm) => vm.sys_sleep(BigInt(millis)), (asyncResultValue) => {
+            if (asyncResultValue === "Empty") {
+                // Empty
                 return undefined;
             }
-            return serde.deserialize(result);
+            else if ("Failure" in asyncResultValue) {
+                throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                    errorCode: asyncResultValue.Failure.code,
+                });
+            }
+            throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
         });
-        // This needs to be done after handling the message in the state machine
-        // otherwise the index is not yet incremented.
-        const encodedEntryIndex = node_buffer_1.Buffer.alloc(4 /* Size of u32 */);
-        encodedEntryIndex.writeUInt32BE(this.stateMachine.getUserCodeJournalIndex());
+    }
+    // -- Awakeables
+    awakeable(serde) {
+        let awakeable;
+        try {
+            awakeable = this.coreVm.sys_awakeable();
+        }
+        catch (e) {
+            this.handleInvocationEndError(e);
+            return {
+                id: "invalid",
+                promise: new LazyContextPromise(0, this, () => pendingPromise()),
+            };
+        }
         return {
-            id: protocol_js_1.AWAKEABLE_IDENTIFIER_PREFIX +
-                node_buffer_1.Buffer.concat([this.request().id, encodedEntryIndex]).toString("base64url"),
-            promise: this.markCombineablePromise(promise),
+            id: awakeable.id,
+            promise: new LazyContextPromise(awakeable.handle, this, () => this.pollAsyncResult(awakeable.handle, (asyncResultValue) => {
+                if (typeof asyncResultValue === "object" &&
+                    "Success" in asyncResultValue) {
+                    if (!serde) {
+                        return (0, rpc_js_1.defaultSerde)().deserialize(asyncResultValue.Success);
+                    }
+                    if (asyncResultValue.Success.length === 0) {
+                        return undefined;
+                    }
+                    return serde.deserialize(asyncResultValue.Success);
+                }
+                else if (typeof asyncResultValue === "object" &&
+                    "Failure" in asyncResultValue) {
+                    throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                        errorCode: asyncResultValue.Failure.code,
+                    });
+                }
+                throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
+            })),
         };
     }
     resolveAwakeable(id, payload, serde) {
-        // We coerce undefined to null as null can be stringified by JSON.stringify
-        let value;
-        if (serde) {
-            value =
-                payload == undefined ? new Uint8Array() : serde.serialize(payload);
-        }
-        else {
-            value =
-                payload != undefined
-                    ? (0, rpc_js_1.defaultSerde)().serialize(payload)
-                    : (0, rpc_js_1.defaultSerde)().serialize(null);
-        }
-        this.checkState("resolveAwakeable");
-        this.completeAwakeable(id, {
-            result: {
-                case: "value",
-                value,
-            },
+        this.processNonCompletableEntry((vm) => {
+            // We coerce undefined to null as null can be stringified by JSON.stringify
+            let value;
+            if (serde) {
+                value =
+                    payload === undefined ? new Uint8Array() : serde.serialize(payload);
+            }
+            else {
+                value =
+                    payload !== undefined
+                        ? (0, rpc_js_1.defaultSerde)().serialize(payload)
+                        : (0, rpc_js_1.defaultSerde)().serialize(null);
+            }
+            vm.sys_complete_awakeable_success(id, value);
         });
     }
     rejectAwakeable(id, reason) {
-        this.checkState("rejectAwakeable");
-        this.completeAwakeable(id, {
-            result: {
-                case: "failure",
-                value: { code: errors_js_1.UNKNOWN_ERROR_CODE, message: reason },
-            },
+        this.processNonCompletableEntry((vm) => {
+            vm.sys_complete_awakeable_failure(id, {
+                code: errors_js_1.UNKNOWN_ERROR_CODE,
+                message: reason,
+            });
         });
     }
-    completeAwakeable(id, base) {
-        base.id = id;
-        this.stateMachine
-            .handleUserCodeMessage(protocol_js_1.COMPLETE_AWAKEABLE_ENTRY_MESSAGE_TYPE, new protocol_pb_js_1.CompleteAwakeableEntryMessage(base))
-            .catch((e) => this.stateMachine.handleDanglingPromiseError(e));
+    promise(name, serde) {
+        return new DurablePromiseImpl(this, name, serde);
     }
     // Used by static methods of CombineablePromise
-    static createCombinator(combinatorConstructor, promises) {
+    static createCombinator(combinatorType, promises) {
+        // Extract context from first promise
         const self = extractContext(promises[0]);
         if (!self) {
-            throw errors_js_1.RetryableError.internal("Not a combinable promise");
+            throw new Error("Not a combinable promise");
         }
-        const outPromises = [];
+        // Collect first the promises downcasted to the internal promise type
+        const castedPromises = [];
         for (const promise of promises) {
             if (extractContext(promise) !== self) {
-                throw errors_js_1.RetryableError.internal("You're mixing up CombineablePromises from different RestateContext. This is not supported.");
+                self.handleInvocationEndError(new Error("You're mixing up CombineablePromises from different RestateContext. This is not supported."));
+                return pendingPromise();
+            }
+            castedPromises.push(promise);
+        }
+        const handles = new Uint32Array(castedPromises.map((p) => p.asyncResultHandle));
+        // From now on, lazily executes on await
+        return new LazyPromise(async () => {
+            let combinatorResultHandle;
+            try {
+                // Take output
+                const nextOutput1 = self.coreVm.take_output();
+                if (nextOutput1 instanceof Uint8Array) {
+                    await self.outputWriter.write(nextOutput1);
+                }
+                for (;;) {
+                    switch (combinatorType) {
+                        case "All":
+                            combinatorResultHandle =
+                                self.coreVm.sys_try_complete_all_combinator(handles);
+                            break;
+                        case "Any":
+                            combinatorResultHandle =
+                                self.coreVm.sys_try_complete_any_combinator(handles);
+                            break;
+                        case "AllSettled":
+                            combinatorResultHandle =
+                                self.coreVm.sys_try_complete_all_settled_combinator(handles);
+                            break;
+                        case "Race":
+                        case "OrTimeout":
+                            combinatorResultHandle =
+                                self.coreVm.sys_try_complete_race_combinator(handles);
+                            break;
+                    }
+                    // We got a result, we're done in this loop
+                    if (combinatorResultHandle !== undefined) {
+                        break;
+                    }
+                    // No result yet, await the next read
+                    await self.awaitNextRead();
+                }
+                // We got a result, we need to take_output to write the combinator entry, then we need to poll the result
+                const nextOutput = self.coreVm.take_output();
+                if (nextOutput instanceof Uint8Array) {
+                    await self.outputWriter.write(nextOutput);
+                }
+            }
+            catch (e) {
+                if (e instanceof errors_js_1.TerminalError) {
+                    // All good, this is a recorded failure
+                    throw e;
+                }
+                // Not good, this is a retryable error.
+                self.handleInvocationEndError(e);
+                return await pendingPromise();
             }
-            const index = promise
-                .journalIndex;
-            outPromises.push({
-                id: (0, promise_combinator_tracker_js_1.newJournalEntryPromiseId)(index),
-                promise: promise,
+            const handlesResult = await self.pollAsyncResult(combinatorResultHandle, (asyncResultValue) => {
+                if (typeof asyncResultValue === "object" &&
+                    "CombinatorResult" in asyncResultValue) {
+                    return asyncResultValue.CombinatorResult;
+                }
+                throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
             });
+            const promisesMap = new Map(castedPromises.map((p) => [p.asyncResultHandle, p]));
+            // Now all we need to do is to construct the final output based on the handles,
+            // this depends on combinators themselves.
+            switch (combinatorType) {
+                case "All":
+                    return this.extractAllCombinatorResult(handlesResult, promisesMap);
+                case "Any":
+                    return this.extractAnyCombinatorResult(handlesResult, promisesMap);
+                case "AllSettled":
+                    return this.extractAllSettledCombinatorResult(handlesResult, promisesMap);
+                case "Race":
+                    // Just one promise succeeded
+                    return promisesMap.get(handlesResult[0]);
+                case "OrTimeout":
+                    // The sleep promise is always the second one in the list.
+                    if (handlesResult[0] === castedPromises[1].asyncResultHandle) {
+                        return Promise.reject(new errors_js_1.TimeoutError());
+                    }
+                    else {
+                        return promisesMap.get(handlesResult[0]);
+                    }
+            }
+        });
+    }
+    static async extractAllCombinatorResult(handlesResult, promisesMap) {
+        // The result can either all values, or one error
+        const resultValues = [];
+        for (const handle of handlesResult) {
+            try {
+                resultValues.push(await promisesMap.get(handle));
+            }
+            catch (e) {
+                return Promise.reject(e);
+            }
+        }
+        return Promise.resolve(resultValues);
+    }
+    static async extractAnyCombinatorResult(handlesResult, promisesMap) {
+        // The result can either be one value, or a list of errors
+        const resultFailures = [];
+        for (const handle of handlesResult) {
+            try {
+                return Promise.resolve(await promisesMap.get(handle));
+            }
+            catch (e) {
+                resultFailures.push(e);
+            }
+        }
+        // Giving back the cause here is completely fine, because all these errors in Aggregate error are Terminal errors!
+        return Promise.reject(new errors_js_1.TerminalError("All input promises failed", {
+            cause: new AggregateError(resultFailures),
+        }));
+    }
+    static async extractAllSettledCombinatorResult(handlesResult, promisesMap) {
+        const resultValues = [];
+        for (const handle of handlesResult) {
+            try {
+                resultValues.push(await promisesMap.get(handle));
+            }
+            catch (e) {
+                resultValues.push(e);
+            }
         }
-        return self.stateMachine.createCombinator(combinatorConstructor, outPromises);
+        return Promise.resolve(resultValues);
     }
     // -- Various private methods
-    checkNotExecutingRun(callType) {
-        if (this.executingRun) {
-            throw new errors_js_1.TerminalError(`Invoked a RestateContext method (${callType}) while a run() is still executing.
-          Make sure you await the ctx.run() call before using any other RestateContext method.`, { errorCode: errors_js_1.INTERNAL_ERROR_CODE });
+    processNonCompletableEntry(vmCall) {
+        try {
+            vmCall(this.coreVm);
+        }
+        catch (e) {
+            this.handleInvocationEndError(e);
         }
     }
-    checkState(callType) {
-        this.checkNotExecutingRun(callType);
+    processCompletableEntry(vmCall, transformer) {
+        let handle;
+        try {
+            handle = vmCall(this.coreVm);
+        }
+        catch (e) {
+            this.handleInvocationEndError(e);
+            return new LazyContextPromise(0, this, () => pendingPromise());
+        }
+        return new LazyContextPromise(handle, this, () => this.pollAsyncResult(handle, transformer));
+    }
+    async pollAsyncResult(handle, transformer) {
+        try {
+            // Take output
+            const nextOutput = this.coreVm.take_output();
+            if (nextOutput instanceof Uint8Array) {
+                await this.outputWriter.write(nextOutput);
+            }
+            // Now loop waiting for the async result
+            let asyncResult = this.coreVm.take_async_result(handle);
+            while (asyncResult === "NotReady") {
+                await this.awaitNextRead();
+                // Using notify_await_point immediately before take_async_result
+                // makes sure the state machine will try to suspend only now,
+                // in case there aren't other concurrent tasks trying to poll this async result.
+                this.coreVm.notify_await_point(handle);
+                asyncResult = this.coreVm.take_async_result(handle);
+            }
+            return transformer(asyncResult);
+        }
+        catch (e) {
+            if (e instanceof errors_js_1.TerminalError) {
+                // All good, this is a recorded failure
+                throw e;
+            }
+            // Not good, this is a retryable error.
+            this.handleInvocationEndError(e);
+            return await pendingPromise();
+        }
     }
-    markCombineablePromise(p) {
-        const journalIndex = this.stateMachine.getUserCodeJournalIndex();
-        const orTimeout = (millis) => {
-            const sleepPromise = this.sleepInternal(millis).transform(() => {
-                throw new errors_js_1.TimeoutError();
+    // This function triggers a read on the input reader,
+    // and will notify the caller that a read was executed
+    // and the result was piped in the state machine.
+    awaitNextRead() {
+        if (this.currentRead === undefined) {
+            // Register a new read
+            this.currentRead = this.readNext().finally(() => {
+                this.currentRead = undefined;
             });
-            const sleepPromiseIndex = this.stateMachine.getUserCodeJournalIndex();
-            return this.stateMachine.createCombinator(Promise.race.bind(Promise), [
-                {
-                    id: (0, promise_combinator_tracker_js_1.newJournalEntryPromiseId)(journalIndex),
-                    promise: p,
-                },
-                {
-                    id: (0, promise_combinator_tracker_js_1.newJournalEntryPromiseId)(sleepPromiseIndex),
-                    promise: sleepPromise,
-                },
-            ]);
-        };
-        defineProperty(p, RESTATE_CTX_SYMBOL, this);
-        defineProperty(p, "journalIndex", journalIndex);
-        defineProperty(p, "orTimeout", orTimeout.bind(this));
-        return p;
+        }
+        // eslint-disable-next-line @typescript-eslint/no-misused-promises
+        return new Promise((resolve) => this.currentRead?.finally(resolve));
+    }
+    async readNext() {
+        // Take input, and notify it to the vm
+        let nextValue;
+        try {
+            nextValue = await this.inputReader.read();
+        }
+        catch (e) {
+            this.handleInvocationEndError(e);
+            return pendingPromise();
+        }
+        if (nextValue.value !== undefined) {
+            this.coreVm.notify_input(nextValue.value);
+        }
+        if (nextValue.done) {
+            this.coreVm.notify_input_closed();
+        }
+    }
+    handleInvocationEndError(e) {
+        const error = (0, errors_js_1.ensureError)(e);
+        if (!(error instanceof errors_js_1.RestateError) ||
+            error.code !== errors_js_1.SUSPENDED_ERROR_CODE) {
+            this.console.warn("Function completed with an error.\n", error);
+        }
+        this.coreVm.notify_error(error.message, error.stack);
+        // From now on, no progress will be made.
+        this.invocationEndPromise.resolve();
     }
 }
 exports.ContextImpl = ContextImpl;
-// wraps defineProperty such that it informs tsc of the correct type of its output
-function defineProperty(obj, prop, value) {
-    Object.defineProperty(obj, prop, { value });
-}
 function unpack(a, b) {
-    if (typeof a == "string") {
+    if (typeof a === "string") {
         if (typeof b !== "function") {
             throw new TypeError("");
         }
@@ -452,56 +613,108 @@ class DurablePromiseImpl {
     }
     [Symbol.toStringTag] = "DurablePromise";
     get() {
-        const msg = new protocol_pb_js_1.GetPromiseEntryMessage({
-            key: this.name,
-        });
-        return this.ctx.markCombineablePromise(this.ctx.stateMachine
-            .handleUserCodeMessage(protocol_js_1.GET_PROMISE_MESSAGE_TYPE, msg)
-            .transform((v) => {
-            if (!v) {
-                return undefined;
+        return this.ctx.processCompletableEntry((vm) => vm.sys_get_promise(this.name), (asyncResultValue) => {
+            if (typeof asyncResultValue === "object" &&
+                "Success" in asyncResultValue) {
+                return this.serde.deserialize(asyncResultValue.Success);
             }
-            if (v instanceof protocol_pb_js_1.Empty) {
-                return undefined;
+            else if (typeof asyncResultValue === "object" &&
+                "Failure" in asyncResultValue) {
+                throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                    errorCode: asyncResultValue.Failure.code,
+                });
             }
-            return this.serde.deserialize(v);
-        }));
+            throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
+        });
     }
     peek() {
-        const msg = new protocol_pb_js_1.PeekPromiseEntryMessage({
-            key: this.name,
-        });
-        return this.ctx.stateMachine
-            .handleUserCodeMessage(protocol_js_1.PEEK_PROMISE_MESSAGE_TYPE, msg)
-            .transform((v) => {
-            if (!v || v instanceof protocol_pb_js_1.Empty) {
+        return this.ctx.processCompletableEntry((vm) => vm.sys_peek_promise(this.name), (asyncResultValue) => {
+            if (asyncResultValue === "Empty") {
                 return undefined;
             }
-            return this.serde.deserialize(v);
+            else if (typeof asyncResultValue === "object" &&
+                "Success" in asyncResultValue) {
+                return this.serde.deserialize(asyncResultValue.Success);
+            }
+            else if (typeof asyncResultValue === "object" &&
+                "Failure" in asyncResultValue) {
+                throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                    errorCode: asyncResultValue.Failure.code,
+                });
+            }
+            throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
         });
     }
     resolve(value) {
-        const buffer = value != undefined ? this.serde.serialize(value) : new Uint8Array();
-        const msg = new protocol_pb_js_1.CompletePromiseEntryMessage({
-            key: this.name,
-            completion: {
-                case: "completionValue",
-                value: buffer,
-            },
+        return this.ctx.processCompletableEntry((vm) => vm.sys_complete_promise_success(this.name, this.serde.serialize(value)), (asyncResultValue) => {
+            if (asyncResultValue === "Empty") {
+                return undefined;
+            }
+            else if (typeof asyncResultValue === "object" &&
+                "Failure" in asyncResultValue) {
+                throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                    errorCode: asyncResultValue.Failure.code,
+                });
+            }
+            throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
         });
-        return this.ctx.stateMachine.handleUserCodeMessage(protocol_js_1.COMPLETE_PROMISE_MESSAGE_TYPE, msg);
     }
     reject(errorMsg) {
-        const msg = new protocol_pb_js_1.CompletePromiseEntryMessage({
-            key: this.name,
-            completion: {
-                case: "completionFailure",
-                value: {
-                    message: errorMsg,
-                },
-            },
+        return this.ctx.processCompletableEntry((vm) => vm.sys_complete_promise_failure(this.name, {
+            code: errors_js_1.INTERNAL_ERROR_CODE,
+            message: errorMsg,
+        }), (asyncResultValue) => {
+            if (asyncResultValue === "Empty") {
+                return undefined;
+            }
+            else if (typeof asyncResultValue === "object" &&
+                "Failure" in asyncResultValue) {
+                throw new errors_js_1.TerminalError(asyncResultValue.Failure.message, {
+                    errorCode: asyncResultValue.Failure.code,
+                });
+            }
+            throw new Error(`Unexpected variant in async result: ${JSON.stringify(asyncResultValue)}`);
         });
-        return this.ctx.stateMachine.handleUserCodeMessage(protocol_js_1.COMPLETE_PROMISE_MESSAGE_TYPE, msg);
     }
 }
+class LazyPromise {
+    executor;
+    _promise;
+    constructor(executor) {
+        this.executor = executor;
+    }
+    then(onfulfilled, onrejected) {
+        this._promise = this._promise || this.executor();
+        return this._promise.then(onfulfilled, onrejected);
+    }
+    catch(onrejected) {
+        this._promise = this._promise || this.executor();
+        return this._promise.catch(onrejected);
+    }
+    finally(onfinally) {
+        this._promise = this._promise || this.executor();
+        return this._promise.finally(onfinally);
+    }
+    [Symbol.toStringTag] = "LazyPromise";
+}
+class LazyContextPromise extends LazyPromise {
+    asyncResultHandle;
+    [RESTATE_CTX_SYMBOL];
+    constructor(asyncResultHandle, ctx, executor) {
+        super(executor);
+        this.asyncResultHandle = asyncResultHandle;
+        this[RESTATE_CTX_SYMBOL] = ctx;
+    }
+    orTimeout(millis) {
+        return ContextImpl.createCombinator("OrTimeout", [
+            this,
+            this[RESTATE_CTX_SYMBOL].sleep(millis),
+        ]);
+    }
+}
+// A promise that is never completed
+function pendingPromise() {
+    // eslint-disable-next-line @typescript-eslint/no-empty-function
+    return new Promise(() => { });
+}
 //# sourceMappingURL=context_impl.js.map