fix(ai): preserve full RPC paths with the native client

packages/ai/src/ai/tools/utils/rpc.test.ts

@@ -1,23 +1,36 @@
 import { describe, expect, it, mock } from "bun:test";
+import { ORPCError, os } from "@orpc/server";
+import { z } from "zod";
 import type { AppContext } from "../../config/context";
 
 let observedSignal: AbortSignal | undefined;
 
-mock.module("../../../lib/orpc-server", () => ({
-	getServerRPCClient: async () => ({
+const createRPCContext = mock(
+	async (
+		opts: { headers: Headers },
+		serviceAuth?: AppContext["serviceAuth"]
+	) => ({
+		...opts,
+		serviceAuth,
+	})
+);
+const procedure = os.$context<Awaited<ReturnType<typeof createRPCContext>>>();
+
+mock.module("@databuddy/rpc", () => ({
+	createRPCContext,
+	appRouter: {
 		links: {
-			create: async () => {
-				throw new Error("Live RPC should not be called in this test");
-			},
-			list: async (
-				_input: unknown,
-				options?: { signal?: AbortSignal }
-			) => {
-				observedSignal = options?.signal;
-				return [];
-			},
+			create: procedure.handler(() => {
+				throw new ORPCError("FORBIDDEN");
+			}),
+			list: procedure
+				.input(z.object({ websiteId: z.string() }))
+				.handler(({ context, path, signal }) => {
+					observedSignal = signal;
+					return { context, path };
+				}),
 		},
-	}),
+	},
 }));
 
 const { callRPCProcedure } = await import("./rpc");
@@ -28,12 +41,13 @@ const BASE_CONTEXT: AppContext = {
 	requestHeaders: new Headers(),
 	timezone: "UTC",
 	userId: "eval-user",
-	websiteDomain: "databuddy.cc",
+	websiteDomain: "example.com",
 	websiteId: "website_123",
 };
 
 describe("AI tool RPC helper", () => {
 	it("blocks mutation RPC calls in dry-run mode", async () => {
+		const callsBefore = createRPCContext.mock.calls.length;
 		const result = await callRPCProcedure(
 			"links",
 			"create",
@@ -54,19 +68,36 @@ describe("AI tool RPC helper", () => {
 			{ ...BASE_CONTEXT, mutationMode: "dry-run" }
 		);
 		expect(reply).toMatchObject({ dryRun: true, mutationBlocked: true });
+		expect(createRPCContext).toHaveBeenCalledTimes(callsBefore);
 	});
 
-	it("forwards cancellation to the ORPC client", async () => {
+	it("preserves the full procedure path, authentication context, and cancellation", async () => {
 		const controller = new AbortController();
+		const serviceAuth = { apiKey: null, session: null };
 
-		await callRPCProcedure(
+		const result = await callRPCProcedure(
 			"links",
 			"list",
 			{ websiteId: "website_123" },
-			BASE_CONTEXT,
+			{ ...BASE_CONTEXT, serviceAuth },
 			controller.signal
 		);
 
+		expect(result).toEqual({
+			context: { headers: BASE_CONTEXT.requestHeaders, serviceAuth },
+			path: ["links", "list"],
+		});
 		expect(observedSignal).toBe(controller.signal);
 	});
+
+	it.each([
+		["missing", "list", "Router missing not found"],
+		["links", "missing", "Procedure links.missing not found or not callable."],
+		["links", "create", "You don't have permission to access this resource."],
+		["links", "list", "Invalid request: Input validation failed"],
+	])("preserves the error for %s.%s", async (router, method, message) => {
+		await expect(
+			callRPCProcedure(router, method, {}, BASE_CONTEXT)
+		).rejects.toThrow(message);
+	});
 });

packages/ai/src/ai/tools/utils/rpc.ts

@@ -1,5 +1,4 @@
-import { ORPCError } from "@orpc/server";
-import type { PreResolvedAuth } from "@databuddy/rpc";
+import { createRouterClient, ORPCError } from "@orpc/server";
 import type { AppContext } from "../../config/context";
 import { createToolLogger } from "./logger";
 
@@ -15,7 +14,7 @@ export async function callRPCProcedure(
 	abortSignal?: AbortSignal
 ) {
 	try {
-		if (context.mutationMode === "dry-run" && isMutationMethod(method)) {
+		if (context.mutationMode === "dry-run" && MUTATION_METHOD_RE.test(method)) {
 			return {
 				dryRun: true,
 				message: `Dry-run mode blocked ${routerName}.${method}; no data was changed.`,
@@ -25,9 +24,10 @@ export async function callRPCProcedure(
 		}
 
 		const headers = context.requestHeaders ?? new Headers();
-		const preResolved = resolvePreResolvedAuth(context);
-		const { getServerRPCClient } = await import("../../../lib/orpc-server");
-		const client = await getServerRPCClient(headers, preResolved);
+		const { appRouter, createRPCContext } = await import("@databuddy/rpc");
+		const client = createRouterClient(appRouter, {
+			context: await createRPCContext({ headers }, context.serviceAuth),
+		});
 
 		const router = client[routerName as keyof typeof client] as
 			| Record<
@@ -95,13 +95,3 @@ export async function callRPCProcedure(
 		throw new Error("An unexpected error occurred. Please try again.");
 	}
 }
-
-function isMutationMethod(method: string): boolean {
-	return MUTATION_METHOD_RE.test(method);
-}
-
-function resolvePreResolvedAuth(
-	context: AppContext
-): PreResolvedAuth | undefined {
-	return context.serviceAuth;
-}