feat(proto): auto-discover payload-visitor roots from coresdk

packages/proto/scripts/gen-payload-visitor.ts

@@ -3,18 +3,16 @@ import { writeFileSync } from 'node:fs';
 import * as protobuf from 'protobufjs';
 import * as prettier from 'prettier';
 
-// Generates `packages/proto/src/payload-visitor.generated.ts`: a synchronous walk of the
-// WorkflowActivation / WorkflowActivationCompletion message trees that invokes a payload visitor at
-// every Payload-bearing field, threading a per-message context. Runs in proto's build (and via
-// `pnpm gen:payload-visitor`).
+// Generates `packages/proto/src/payload-visitor.generated.ts`: a synchronous walk of the worker
+// boundary message trees (WorkflowActivation, ActivityTask, NexusTaskCompletion, ...) that invokes a
+// payload visitor at every Payload-bearing field, threading a per-message context. Runs in proto's
+// build (and via `pnpm gen:payload-visitor`).
 
 const PAYLOAD = 'temporal.api.common.v1.Payload';
 const ANY = 'google.protobuf.Any';
 
-const ROOTS = [
-  { type: 'coresdk.workflow_activation.WorkflowActivation', entry: 'walkWorkflowActivation' },
-  { type: 'coresdk.workflow_completion.WorkflowActivationCompletion', entry: 'walkWorkflowActivationCompletion' },
-] as const;
+// Namespaces scanned for root messages.
+const ROOT_NAMESPACES = ['coresdk'] as const;
 
 const jsonModule = require(resolve(__dirname, '../protos/json-module.js'));
 const root = protobuf.Root.fromJSON(jsonModule);
@@ -40,7 +38,23 @@ function tsType(type: protobuf.Type): string {
   return segments.join('.');
 }
 
-// All message types reachable from the roots, excluding the terminal Payload and the opaque Any.
+/** Collect every message type declared (recursively) under a namespace. */
+function collectMessages(obj: protobuf.ReflectionObject, out: protobuf.Type[]): void {
+  if (obj instanceof protobuf.Type) out.push(obj);
+  if (obj instanceof protobuf.Namespace) {
+    for (const child of obj.nestedArray) collectMessages(child, out);
+  }
+}
+
+// Candidate roots: every message declared under the scanned namespaces.
+const candidates: protobuf.Type[] = [];
+for (const name of ROOT_NAMESPACES) {
+  const ns = root.lookup(name);
+  if (!ns) throw new Error(`Root namespace not found: ${name}`);
+  collectMessages(ns, candidates);
+}
+
+// All message types reachable from the candidates, excluding the terminal Payload and the opaque Any.
 const reachableTypes = (): Map<string, protobuf.Type> => {
   const types = new Map<string, protobuf.Type>();
   const discover = (type: protobuf.Type): void => {
@@ -52,15 +66,14 @@ const reachableTypes = (): Map<string, protobuf.Type> => {
       if (message) discover(message);
     }
   };
-  for (const { type } of ROOTS) discover(root.lookupType(type));
+  for (const candidate of candidates) discover(candidate);
   return types;
 };
 
 const types = reachableTypes();
 
-// Which of those types can contain a Payload, directly or nested inside another message? We work
-// backwards from Payload. First index the reverse references: `referrers.get(X)` is every type
-// that has a field of type X.
+// Reverse reference index: `referrers.get(X)` is every type that has a field of type X. A type absent
+// from this map is embedded by nothing — that is the "is a root" test below.
 const referrers = new Map<string, string[]>();
 for (const type of types.values()) {
   for (const field of type.fieldsArray) {
@@ -87,6 +100,20 @@ while (toVisit.length > 0) {
 // A message contains a Payload if it is one, or if it was marked above.
 const hasPayload = (type: protobuf.Type): boolean => fqn(type) === PAYLOAD || reachesPayload.has(fqn(type));
 
+// Roots: candidate messages that carry a payload and are embedded by no other message.
+const roots = candidates.filter((type) => hasPayload(type) && !referrers.has(fqn(type))).sort(byFqn);
+
+// Public entry-point name for a root, derived from its message name (e.g. `walkWorkflowActivation`).
+const entryName = (type: protobuf.Type): string => `walk${type.name}`;
+
+// Guard: two roots that share a simple name would collide as exported entry points.
+const seenEntries = new Set<string>();
+for (const rootType of roots) {
+  const name = entryName(rootType);
+  if (seenEntries.has(name)) throw new Error(`Duplicate root entry name '${name}' (from ${fqn(rootType)})`);
+  seenEntries.add(name);
+}
+
 /** Every reachable message type that needs a walker (i.e. can contain a Payload), sorted by name. */
 const collectWalkers = (): protobuf.Type[] => [...types.values()].filter(hasPayload).sort(byFqn);
 
@@ -150,12 +177,13 @@ function emit(): string {
   const walkers = collectWalkers();
   const namespaces = [...new Set(walkers.map(topLevelNamespace))].sort();
 
-  const entries = ROOTS.map(({ type, entry }) => {
-    const t = root.lookupType(type);
+  const entries = roots.map((type) => {
     return [
-      `export function ${entry}<Ctx>(root: ${tsType(t)}, env: WalkEnv<Ctx>, context: Ctx): Promise<unknown>[] {`,
+      `export function ${entryName(type)}<Ctx>(root: ${tsType(
+        type
+      )}, env: WalkEnv<Ctx>, context: Ctx): Promise<unknown>[] {`,
       `  const pending: Promise<unknown>[] = [];`,
-      `  ${fnName(t)}(root, env, context, pending);`,
+      `  ${fnName(type)}(root, env, context, pending);`,
       `  return pending;`,
       `}`,
     ].join('\n');