Convert rpc handler responses into a JSON-safe 'any' type (#203)

We pass directly to the google.protobuf.Value conversion function, which is fine, except that this is not as clever as JSON.stringify, and will treat objects like Date as empty objects. To match stringify, we should respect the toJSON() function which may be present on the value to be converted.

https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/JSON/stringify#description

Author: jackkleeman

src/server/base_restate_server.ts

@@ -49,6 +49,7 @@ import { RpcContextImpl } from "../restate_context_impl";
 import { verifyAssumptions } from "../utils/assumpsions";
 import { TerminalError } from "../public_api";
 import { isEventHandler } from "../types/router";
+import { jsonSafeAny } from "../utils/utils";
 
 export interface ServiceOpts {
   descriptor: ProtoMetadata;
@@ -175,7 +176,9 @@ export abstract class BaseRestateServer {
 
     const decoder = RpcRequest.decode;
     const encoder = (message: RpcResponse) =>
-      RpcResponse.encode(message).finish();
+      RpcResponse.encode({
+        response: jsonSafeAny("", message.response),
+      }).finish();
 
     const method = new GrpcServiceMethod<RpcRequest, RpcResponse>(
       route,

src/utils/utils.ts

@@ -69,6 +69,26 @@ export function jsonDeserialize<T>(json: string): T {
   ) as T;
 }
 
+// When using google.protobuf.Value in RPC handler responses, we want to roughly match the behaviour of JSON.stringify
+// for example in converting Date objects to a UTC string
+export function jsonSafeAny(key: string, value: any): any {
+  if (typeof value.toJSON == "function") {
+    return value.toJSON(key) as any;
+  } else if (globalThis.Array.isArray(value)) {
+    // in place replace
+    value.forEach((_, i) => (value[i] = jsonSafeAny(i.toString(), value[i])));
+    return value;
+  } else if (typeof value === "object") {
+    Object.keys(value).forEach((key) => {
+      value[key] = jsonSafeAny(key, value[key]);
+    });
+    return value;
+  } else {
+    // primitive that doesn't have a toJSON method, with no children
+    return value;
+  }
+}
+
 export function printMessageAsJson(obj: any): string {
   const newObj = { ...(obj as Record<string, unknown>) };
   for (const [key, value] of Object.entries(newObj)) {

test/utils.test.ts

@@ -12,6 +12,7 @@
 import { describe, expect } from "@jest/globals";
 import {
   jsonDeserialize,
+  jsonSafeAny,
   jsonSerialize,
   printMessageAsJson,
 } from "../src/utils/utils";
@@ -127,3 +128,73 @@ describe("rand", () => {
     expect(actual).toStrictEqual(expected);
   });
 });
+
+describe("jsonSafeAny", () => {
+  it("handles dates", () => {
+    expect(jsonSafeAny("", new Date(1701878170682))).toStrictEqual(
+      "2023-12-06T15:56:10.682Z"
+    );
+    expect(jsonSafeAny("", { date: new Date(1701878170682) })).toStrictEqual({
+      date: "2023-12-06T15:56:10.682Z",
+    });
+    expect(
+      jsonSafeAny("", {
+        dates: [new Date(1701878170682), new Date(1701878170683)],
+      })
+    ).toStrictEqual({
+      dates: ["2023-12-06T15:56:10.682Z", "2023-12-06T15:56:10.683Z"],
+    });
+  });
+  it("handles urls", () => {
+    expect(jsonSafeAny("", new URL("https://restate.dev"))).toStrictEqual(
+      "https://restate.dev/"
+    );
+  });
+  it("handles patched BigInts", () => {
+    // by default should do nothing
+    expect(jsonSafeAny("", BigInt("9007199254740991"))).toStrictEqual(
+      BigInt("9007199254740991")
+    );
+    (BigInt.prototype as any).toJSON = function () {
+      return this.toString();
+    };
+    expect(jsonSafeAny("", BigInt("9007199254740991"))).toStrictEqual(
+      "9007199254740991"
+    );
+    delete (BigInt.prototype as any).toJSON;
+  });
+  it("handles custom types", () => {
+    const numberType = {
+      toJSON(key: string): number {
+        return 1;
+      },
+    };
+    const stringType = {
+      toJSON(key: string): string {
+        return "foo";
+      },
+    };
+    expect(jsonSafeAny("", numberType)).toStrictEqual(1);
+    expect(jsonSafeAny("", stringType)).toStrictEqual("foo");
+  });
+  it("provides the correct key", () => {
+    const keys: string[] = [];
+    const typ = {
+      toJSON(key: string): string {
+        keys.push(key);
+        return "";
+      },
+    };
+    expect(jsonSafeAny("", typ)).toStrictEqual("");
+    expect(jsonSafeAny("", { key: typ })).toStrictEqual({ key: "" });
+    expect(jsonSafeAny("", { key: [typ] })).toStrictEqual({ key: [""] });
+    expect(jsonSafeAny("", { key: [0, typ] })).toStrictEqual({ key: [0, ""] });
+    expect(jsonSafeAny("", { key: [0, { key2: typ }] })).toStrictEqual({
+      key: [0, { key2: "" }],
+    });
+    expect(jsonSafeAny("", { key: [0, { key2: [typ] }] })).toStrictEqual({
+      key: [0, { key2: [""] }],
+    });
+    expect(keys).toStrictEqual(["", "key", "0", "1", "key2", "0"]);
+  });
+});