feat(grpc): generate an example message from the method schema (#613)

apps/yaak-client/components/GrpcEditor.tsx

@@ -1,9 +1,8 @@
 import { linter } from "@codemirror/lint";
 import type { EditorView } from "@codemirror/view";
 import { jsoncLanguage } from "@shopify/lang-jsonc";
-import type { GrpcRequest } from "@yaakapp-internal/models";
-import { FormattedError, InlineCode, VStack } from "@yaakapp-internal/ui";
-import classNames from "classnames";
+import { type GrpcRequest, patchModel } from "@yaakapp-internal/models";
+import { Banner, FormattedError, Icon, InlineCode, VStack } from "@yaakapp-internal/ui";
 import {
   handleRefresh,
   jsonCompletion,
@@ -11,12 +10,20 @@ import {
   stateExtensions,
   updateSchema,
 } from "codemirror-json-schema";
+import type { JSONSchema7 } from "json-schema";
+import type { ReactNode } from "react";
 import { useCallback, useEffect, useMemo, useState } from "react";
 import type { ReflectResponseService } from "../hooks/useGrpc";
+import { wasUpdatedExternally } from "../hooks/useRequestUpdateKey";
 import { showAlert } from "../lib/alert";
+import { showConfirm } from "../lib/confirm";
 import { showDialog } from "../lib/dialog";
+import type { JsonSchema } from "../lib/jsonSchemaExample";
+import { buildExampleFromSchema } from "../lib/jsonSchemaExample";
 import { pluralizeCount } from "../lib/pluralize";
+import { queryClient } from "../lib/queryClient";
 import { Button } from "./core/Button";
+import { Dropdown } from "./core/Dropdown";
 import type { EditorProps } from "./core/Editor/Editor";
 import { Editor } from "./core/Editor/LazyEditor";
 import { GrpcProtoSelectionDialog } from "./GrpcProtoSelectionDialog";
@@ -29,6 +36,11 @@ type Props = Pick<EditorProps, "heightMode" | "onChange" | "className" | "forceU
   protoFiles: string[];
 };
 
+type MethodSchema =
+  | { type: "none" }
+  | { type: "schema"; schema: JsonSchema }
+  | { type: "error"; id: string; title: string; body: ReactNode; log: unknown[] };
+
 export function GrpcEditor({
   services,
   reflectionError,
@@ -42,21 +54,16 @@ export function GrpcEditor({
     setEditorView(h);
   }, []);
 
-  // Find the schema for the selected service and method and update the editor
-  useEffect(() => {
-    if (
-      editorView == null ||
-      services === null ||
-      request.service === null ||
-      request.method === null
-    ) {
-      return;
+  // Find the schema for the selected service and method
+  const methodSchema = useMemo<MethodSchema>(() => {
+    if (services === null || request.service === null || request.method === null) {
+      return { type: "none" };
     }
 
     const s = services.find((s) => s.name === request.service);
     if (s == null) {
-      console.log("Failed to find service", { service: request.service, services });
-      showAlert({
+      return {
+        type: "error",
         id: "grpc-find-service-error",
         title: "Couldn't Find Service",
         body: (
@@ -64,14 +71,14 @@ export function GrpcEditor({
             Failed to find service <InlineCode>{request.service}</InlineCode> in schema
           </>
         ),
-      });
-      return;
+        log: ["Failed to find service", { service: request.service, services }],
+      };
     }
 
     const schema = s.methods.find((m) => m.name === request.method)?.schema;
-    if (request.method != null && schema == null) {
-      console.log("Failed to find method", { method: request.method, methods: s?.methods });
-      showAlert({
+    if (schema == null) {
+      return {
+        type: "error",
         id: "grpc-find-schema-error",
         title: "Couldn't Find Method",
         body: (
@@ -80,18 +87,15 @@ export function GrpcEditor({
             <InlineCode>{request.service}</InlineCode> in schema
           </>
         ),
-      });
-      return;
-    }
-
-    if (schema == null) {
-      return;
+        log: ["Failed to find method", { method: request.method, methods: s.methods }],
+      };
     }
 
     try {
-      updateSchema(editorView, JSON.parse(schema));
+      return { type: "schema", schema: JSON.parse(schema) as JsonSchema };
     } catch (err) {
-      showAlert({
+      return {
+        type: "error",
         id: "grpc-parse-schema-error",
         title: "Failed to Parse Schema",
         body: (
@@ -103,9 +107,22 @@ export function GrpcEditor({
             <FormattedError>{String(err)}</FormattedError>
           </VStack>
         ),
-      });
+        log: ["Failed to parse schema", err],
+      };
     }
-  }, [editorView, services, request.method, request.service]);
+  }, [services, request.method, request.service]);
+
+  useEffect(() => {
+    if (methodSchema.type !== "error") return;
+    console.log(...methodSchema.log);
+    showAlert({ id: methodSchema.id, title: methodSchema.title, body: methodSchema.body });
+  }, [methodSchema]);
+
+  // Update the editor whenever the schema changes
+  useEffect(() => {
+    if (editorView == null || methodSchema.type !== "schema") return;
+    updateSchema(editorView, methodSchema.schema as JSONSchema7);
+  }, [editorView, methodSchema]);
 
   const extraExtensions = useMemo(
     () => [
@@ -124,45 +141,145 @@ export function GrpcEditor({
   const reflectionUnavailable = reflectionError?.match(/unimplemented/i);
   reflectionError = reflectionUnavailable ? undefined : reflectionError;
 
+  const handleGenerateExample = useCallback(async () => {
+    if (methodSchema.type !== "schema") return;
+
+    if (request.message.trim() !== "") {
+      const confirmed = await showConfirm({
+        id: "grpc-generate-example",
+        title: "Generate Example",
+        description: "The current message will be replaced with an example.",
+        confirmText: "Generate",
+      });
+      if (!confirmed) return;
+    }
+
+    const message = JSON.stringify(buildExampleFromSchema(methodSchema.schema), null, 2);
+    await patchModel(request, { message });
+
+    // Force the editor to pick up the new message
+    wasUpdatedExternally(request.id);
+  }, [methodSchema, request]);
+
+  // The reflect query is keyed by request, url and proto files, so a prefix invalidate
+  // reaches it without threading a refetch down from the connection layout.
+  const handleReloadSchema = useCallback(
+    () => queryClient.invalidateQueries({ queryKey: ["grpc_reflect", request.id] }),
+    [request.id],
+  );
+
+  const handleShowReflectionError = useCallback(() => {
+    showDialog({
+      id: "grpc-reflection-error",
+      title: "Reflection Failed",
+      size: "sm",
+      render: ({ hide }) => (
+        <>
+          <FormattedError>{reflectionError ?? "unknown"}</FormattedError>
+          <div className="w-full my-4">
+            <Button
+              className="ml-auto"
+              color="primary"
+              size="sm"
+              onClick={async () => {
+                hide();
+                await handleReloadSchema();
+              }}
+            >
+              Retry
+            </Button>
+          </div>
+        </>
+      ),
+    });
+  }, [handleReloadSchema, reflectionError]);
+
   const actions = useMemo(
     () => [
-      <div key="reflection" className={classNames(services == null && "opacity-100!")}>
-        <Button
-          size="xs"
-          color={
-            reflectionLoading
-              ? "secondary"
-              : reflectionUnavailable
-                ? "info"
-                : reflectionError
-                  ? "danger"
-                  : "secondary"
-          }
-          isLoading={reflectionLoading}
-          onClick={() => {
-            showDialog({
-              title: "Configure Schema",
-              size: "md",
-              id: "reflection-failed",
-              render: ({ hide }) => <GrpcProtoSelectionDialog onDone={hide} />,
-            });
-          }}
+      // Matches the GraphQL editor: one always-visible control labelled by schema state,
+      // with everything schema-related behind it.
+      <div key="schema" className="opacity-100!">
+        <Dropdown
+          items={[
+            {
+              // Hidden for servers without reflection, which isn't an error
+              hidden: !reflectionError,
+              type: "content",
+              label: (
+                <Banner color="danger">
+                  <p className="mb-1">Reflection failed</p>
+                  <Button
+                    size="xs"
+                    color="danger"
+                    variant="border"
+                    onClick={handleShowReflectionError}
+                  >
+                    View Error
+                  </Button>
+                </Banner>
+              ),
+            },
+            {
+              label: "Generate Example Message",
+              leftSlot: <Icon icon="magic_wand" />,
+              disabled: methodSchema.type !== "schema",
+              onSelect: handleGenerateExample,
+            },
+            { type: "separator" },
+            {
+              label: "Reload Schema",
+              leftSlot: <Icon icon="refresh" spin={reflectionLoading} />,
+              keepOpenOnSelect: true,
+              onSelect: handleReloadSchema,
+            },
+            {
+              label: protoFiles.length > 0 ? "Select Proto Files\u2026" : "Configure Schema\u2026",
+              leftSlot: <Icon icon="settings" />,
+              onSelect: () => {
+                showDialog({
+                  title: "Configure Schema",
+                  size: "md",
+                  id: "grpc-configure-schema",
+                  render: ({ hide }) => <GrpcProtoSelectionDialog onDone={hide} />,
+                });
+              },
+            },
+          ]}
         >
-          {reflectionLoading
-            ? "Inspecting Schema"
-            : reflectionUnavailable
-              ? "Select Proto Files"
-              : reflectionError
-                ? "Server Error"
-                : protoFiles.length > 0
-                  ? pluralizeCount("File", protoFiles.length)
-                  : services != null && protoFiles.length === 0
-                    ? "Schema Detected"
-                    : "Select Schema"}
-        </Button>
+          <Button
+            size="sm"
+            variant="border"
+            title="Schema"
+            forDropdown
+            isLoading={reflectionLoading}
+            color={reflectionUnavailable ? "info" : reflectionError ? "danger" : "default"}
+          >
+            {reflectionLoading
+              ? "Inspecting Schema"
+              : reflectionUnavailable
+                ? "Select Proto Files"
+                : reflectionError
+                  ? "Server Error"
+                  : protoFiles.length > 0
+                    ? pluralizeCount("File", protoFiles.length)
+                    : services != null
+                      ? "Schema Detected"
+                      : "Select Schema"}
+          </Button>
+        </Dropdown>
       </div>,
     ],
-    [protoFiles.length, reflectionError, reflectionLoading, reflectionUnavailable, services],
+    [
+      handleGenerateExample,
+      handleReloadSchema,
+      handleShowReflectionError,
+      methodSchema.type,
+      protoFiles.length,
+      reflectionError,
+      reflectionLoading,
+      reflectionUnavailable,
+      services,
+    ],
   );
 
   return (

apps/yaak-client/lib/jsonSchemaExample.test.ts

@@ -0,0 +1,251 @@
+import { describe, expect, test } from "vite-plus/test";
+import type { JsonSchema } from "./jsonSchemaExample";
+import { buildExampleFromSchema } from "./jsonSchemaExample";
+
+describe("buildExampleFromSchema", () => {
+  test("fills scalar fields with placeholders", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        name: { type: "string" },
+        age: { type: "number", format: "int32" },
+        active: { type: "boolean" },
+        data: { type: "string", format: "byte" },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({
+      name: "",
+      age: 0,
+      active: false,
+      data: "",
+    });
+  });
+
+  test("encodes 64-bit integers as strings", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        id: { type: "string", format: "int64" },
+        count: { type: "string", format: "uint64" },
+        offset: { type: "string", format: "sfixed64" },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({ id: "0", count: "0", offset: "0" });
+  });
+
+  test("fills date-time with a parseable timestamp", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: { createdAt: { type: "string", format: "date-time" } },
+    };
+
+    const example = buildExampleFromSchema(schema) as { createdAt: string };
+    expect(Number.isNaN(Date.parse(example.createdAt))).toBe(false);
+  });
+
+  test("fills a duration with a value that parses", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: { timeout: { type: "string", format: "duration" } },
+    };
+
+    // An empty string fails protobuf's Duration parsing, so the message wouldn't send
+    expect(buildExampleFromSchema(schema)).toEqual({ timeout: "0s" });
+  });
+
+  test("expands nested messages through $defs", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: { user: { $ref: "#/$defs/example.User" } },
+      $defs: {
+        "example.User": {
+          type: "object",
+          properties: {
+            name: { type: "string" },
+            address: { $ref: "#/$defs/example.Address" },
+          },
+        },
+        "example.Address": {
+          type: "object",
+          properties: { city: { type: "string" } },
+        },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({
+      user: { name: "", address: { city: "" } },
+    });
+  });
+
+  test("gives repeated fields a single placeholder item", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        tags: { type: "array", items: { type: "string" } },
+        users: { type: "array", items: { $ref: "#/$defs/example.User" } },
+        unknown: { type: "array" },
+      },
+      $defs: {
+        "example.User": { type: "object", properties: { name: { type: "string" } } },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({
+      tags: [""],
+      users: [{ name: "" }],
+      unknown: [],
+    });
+  });
+
+  test("uses the first value of an enum", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        status: { type: "string", enum: ["STATUS_UNSPECIFIED", "STATUS_ACTIVE"] },
+        empty: { type: "string", enum: [] },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({ status: "STATUS_UNSPECIFIED", empty: "" });
+  });
+
+  test("gives maps a single placeholder entry", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        labels: { type: "object", additionalProperties: { type: "string" } },
+        users: { type: "object", additionalProperties: { $ref: "#/$defs/example.User" } },
+      },
+      $defs: {
+        "example.User": { type: "object", properties: { name: { type: "string" } } },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({
+      labels: { key: "" },
+      users: { key: { name: "" } },
+    });
+  });
+
+  test("stops at the root self-reference", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        value: { type: "string" },
+        children: { type: "array", items: { $ref: "#" } },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({ value: "", children: [{}] });
+  });
+
+  test("stops at a cycle between messages", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: { node: { $ref: "#/$defs/example.Node" } },
+      $defs: {
+        "example.Node": {
+          type: "object",
+          properties: {
+            name: { type: "string" },
+            parent: { $ref: "#/$defs/example.Node" },
+            leaf: { $ref: "#/$defs/example.Leaf" },
+          },
+        },
+        "example.Leaf": {
+          type: "object",
+          properties: { node: { $ref: "#/$defs/example.Node" } },
+        },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({
+      node: { name: "", parent: {}, leaf: { node: {} } },
+    });
+  });
+
+  test("expands the same message twice when it is not on the same path", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        from: { $ref: "#/$defs/example.User" },
+        to: { $ref: "#/$defs/example.User" },
+      },
+      $defs: {
+        "example.User": { type: "object", properties: { name: { type: "string" } } },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({ from: { name: "" }, to: { name: "" } });
+  });
+
+  test("fills every branch of a flattened oneof", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        id: { type: "string" },
+        text: { type: "string" },
+        image: { $ref: "#/$defs/example.Image" },
+      },
+      $defs: {
+        "example.Image": { type: "object", properties: { url: { type: "string" } } },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({
+      id: "",
+      text: "",
+      image: { url: "" },
+    });
+  });
+
+  test("stops expanding once the node budget runs out", () => {
+    // Every level references the next one twice, so an unbounded walk would build 2^depth
+    // nodes without ever repeating a ref on the same path.
+    const depth = 16;
+    const $defs: Record<string, JsonSchema> = { [`d${depth}`]: { type: "string" } };
+    for (let i = 0; i < depth; i++) {
+      $defs[`d${i}`] = {
+        type: "object",
+        properties: {
+          a: { $ref: `#/$defs/d${i + 1}` },
+          b: { $ref: `#/$defs/d${i + 1}` },
+        },
+      };
+    }
+
+    const example = buildExampleFromSchema({
+      type: "object",
+      properties: { root: { $ref: "#/$defs/d0" } },
+      $defs,
+    });
+
+    // 2 ** 16 nodes unbounded; the budget holds it to a couple of thousand
+    expect(countNodes(example)).toBeLessThan(10_000);
+  });
+
+  test("handles messages without a known type", () => {
+    const schema: JsonSchema = {
+      type: "object",
+      properties: {
+        empty: {},
+        struct: { type: "object" },
+        missing: { $ref: "#/$defs/example.Nope" },
+      },
+    };
+
+    expect(buildExampleFromSchema(schema)).toEqual({ empty: null, struct: {}, missing: {} });
+  });
+});
+
+function countNodes(value: unknown): number {
+  if (Array.isArray(value)) {
+    return 1 + value.reduce((total: number, v) => total + countNodes(v), 0);
+  }
+  if (value !== null && typeof value === "object") {
+    return 1 + Object.values(value).reduce((total: number, v) => total + countNodes(v), 0);
+  }
+  return 1;
+}

apps/yaak-client/lib/jsonSchemaExample.ts

@@ -0,0 +1,121 @@
+/**
+ * Subset of JSON Schema emitted by the gRPC reflection layer for a method's
+ * input message. See `message_to_json_schema` in the `yaak-grpc` crate.
+ */
+export type JsonSchema = {
+  type?: string;
+  format?: string;
+  properties?: Record<string, JsonSchema>;
+  items?: JsonSchema;
+  additionalProperties?: JsonSchema;
+  enum?: unknown[];
+  $defs?: Record<string, JsonSchema>;
+  $ref?: string;
+};
+
+const DEFS_PREFIX = "#/$defs/";
+const ROOT_REF = "#";
+
+// Protobuf 64-bit integers are encoded as strings in the JSON mapping
+const STRING_NUMBER_FORMATS = ["int64", "uint64", "sint64", "fixed64", "sfixed64"];
+
+// Refs on sibling branches each expand their own subtree, so a schema that references the
+// same messages repeatedly can produce exponentially many nodes without ever cycling.
+const MAX_NODES = 5000;
+
+type Budget = { remaining: number };
+
+/** Build a sample message with placeholder values for every field in the schema */
+export function buildExampleFromSchema(schema: JsonSchema): unknown {
+  // The root is already being built, so a `#` ref anywhere below it is a cycle
+  return buildValue(schema, schema, new Set([ROOT_REF]), { remaining: MAX_NODES });
+}
+
+function buildValue(
+  schema: JsonSchema,
+  root: JsonSchema,
+  refPath: Set<string>,
+  budget: Budget,
+): unknown {
+  if (schema == null || typeof schema !== "object" || budget.remaining <= 0) {
+    return null;
+  }
+
+  budget.remaining -= 1;
+
+  if (typeof schema.$ref === "string") {
+    if (refPath.has(schema.$ref)) {
+      return {};
+    }
+    const resolved = resolveRef(schema.$ref, root);
+    if (resolved == null) {
+      return {};
+    }
+    return buildValue(resolved, root, new Set(refPath).add(schema.$ref), budget);
+  }
+
+  if (Array.isArray(schema.enum)) {
+    return schema.enum[0] ?? "";
+  }
+
+  switch (schema.type) {
+    case "object":
+      return buildObject(schema, root, refPath, budget);
+    case "array":
+      return schema.items == null ? [] : [buildValue(schema.items, root, refPath, budget)];
+    case "string":
+      return buildString(schema.format);
+    case "number":
+      return 0;
+    case "boolean":
+      return false;
+    default:
+      return null;
+  }
+}
+
+function buildObject(
+  schema: JsonSchema,
+  root: JsonSchema,
+  refPath: Set<string>,
+  budget: Budget,
+): unknown {
+  if (schema.properties != null && typeof schema.properties === "object") {
+    const example: Record<string, unknown> = {};
+    for (const [name, propertySchema] of Object.entries(schema.properties)) {
+      example[name] = buildValue(propertySchema, root, refPath, budget);
+    }
+    return example;
+  }
+
+  // Maps have no properties, only a value schema
+  if (schema.additionalProperties != null) {
+    return { key: buildValue(schema.additionalProperties, root, refPath, budget) };
+  }
+
+  return {};
+}
+
+function buildString(format: string | undefined): string {
+  if (format === "date-time") {
+    return new Date().toISOString();
+  }
+  // Duration JSON is a decimal string with an `s` suffix, and an empty one fails to parse
+  if (format === "duration") {
+    return "0s";
+  }
+  if (format != null && STRING_NUMBER_FORMATS.includes(format)) {
+    return "0";
+  }
+  return "";
+}
+
+function resolveRef(ref: string, root: JsonSchema): JsonSchema | null {
+  if (ref === ROOT_REF) {
+    return root;
+  }
+  if (!ref.startsWith(DEFS_PREFIX)) {
+    return null;
+  }
+  return root.$defs?.[ref.slice(DEFS_PREFIX.length)] ?? null;
+}

crates/yaak-grpc/src/json_schema.rs

@@ -210,7 +210,7 @@ fn field_to_type_or_ref(root_name: &str, field: FieldDescriptor) -> JsonSchemaEn
                 // [Protocol Buffers Well-Known Types]: https://protobuf.dev/reference/protobuf/google.protobuf/
                 "google.protobuf.FieldMask" => JsonSchemaEntry::string(),
                 "google.protobuf.Timestamp" => JsonSchemaEntry::string_with_format("date-time"),
-                "google.protobuf.Duration" => JsonSchemaEntry::string(),
+                "google.protobuf.Duration" => JsonSchemaEntry::string_with_format("duration"),
                 "google.protobuf.StringValue" => JsonSchemaEntry::string(),
                 "google.protobuf.BytesValue" => JsonSchemaEntry::string_with_format("byte"),
                 "google.protobuf.Int32Value" => JsonSchemaEntry::number("int32"),