From 235339a7bea0ed28a31422ea50696a3ff385a28a Mon Sep 17 00:00:00 2001 From: iza <59828082+izadoesdev@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:46:48 +0300 Subject: [PATCH] fix(ai): preserve full RPC paths with the native client --- packages/ai/src/ai/tools/utils/rpc.test.ts | 65 ++++++++++++++++------ packages/ai/src/ai/tools/utils/rpc.ts | 22 ++------ packages/ai/src/lib/orpc-server.ts | 24 -------- 3 files changed, 54 insertions(+), 57 deletions(-) delete mode 100644 packages/ai/src/lib/orpc-server.ts diff --git a/packages/ai/src/ai/tools/utils/rpc.test.ts b/packages/ai/src/ai/tools/utils/rpc.test.ts index 75ee92983a..5a7d7a7479 100644 --- a/packages/ai/src/ai/tools/utils/rpc.test.ts +++ b/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>>(); + +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); + }); }); diff --git a/packages/ai/src/ai/tools/utils/rpc.ts b/packages/ai/src/ai/tools/utils/rpc.ts index f9c9faed4e..1a7957a347 100644 --- a/packages/ai/src/ai/tools/utils/rpc.ts +++ b/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; -} diff --git a/packages/ai/src/lib/orpc-server.ts b/packages/ai/src/lib/orpc-server.ts deleted file mode 100644 index 5e19870737..0000000000 --- a/packages/ai/src/lib/orpc-server.ts +++ /dev/null @@ -1,24 +0,0 @@ -import type { AppRouter, PreResolvedAuth } from "@databuddy/rpc"; -import { appRouter, createRPCContext } from "@databuddy/rpc"; -import type { RouterClient } from "@orpc/server"; -import { createRouterClient } from "@orpc/server"; - -export async function getServerRPCClient( - headers: Headers, - preResolved?: PreResolvedAuth -): Promise> { - const rpcContext = await createRPCContext({ headers }, preResolved); - - const client = {} as RouterClient; - - for (const [routerName, router] of Object.entries(appRouter)) { - if (router && typeof router === "object") { - const routerClient = createRouterClient(router as any, { - context: rpcContext, - }); - (client as Record)[routerName] = routerClient; - } - } - - return client; -}