Add support for headers for service-to-service calls (#467)

Author: jackkleeman

package-lock.json

@@ -46,7 +46,9 @@
     "prettier": "^2.8.4",
     "release-it": "^17.3.0",
     "typescript": "^5.4.5",
-    "vitest": "^1.6.0"
+    "vitest": "^1.6.0",
+    "wasm-pack": "^0.13.1",
+    "wasm-pack-inline": "^0.1.2"
   },
   "engines": {
     "node": ">= 18.13"

package.json

@@ -228,6 +228,7 @@ export type GenericCall<REQ, RES> = {
   method: string;
   parameter: REQ;
   key?: string;
+  headers?: Record<string, string>;
   inputSerde?: Serde<REQ>;
   outputSerde?: Serde<RES>;
 };
@@ -242,6 +243,7 @@ export type GenericSend<REQ> = {
   method: string;
   parameter: REQ;
   key?: string;
+  headers?: Record<string, string>;
   inputSerde?: Serde<REQ>;
   delay?: number;
 };

packages/restate-sdk/src/context.ts

@@ -60,6 +60,7 @@ import type {
 } from "node:stream/web";
 import type { ReadableStreamReadResult } from "stream/web";
 import type { CompletablePromise } from "./utils/completable_promise.js";
+import { WasmHeader } from "./endpoint/handlers/vm/sdk_shared_core_wasm_bindings.js";
 
 export type InternalCombineablePromise<T> = CombineablePromise<T> & {
   asyncResultHandle: number;
@@ -213,7 +214,17 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
     return this.processCompletableEntry(
       (vm) => {
         const parameter = requestSerde.serialize(call.parameter);
-        return vm.sys_call(call.service, call.method, parameter, call.key);
+        return vm.sys_call(
+          call.service,
+          call.method,
+          parameter,
+          call.key,
+          call.headers
+            ? Object.entries(call.headers).map(
+                ([key, value]) => new WasmHeader(key, value)
+              )
+            : []
+        );
       },
       (asyncResultValue) => {
         if (
@@ -248,7 +259,18 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
         delay = BigInt(send.delay);
       }
 
-      vm.sys_send(send.service, send.method, parameter, send.key, delay);
+      void vm.sys_send(
+        send.service,
+        send.method,
+        parameter,
+        send.key,
+        send.headers
+          ? Object.entries(send.headers).map(
+              ([key, value]) => new WasmHeader(key, value)
+            )
+          : [],
+        delay
+      );
     });
   }
 
@@ -736,6 +758,7 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
         | { Success: Uint8Array }
         | { Failure: vm.WasmFailure }
         | { StateKeys: string[] }
+        | { InvocationId: string }
         | { CombinatorResult: number[] }
     ) => T
   ): LazyContextPromise<T> {
@@ -759,6 +782,7 @@ export class ContextImpl implements ObjectContext, WorkflowContext {
         | { Success: Uint8Array }
         | { Failure: vm.WasmFailure }
         | { StateKeys: string[] }
+        | { InvocationId: string }
         | { CombinatorResult: number[] }
     ) => T
   ): Promise<T> {

packages/restate-sdk/src/context_impl.ts

@@ -22,6 +22,8 @@ export interface WasmFailure {
   message: string;
 }
 
+export type WasmSendHandle = number;
+
 export interface WasmExponentialRetryConfig {
   initial_interval: number | undefined;
   factor: number;
@@ -41,6 +43,7 @@ export type WasmAsyncResultValue =
   | { Success: Uint8Array }
   | { Failure: WasmFailure }
   | { StateKeys: string[] }
+  | { InvocationId: string }
   | { CombinatorResult: WasmAsyncResultHandle[] };
 
 export type WasmRunEnterResult =
@@ -181,29 +184,34 @@ export class WasmVM {
    * @param {string} service
    * @param {string} handler
    * @param {Uint8Array} buffer
-   * @param {string | undefined} [key]
+   * @param {string | undefined} key
+   * @param {(WasmHeader)[]} headers
    * @returns {number}
    */
   sys_call(
     service: string,
     handler: string,
     buffer: Uint8Array,
-    key?: string
+    key: string | undefined,
+    headers: WasmHeader[]
   ): number;
   /**
    * @param {string} service
    * @param {string} handler
    * @param {Uint8Array} buffer
-   * @param {string | undefined} [key]
+   * @param {string | undefined} key
+   * @param {(WasmHeader)[]} headers
    * @param {bigint | undefined} [delay]
+   * @returns {WasmSendHandle}
    */
   sys_send(
     service: string,
     handler: string,
     buffer: Uint8Array,
-    key?: string,
+    key: string | undefined,
+    headers: WasmHeader[],
     delay?: bigint
-  ): void;
+  ): WasmSendHandle;
   /**
    * @returns {WasmAwakeable}
    */

packages/restate-sdk/src/endpoint/handlers/vm/sdk_shared_core_wasm_bindings.d.ts

@@ -12,7 +12,7 @@ default = ["console_error_panic_hook"]
 
 [dependencies]
 wasm-bindgen = "0.2.84"
-restate-sdk-shared-core = { git = "https://github.com/restatedev/sdk-shared-core.git", rev = "771df58ed7fa8976f18a865496e1f3ee0dae73f4", features = ["request_identity"] }
+restate-sdk-shared-core = { git = "https://github.com/restatedev/sdk-shared-core.git", rev = "3c342c43e91514aaa20c7ee42295162842b223ab", features = ["request_identity"] }
 console_error_panic_hook = { version = "0.1.7", optional = true }
 serde = { version = "1.0.210", features = ["derive"] }
 serde-wasm-bindgen = "0.6.5"

packages/restate-sdk/src/endpoint/handlers/vm/sdk_shared_core_wasm_bindings.js

@@ -1,4 +1,4 @@
 To build:
 
-* wasm-pack build --target web
+* npx wasm-pack build --target web
 * npx wasm-pack-inline ./pkg --dir ../packages/restate-sdk/src/endpoint/handlers/vm --name sdk_shared_core_wasm_bindings

sdk-shared-core-wasm-bindings/Cargo.lock

@@ -2,8 +2,8 @@ use js_sys::Uint8Array;
 use restate_sdk_shared_core::{
     AsyncResultAccessTracker, AsyncResultCombinator, AsyncResultHandle, AsyncResultState, CoreVM,
     Error, Header, HeaderMap, IdentityVerifier, Input, NonEmptyValue, ResponseHead, RetryPolicy,
-    RunEnterResult, RunExitResult, SuspendedOrVMError, TakeOutputResult, Target, TerminalFailure,
-    VMOptions, Value, VM,
+    RunEnterResult, RunExitResult, SendHandle, SuspendedOrVMError, TakeOutputResult, Target,
+    TerminalFailure, VMOptions, Value, VM,
 };
 use serde::{Deserialize, Serialize};
 use std::convert::{Infallible, Into};
@@ -41,14 +41,9 @@ extern "C" {
     fn vm_log(level: LogLevel, s: &str);
 }
 
+#[derive(Default)]
 pub struct MakeWebConsoleWriter {}
 
-impl Default for MakeWebConsoleWriter {
-    fn default() -> Self {
-        MakeWebConsoleWriter {}
-    }
-}
-
 impl<'a> MakeWriter<'a> for MakeWebConsoleWriter {
     type Writer = ConsoleWriter;
 
@@ -155,6 +150,15 @@ impl From<Header> for WasmHeader {
     }
 }
 
+impl From<WasmHeader> for Header {
+    fn from(h: WasmHeader) -> Self {
+        Header {
+            key: h.key.into(),
+            value: h.value.into(),
+        }
+    }
+}
+
 #[wasm_bindgen(getter_with_clone)]
 pub struct WasmResponseHead {
     #[wasm_bindgen(readonly)]
@@ -203,6 +207,16 @@ impl From<WasmFailure> for JsValue {
     }
 }
 
+#[derive(Tsify, Serialize, Deserialize)]
+#[tsify(into_wasm_abi, from_wasm_abi)]
+pub struct WasmSendHandle(u32);
+
+impl From<SendHandle> for WasmSendHandle {
+    fn from(value: SendHandle) -> Self {
+        WasmSendHandle(value.into())
+    }
+}
+
 #[derive(Tsify, Serialize, Deserialize)]
 #[tsify(into_wasm_abi, from_wasm_abi)]
 pub struct WasmExponentialRetryConfig {
@@ -285,6 +299,7 @@ pub enum WasmAsyncResultValue {
     ),
     Failure(WasmFailure),
     StateKeys(Vec<String>),
+    InvocationId(String),
     CombinatorResult(Vec<WasmAsyncResultHandle>),
 }
 
@@ -377,6 +392,10 @@ impl WasmVM {
                 Some(Value::Success(b)) => WasmAsyncResultValue::Success(b.to_vec().into()),
                 Some(Value::Failure(f)) => WasmAsyncResultValue::Failure(f.into()),
                 Some(Value::StateKeys(keys)) => WasmAsyncResultValue::StateKeys(keys),
+                Some(Value::InvocationId(invocation_id)) => {
+                    WasmAsyncResultValue::InvocationId(invocation_id)
+                }
+
                 Some(Value::CombinatorResult(handles)) => WasmAsyncResultValue::CombinatorResult(
                     handles.into_iter().map(Into::into).collect(),
                 ),
@@ -435,38 +454,46 @@ impl WasmVM {
         handler: String,
         buffer: js_sys::Uint8Array,
         key: Option<String>,
+        headers: Vec<WasmHeader>,
     ) -> Result<WasmAsyncResultHandle, WasmFailure> {
         self.vm
             .sys_call(
                 Target {
                     service,
                     handler,
                     key,
+                    idempotency_key: None,
+                    headers: headers.into_iter().map(Header::from).collect(),
                 },
                 buffer.to_vec().into(),
             )
             .map(Into::into)
             .map_err(Into::into)
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn sys_send(
         &mut self,
         service: String,
         handler: String,
         buffer: Uint8Array,
         key: Option<String>,
+        headers: Vec<WasmHeader>,
         delay: Option<u64>,
-    ) -> Result<(), WasmFailure> {
+    ) -> Result<WasmSendHandle, WasmFailure> {
         self.vm
             .sys_send(
                 Target {
                     service,
                     handler,
                     key,
+                    idempotency_key: None,
+                    headers: headers.into_iter().map(Header::from).collect(),
                 },
                 buffer.to_vec().into(),
                 delay.map(|delay| duration_since_unix_epoch() + Duration::from_millis(delay)),
             )
+            .map(Into::into)
             .map_err(Into::into)
     }