feat: update setMetadata to accept binary headers

Author: chris-olszewski

packages/core-bridge/src/client.rs

@@ -88,13 +88,23 @@ pub fn client_new(
 #[js_function]
 pub fn client_update_headers(
     client: OpaqueInboundHandle<Client>,
-    headers: HashMap<String, String>,
+    headers: HashMap<String, MetadataValue>,
 ) -> BridgeResult<()> {
+    let (ascii_headers, bin_headers) = config::partition_headers(Some(headers));
     client
         .borrow()?
         .core_client
         .get_client()
-        .set_headers(headers)
+        .set_headers(ascii_headers.unwrap_or_default())
+        .map_err(|err| BridgeError::TypeError {
+            message: format!("Invalid metadata key: {err}"),
+            field: None,
+        })?;
+    client
+        .borrow()?
+        .core_client
+        .get_client()
+        .set_binary_headers(bin_headers.unwrap_or_default())
         .map_err(|err| BridgeError::TypeError {
             message: format!("Invalid metadata key: {err}"),
             field: None,
@@ -742,7 +752,7 @@ mod config {
         }
     }
 
-    fn partition_headers(
+    pub(super) fn partition_headers(
         headers: Option<HashMap<String, MetadataValue>>,
     ) -> (
         Option<HashMap<String, String>>,

packages/core-bridge/ts/native.ts

@@ -100,7 +100,7 @@ export interface OtelMetricsExporterOptions {
 
 export declare function newClient(runtime: Runtime, clientOptions: ClientOptions): Promise<Client>;
 
-export declare function clientUpdateHeaders(client: Client, headers: Record<string, string>): void;
+export declare function clientUpdateHeaders(client: Client, headers: Record<string, MetadataValue>): void;
 
 export declare function clientUpdateApiKey(client: Client, apiKey: string): void;
 

packages/test/src/test-native-connection.ts

@@ -439,3 +439,39 @@ test('can power workflow client calls', async (t) => {
     await env.teardown();
   }
 });
+
+test('setMetadata accepts binary headers', async (t) => {
+  const requests = new Array<{ metadata: grpc.Metadata; deadline: grpc.Deadline }>();
+  const server = new grpc.Server();
+  server.addService(workflowServiceProtoDescriptor.temporal.api.workflowservice.v1.WorkflowService.service, {
+    getSystemInfo(
+      call: grpc.ServerUnaryCall<
+        temporal.api.workflowservice.v1.IGetSystemInfoRequest,
+        temporal.api.workflowservice.v1.IGetSystemInfoResponse
+      >,
+      callback: grpc.sendUnaryData<temporal.api.workflowservice.v1.IGetSystemInfoResponse>
+    ) {
+      requests.push({ metadata: call.metadata, deadline: call.getDeadline() });
+      callback(null, {});
+    },
+  });
+
+  const port = await util.promisify(server.bindAsync.bind(server))(
+    'localhost:0',
+    grpc.ServerCredentials.createInsecure()
+  );
+  const connection = await NativeConnection.connect({
+    address: `127.0.0.1:${port}`,
+    metadata: { 'start-ascii': 'a', 'start-bin': Buffer.from([0x00]) },
+  });
+
+  await connection.setMetadata({ 'end-bin': Buffer.from([0x01]) });
+
+  await connection.workflowService.getSystemInfo({});
+  t.is(requests.length, 2);
+  t.deepEqual(requests[1].metadata.get('start-bin'), []);
+  t.deepEqual(requests[1].metadata.get('start-ascii'), []);
+  t.deepEqual(requests[1].metadata.get('end-bin'), [Buffer.from([0x01])]);
+  await connection.close();
+  server.forceShutdown();
+});

packages/worker/src/connection.ts

@@ -126,15 +126,7 @@ export class NativeConnection implements ConnectionLike {
     // TODO: add support for abortSignal
 
     const ctx = this.callContextStorage.getStore() ?? {};
-    const metadata =
-      ctx.metadata != null
-        ? Object.fromEntries(
-            Object.entries(ctx.metadata).map(([k, value]) => [
-              k,
-              typeof value === 'string' ? { type: 'ascii' as const, value } : { type: 'binary' as const, value },
-            ])
-          )
-        : {};
+    const metadata = ctx.metadata != null ? tagMetadata(ctx.metadata) : {};
 
     const req = {
       rpc: method.name,
@@ -271,8 +263,8 @@ export class NativeConnection implements ConnectionLike {
    *
    * Use {@link NativeConnectionOptions.metadata} to set the initial metadata for client creation.
    */
-  async setMetadata(metadata: Record<string, string>): Promise<void> {
-    native.clientUpdateHeaders(this.nativeClient, metadata);
+  async setMetadata(metadata: Metadata): Promise<void> {
+    native.clientUpdateHeaders(this.nativeClient, tagMetadata(metadata));
   }
 
   /**
@@ -348,3 +340,12 @@ function getRelativeTimeout(deadline: grpc.Deadline) {
     return timeout;
   }
 }
+
+function tagMetadata(metadata: Metadata): Record<string, native.MetadataValue> {
+  return Object.fromEntries(
+    Object.entries(metadata).map(([k, value]) => [
+      k,
+      typeof value === 'string' ? { type: 'ascii' as const, value } : { type: 'binary' as const, value },
+    ])
+  );
+}