feat: add registerServerFunctionWrap for global server function instrumentation (#1137)

peterp · web-flow · commit bc6a058b9ebf · 2026-04-14T10:10:15.000+02:00
## Summary - Adds `registerServerFunctionWrap()` — a global hook that wraps every `serverAction` / `serverQuery` handler for observability (Sentry spans, OpenTelemetry, timing logs, etc.) - Register once in your worker entry point, every server function is instrumented automatically — no per-action boilerplate, no import shims, no lint rules - Throws if called more than once to prevent silent overwrites ## Usage ```ts // worker.tsx import { registerServerFunctionWrap } from "rwsdk/worker"; import * as Sentry from "@sentry/cloudflare"; registerServerFunctionWrap((fn, args, type) => Sentry.startSpan( { name: fn.name, op: `function.rsc_${type}` }, () => fn(...args) ) ); ``` ## Test plan - [x] 9 new tests covering: action wrapping, query wrapping, fn.name/args passed correctly, interruptors run outside wrapper, wrapper skipped on interruptor short-circuit, double registration throws, no-op without registration, return value transformation, array-style `[interruptors..., handler]` - [x] All 490 existing tests still pass Closes #1135 🤖 Generated with [Claude Code](https://claude.com/claude-code)
diff --git a/sdk/src/runtime/server.test.ts b/sdk/src/runtime/server.test.ts
@@ -0,0 +1,176 @@
+import { describe, expect, it, vi, beforeEach } from "vitest";
+
+let mockRequestInfo: { request: Request; ctx: Record<string, any> };
+
+vi.mock("./requestInfo/worker", () => ({
+  get requestInfo() {
+    return mockRequestInfo;
+  },
+}));
+
+import {
+  serverAction,
+  serverQuery,
+  registerServerFunctionWrap,
+  __resetServerFunctionWrap,
+} from "./server";
+
+describe("registerServerFunctionWrap", () => {
+  beforeEach(() => {
+    __resetServerFunctionWrap();
+    mockRequestInfo = {
+      request: new Request("https://test.example/"),
+      ctx: {},
+    };
+  });
+
+  it("wraps serverAction handlers", async () => {
+    const wrapSpy = vi.fn((fn, args, _type) => fn(...args));
+    registerServerFunctionWrap(wrapSpy);
+
+    const action = serverAction(async function createThing(name: string) {
+      return `created ${name}`;
+    });
+
+    const result = await action("foo");
+
+    expect(wrapSpy).toHaveBeenCalledOnce();
+    expect(wrapSpy.mock.calls[0][2]).toBe("action");
+    expect(result).toBe("created foo");
+  });
+
+  it("wraps serverQuery handlers", async () => {
+    const wrapSpy = vi.fn((fn, args, _type) => fn(...args));
+    registerServerFunctionWrap(wrapSpy);
+
+    const query = serverQuery(async function getThings() {
+      return [1, 2, 3];
+    });
+
+    const result = await query();
+
+    expect(wrapSpy).toHaveBeenCalledOnce();
+    expect(wrapSpy.mock.calls[0][2]).toBe("query");
+    expect(result).toEqual([1, 2, 3]);
+  });
+
+  it("passes the main function and args to the wrapper", async () => {
+    const wrapSpy = vi.fn((fn, args, _type) => fn(...args));
+    registerServerFunctionWrap(wrapSpy);
+
+    const action = serverAction(
+      async function multiply(a: number, b: number) {
+        return a * b;
+      },
+    );
+
+    await action(3, 7);
+
+    const [fn, args] = wrapSpy.mock.calls[0];
+    expect(fn.name).toBe("multiply");
+    expect(args).toEqual([3, 7]);
+  });
+
+  it("interruptors run outside the wrapper", async () => {
+    const order: string[] = [];
+
+    registerServerFunctionWrap((fn, args, _type) => {
+      order.push("wrap:start");
+      const result = fn(...args);
+      order.push("wrap:end");
+      return result;
+    });
+
+    const interruptor = async () => {
+      order.push("interruptor");
+    };
+
+    const action = serverAction([
+      interruptor,
+      async function doWork() {
+        order.push("handler");
+        return "done";
+      },
+    ]);
+
+    await action();
+
+    expect(order).toEqual([
+      "interruptor",
+      "wrap:start",
+      "handler",
+      "wrap:end",
+    ]);
+  });
+
+  it("wrapper is not called when an interruptor short-circuits", async () => {
+    const wrapSpy = vi.fn((fn, args, _type) => fn(...args));
+    registerServerFunctionWrap(wrapSpy);
+
+    const action = serverAction([
+      async () => new Response("blocked", { status: 403 }),
+      async function neverCalled() {
+        return "nope";
+      },
+    ]);
+
+    const result = await action();
+
+    expect(wrapSpy).not.toHaveBeenCalled();
+    expect(result).toBeInstanceOf(Response);
+  });
+
+  it("throws if called more than once", () => {
+    registerServerFunctionWrap(async (fn, args) => fn(...args));
+
+    expect(() => {
+      registerServerFunctionWrap(async (fn, args) => fn(...args));
+    }).toThrow("registerServerFunctionWrap() has already been called");
+  });
+
+  it("without registration, handlers work normally", async () => {
+    const action = serverAction(async function echo(msg: string) {
+      return msg;
+    });
+
+    const result = await action("hello");
+
+    expect(result).toBe("hello");
+  });
+
+  it("wrapper can transform the return value", async () => {
+    registerServerFunctionWrap(async (fn, args, _type) => {
+      const result = await fn(...args);
+      return { wrapped: true, result };
+    });
+
+    const action = serverAction(async function getValue() {
+      return 42;
+    });
+
+    const result = await action();
+
+    expect(result).toEqual({ wrapped: true, result: 42 });
+  });
+
+  it("works with array-style serverAction (interruptors + handler)", async () => {
+    const wrapSpy = vi.fn((fn, args, _type) => fn(...args));
+    registerServerFunctionWrap(wrapSpy);
+
+    const action = serverAction([
+      async ({ ctx }: any) => {
+        ctx.authed = true;
+      },
+      async function save(data: string) {
+        return `saved ${data}`;
+      },
+    ]);
+
+    const result = await action("test");
+
+    expect(result).toBe("saved test");
+    expect(wrapSpy).toHaveBeenCalledOnce();
+    const [fn] = wrapSpy.mock.calls[0];
+    expect(fn.name).toBe("save");
+  });
+});
diff --git a/sdk/src/runtime/server.ts b/sdk/src/runtime/server.ts
@@ -18,10 +18,57 @@ type WrappedServerFunction<TArgs extends any[] = any[], TResult = any> = {
   method?: "GET" | "POST";
 };
 
+export type ServerFunctionWrap = (
+  fn: Function,
+  args: any[],
+  type: "action" | "query",
+) => Promise<any>;
+
+let globalWrap: ServerFunctionWrap | undefined;
+
+/**
+ * Register a wrapper that runs around every server action and query handler.
+ *
+ * Call this once in your worker entry point. The wrapper receives the main
+ * handler function, its arguments, and the type ("action" or "query").
+ * Interruptors run *outside* the wrapper.
+ *
+ * Throws if called more than once.
+ *
+ * @example
+ * ```ts
+ * import { registerServerFunctionWrap } from "rwsdk/worker";
+ * import * as Sentry from "@sentry/cloudflare";
+ *
+ * registerServerFunctionWrap((fn, args, type) =>
+ *   Sentry.startSpan(
+ *     { name: fn.name, op: `function.rsc_${type}` },
+ *     () => fn(...args)
+ *   )
+ * );
+ * ```
+ */
+export function registerServerFunctionWrap(wrap: ServerFunctionWrap) {
+  if (globalWrap) {
+    throw new Error(
+      "registerServerFunctionWrap() has already been called. " +
+        "Only one wrapper can be registered.",
+    );
+  }
+  globalWrap = wrap;
+}
+
+/**
+ * @internal Reset the global wrap — only for tests.
+ */
+export function __resetServerFunctionWrap() {
+  globalWrap = undefined;
+}
+
 function createServerFunction<TArgs extends any[] = any[], TResult = any>(
   fns: Interruptor<TArgs>[],
   mainFn: ServerFunction<TArgs, TResult>,
-  options?: ServerFunctionOptions,
+  options?: ServerFunctionOptions & { __type?: "action" | "query" },
 ): WrappedServerFunction<TArgs, TResult> {
   const wrapped: WrappedServerFunction<TArgs, TResult> = async (
     ...args: TArgs
@@ -38,6 +85,10 @@ function createServerFunction<TArgs extends any[] = any[], TResult = any>(
       }
     }
 
+    if (globalWrap) {
+      return globalWrap(mainFn, args, options?.__type ?? "action") as Promise<TResult>;
+    }
+
     return mainFn(...args);
   };
 
@@ -81,7 +132,7 @@ export function serverQuery<TArgs extends any[], TResult>(
   }
 
   const method = options?.method ?? "GET"; // Default to GET for query
-  const wrapped = createServerFunction(fns, mainFn, { ...options, method });
+  const wrapped = createServerFunction(fns, mainFn, { ...options, method, __type: "query" });
   wrapped.method = method;
   return wrapped;
 }
@@ -121,7 +172,7 @@ export function serverAction<TArgs extends any[], TResult>(
   }
 
   const method = options?.method ?? "POST"; // Default to POST for action
-  const wrapped = createServerFunction(fns, mainFn, { ...options, method });
+  const wrapped = createServerFunction(fns, mainFn, { ...options, method, __type: "action" });
   wrapped.method = method;
   return wrapped;
 }
