Remove generated UI legacy hook shorthand

Author: RhysSullivan

packages/plugins/dynamic-ui/src/shell/hooks.ts

@@ -8,9 +8,7 @@ import {
   useMutation as useTanStackMutation,
   useQuery as useTanStackQuery,
   useQueryClient as useTanStackQueryClient,
-  type MutationKey,
   type QueryClient as QueryClientType,
-  type QueryKey,
   type UseMutationOptions,
   type UseMutationResult,
   type UseQueryOptions,
@@ -19,30 +17,8 @@ import {
 
 export { QueryClient, QueryClientProvider, mutationOptions, queryOptions, skipToken };
 
-let legacyQueryId = 0;
-let legacyMutationId = 0;
 const invalidationScopes: Array<Array<Promise<unknown>>> = [];
 
-type RefetchableQuery = {
-  refetch: () => unknown;
-};
-
-type LegacyMutationOptions<TData, TError, TVariables, TContext> = UseMutationOptions<
-  TData,
-  TError,
-  TVariables,
-  TContext
-> & {
-  invalidates?: readonly RefetchableQuery[];
-};
-
-const nextLegacyQueryKey = (): QueryKey => ["executor", "legacy-query", ++legacyQueryId];
-const nextLegacyMutationKey = (): MutationKey => [
-  "executor",
-  "legacy-mutation",
-  ++legacyMutationId,
-];
-
 const trackInvalidation = (promise: Promise<unknown>) => {
   for (const scope of invalidationScopes) {
     scope.push(promise);
@@ -105,88 +81,34 @@ const wrapMutationOptions = <TData, TError, TVariables, TContext>(
     : undefined,
 });
 
-/**
- * TanStack Query's `useQuery`, plus compatibility for the original generated
- * UI shorthand: `useQuery(() => tools.namespace.tool(args))`.
- */
-export function useQuery<TData>(fn: () => Promise<TData>): UseQueryResult<TData, Error>;
 export function useQuery<
   TQueryFnData,
   TError = Error,
   TData = TQueryFnData,
-  TQueryKey extends QueryKey = QueryKey,
+  TQueryKey extends readonly unknown[] = readonly unknown[],
 >(
   options: UseQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
   queryClient?: QueryClientType,
 ): UseQueryResult<TData, TError>;
-export function useQuery(
-  optionsOrFn: unknown,
+export function useQuery<
+  TQueryFnData,
+  TError = Error,
+  TData = TQueryFnData,
+  TQueryKey extends readonly unknown[] = readonly unknown[],
+>(
+  options: UseQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
   queryClient?: QueryClientType,
-): UseQueryResult<unknown, Error> {
-  const legacyFnRef = useRef<(() => Promise<unknown>) | null>(null);
-  const legacyQueryKeyRef = useRef<QueryKey | null>(null);
-
-  if (typeof optionsOrFn === "function") {
-    legacyFnRef.current = optionsOrFn as () => Promise<unknown>;
-    legacyQueryKeyRef.current ??= nextLegacyQueryKey();
-    return useTanStackQuery(
-      {
-        queryKey: legacyQueryKeyRef.current,
-        queryFn: () => {
-          const current = legacyFnRef.current;
-          if (!current) throw new Error("Missing legacy query function.");
-          return current();
-        },
-      },
-      queryClient,
-    );
-  }
-
-  return useTanStackQuery(
-    optionsOrFn as UseQueryOptions<unknown, Error, unknown, QueryKey>,
-    queryClient,
-  );
+): UseQueryResult<TData, TError> {
+  return useTanStackQuery(options, queryClient);
 }
 
-/**
- * TanStack Query's `useMutation`, plus compatibility for the original generated
- * UI shorthand: `useMutation((input) => tools.namespace.tool(input), opts)`.
- */
-export function useMutation<TVariables, TData = unknown, TContext = unknown>(
-  fn: (input: TVariables) => Promise<TData>,
-  options?: LegacyMutationOptions<TData, Error, TVariables, TContext>,
-): UseMutationResult<TData, Error, TVariables, TContext>;
 export function useMutation<TData = unknown, TError = Error, TVariables = void, TContext = unknown>(
   options: UseMutationOptions<TData, TError, TVariables, TContext>,
   queryClient?: QueryClientType,
 ): UseMutationResult<TData, TError, TVariables, TContext>;
-export function useMutation(
-  optionsOrFn: unknown,
-  optionsOrQueryClient?: unknown,
-): UseMutationResult<unknown, Error, unknown, unknown> {
-  const legacyMutationKeyRef = useRef<MutationKey | null>(null);
-  legacyMutationKeyRef.current ??= nextLegacyMutationKey();
-
-  if (typeof optionsOrFn === "function") {
-    const legacyOptions =
-      (optionsOrQueryClient as
-        | LegacyMutationOptions<unknown, Error, unknown, unknown>
-        | undefined) ?? {};
-    const { invalidates, onSuccess, ...tanstackOptions } = legacyOptions;
-
-    return useTanStackMutation({
-      mutationKey: legacyOptions.mutationKey ?? legacyMutationKeyRef.current,
-      ...tanstackOptions,
-      mutationFn: optionsOrFn as (input: unknown) => Promise<unknown>,
-      onSuccess: async (data, variables, context, mutationContext) => {
-        await trackMutationCallback(() => onSuccess?.(data, variables, context, mutationContext));
-        await Promise.all(invalidates?.map((query) => query.refetch()) ?? []);
-      },
-    });
-  }
-
-  return useTanStackMutation(
-    wrapMutationOptions(optionsOrFn as UseMutationOptions<unknown, Error, unknown, unknown>),
-    optionsOrQueryClient as QueryClientType | undefined,
-  );
+export function useMutation<TData = unknown, TError = Error, TVariables = void, TContext = unknown>(
+  options: UseMutationOptions<TData, TError, TVariables, TContext>,
+  queryClient?: QueryClientType,
+): UseMutationResult<TData, TError, TVariables, TContext> {
+  return useTanStackMutation(wrapMutationOptions(options), queryClient);
 }

packages/plugins/dynamic-ui/src/shell/mcp-app.browser.test.ts

@@ -172,7 +172,9 @@ const blockedMessages = primitiveNames.map((name) => {
 });
 
 function App() {
-  const { data, error, isLoading } = useQuery(() => tools.inventory.items.listItems({}));
+  const { data, error, isLoading } = useQuery(
+    tools.inventory.items.listItems.queryOptions({})
+  );
   const items = data?.ok ? data.data : data;
   return (
     <Card>
@@ -200,13 +202,15 @@ function App() {
 const generatedApprovalCode = `
 function App() {
   const [status, setStatus] = useState("idle");
-  const createItem = useMutation((input) => tools.inventory.items.createItem(input), {
-    onSuccess: (result) => {
-      const created = result?.ok ? result.data : result;
-      setStatus(created.name + ":" + created.created);
-    },
-    onError: (error) => setStatus(error.message),
-  });
+  const createItem = useMutation(
+    tools.inventory.items.createItem.mutationOptions({
+      onSuccess: (result) => {
+        const created = result?.ok ? result.data : result;
+        setStatus(created.name + ":" + created.created);
+      },
+      onError: (error) => setStatus(error.message),
+    })
+  );
   const ask = async () => {
     setStatus("pending");
     await createItem.mutateAsync({ body: { name: "Approved Widget" } });
@@ -227,13 +231,15 @@ function App() {
 const generatedAutoMutationCode = `
 function App() {
   const [status, setStatus] = useState("idle");
-  const createItem = useMutation((input) => tools.inventory.items.createItem(input), {
-    onSuccess: (result) => {
-      const created = result?.ok ? result.data : result;
-      setStatus(created.name + ":" + created.created);
-    },
-    onError: (error) => setStatus(error.message),
-  });
+  const createItem = useMutation(
+    tools.inventory.items.createItem.mutationOptions({
+      onSuccess: (result) => {
+        const created = result?.ok ? result.data : result;
+        setStatus(created.name + ":" + created.created);
+      },
+      onError: (error) => setStatus(error.message),
+    })
+  );
 
   useEffect(() => {
     createItem.mutate({ body: { name: "Mount Widget" } });
@@ -293,12 +299,14 @@ function App() {
 const generatedSchemaApprovalCode = `
 function App() {
   const [status, setStatus] = useState("idle");
-  const requestDetails = useMutation((input) => tools.profile.submit(input), {
-    onSuccess: (result) => {
-      setStatus(JSON.stringify(result));
-    },
-    onError: (error) => setStatus(error.message),
-  });
+  const requestDetails = useMutation(
+    tools.profile.submit.mutationOptions({
+      onSuccess: (result) => {
+        setStatus(JSON.stringify(result));
+      },
+      onError: (error) => setStatus(error.message),
+    })
+  );
 
   return (
     <Card>