fix: add safe checker on QueryRegistry and use server static path for fe path (#43)

* fix: add safe checker on QueryRegistry and use server static path for d.ts path

Author: ditadi

…s/dev-playground/config/appKitTypes.d.ts …v-playground/client/src/appKitTypes.d.tsapps/dev-playground/config/appKitTypes.d.ts renamed to apps/dev-playground/client/src/appKitTypes.d.ts apps/dev-playground/config/appKitTypes.d.ts renamed to apps/dev-playground/client/src/appKitTypes.d.ts

@@ -1,6 +1,4 @@
 // Auto-generated by AppKit - DO NOT EDIT
-// Generated at: 2025-11-28T12:10:49.910Z
-
 import "@databricks/app-kit-ui/react";
 
 declare module "@databricks/app-kit-ui/react" {

apps/dev-playground/client/src/routes/analytics.route.tsx

@@ -1,4 +1,3 @@
-// @ts-nocheck
 import { useAnalyticsQuery } from "@databricks/app-kit-ui/react";
 import { createFileRoute, retainSearchParams } from "@tanstack/react-router";
 import { useMemo, useState } from "react";

apps/dev-playground/client/tsconfig.app.json

@@ -29,5 +29,5 @@
       "@/*": ["./src/*"]
     }
   },
-  "include": ["src", "../config/**/*.d.ts"]
+  "include": ["src"]
 }

apps/dev-playground/client/tsconfig.json

@@ -7,8 +7,7 @@
   "compilerOptions": {
     "baseUrl": ".",
     "paths": {
-      "@/*": ["./src/*"],
-      "@databricks/app-kit-ui": ["../../../packages/app-kit-ui/dist"]
+      "@/*": ["./src/*"]
     }
   }
 }

packages/app-kit-ui/src/react/hooks/types.ts

@@ -35,24 +35,39 @@ export interface UseAnalyticsQueryResult<T> {
  * }
  * ```
  */
-// biome-ignore lint/suspicious/noEmptyInterface: Required for module augmentation
-export interface QueryRegistry {}
+export interface QueryRegistry {
+  [key: string]: any[];
+}
+
+/** Gets only literal keys from a registry (excludes index signature) */
+export type AugmentedRegistry<T> = keyof {
+  [K in keyof T as string extends K ? never : K]: T[K];
+};
 
 /** Resolves to registry keys if defined, otherwise string */
-export type QueryKey = keyof QueryRegistry extends never
+export type QueryKey = AugmentedRegistry<QueryRegistry> extends never
   ? string
-  : keyof QueryRegistry;
+  : AugmentedRegistry<QueryRegistry>;
+
+/**
+ * Infers result type: uses QueryRegistry if key exists, otherwise falls back to explicit type T
+ */
+export type InferResult<T, K> = K extends AugmentedRegistry<QueryRegistry>
+  ? QueryRegistry[K]
+  : T;
 
-// biome-ignore lint/suspicious/noEmptyInterface: Required for module augmentation
-export interface PluginRegistry {}
+export interface PluginRegistry {
+  [key: string]: Record<string, any>;
+}
 
-export type PluginName = keyof PluginRegistry extends never
+export type PluginName = AugmentedRegistry<PluginRegistry> extends never
   ? string
-  : keyof PluginRegistry;
+  : AugmentedRegistry<PluginRegistry>;
 
-export type PluginRoutes<P extends PluginName> = P extends keyof PluginRegistry
-  ? keyof PluginRegistry[P]
-  : string;
+export type PluginRoutes<P extends PluginName> =
+  P extends AugmentedRegistry<PluginRegistry>
+    ? AugmentedRegistry<PluginRegistry[P]>
+    : string;
 
 export type RouteResponse<
   P extends PluginName,

packages/app-kit-ui/src/react/hooks/use-analytics-query.ts

@@ -1,18 +1,13 @@
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { connectSSE } from "@/js";
 import type {
+  InferResult,
   QueryKey,
-  QueryRegistry,
   UseAnalyticsQueryOptions,
   UseAnalyticsQueryResult,
 } from "./types";
 import { useQueryHMR } from "./use-query-hmr";
 
-/**
- * Infers result type: uses QueryRegistry if key exists, otherwise falls back to explicit type T
- */
-type InferResult<T, K> = K extends keyof QueryRegistry ? QueryRegistry[K] : T;
-
 function getDevMode() {
   const url = new URL(window.location.href);
   const searchParams = url.searchParams;

packages/app-kit/src/core/app-kit.ts

@@ -1,7 +1,5 @@
 import { existsSync, type FSWatcher, watch } from "node:fs";
 import path from "node:path";
-import type { TelemetryConfig } from "../telemetry";
-import { TelemetryManager } from "../telemetry";
 import type {
   BasePlugin,
   InputPluginMap,
@@ -11,7 +9,10 @@ import type {
   PluginMap,
   QuerySchemas,
 } from "shared";
+import type { TelemetryConfig } from "../telemetry";
+import { TelemetryManager } from "../telemetry";
 import { generatePluginTypes } from "./type-generator";
+
 export class AppKit<TPlugins extends InputPluginMap> {
   private static _instance: AppKit<InputPluginMap> | null = null;
   private pluginInstances: Record<string, BasePlugin> = {};
@@ -131,6 +132,10 @@ export class AppKit<TPlugins extends InputPluginMap> {
   private _generatePluginTypes(
     rawPlugins: PluginData<PluginConstructor, unknown, string>[],
   ) {
+    const serverPlugin = rawPlugins.find((p) => p.name === "server");
+    const staticPath = (serverPlugin?.config as { staticPath?: string })
+      ?.staticPath;
+
     const schemaDir = path.join(process.cwd(), "config/queries");
     const querySchemaPath = path.join(schemaDir, "schema.ts");
 
@@ -152,6 +157,7 @@ export class AppKit<TPlugins extends InputPluginMap> {
       generatePluginTypes(
         rawPlugins.map((p) => ({ name: p.name })),
         querySchemas,
+        staticPath,
       );
     };
     generate();

packages/app-kit/src/core/type-generator.ts

@@ -1,8 +1,8 @@
 import fs from "node:fs";
 import path from "node:path";
-import { routeSchemaRegistry } from "../plugin";
 import type { QuerySchemas } from "shared";
 import { createAuxiliaryTypeStore, printNode, zodToTs } from "zod-to-ts";
+import { routeSchemaRegistry } from "../plugin";
 
 interface PluginTypeInfo {
   name: string;
@@ -20,6 +20,7 @@ function indentType(typeString: string, baseIndent: string): string {
 export function generatePluginTypes(
   plugins: PluginTypeInfo[],
   querySchemas: QuerySchemas,
+  fePath?: string,
 ): void {
   const auxiliaryTypeStore = createAuxiliaryTypeStore();
 
@@ -40,11 +41,6 @@ export function generatePluginTypes(
     .filter(Boolean)
     .join("\n");
 
-  const auxiliaryTypes = Array.from(auxiliaryTypeStore.definitions.values())
-    .map((definition) => printNode(definition.node))
-    .join("\n");
-
-  const auxiliaryPreamble = auxiliaryTypes ? `${auxiliaryTypes}\n\n` : "";
   let queryTypes = "";
   if (querySchemas && Object.keys(querySchemas).length > 0) {
     const queryEntries = Object.entries(querySchemas)
@@ -59,25 +55,27 @@ export function generatePluginTypes(
   }
 
   const content = `// Auto-generated by AppKit - DO NOT EDIT
-// Generated at: ${new Date().toISOString()}
-
 import "@databricks/app-kit-ui/react";
 
-${auxiliaryPreamble}declare module "@databricks/app-kit-ui/react" {
+declare module "@databricks/app-kit-ui/react" {
   interface PluginRegistry {
 ${pluginTypes}
   }
 ${queryTypes}
 }
 `;
 
-  const configDir = path.join(process.cwd(), "config");
-  if (!fs.existsSync(configDir)) {
-    fs.mkdirSync(configDir, { recursive: true });
+  const defaultPath = path.join(process.cwd(), "client", "src");
+  const clientSrcDir = fePath ? fePath.replace(/dist$/, "src") : defaultPath;
+
+  if (!fs.existsSync(clientSrcDir)) {
+    console.log(
+      `[AppKit] Client src dir not found: ${clientSrcDir}, skipping type generation`,
+    );
+    return;
   }
 
-  const filePath = path.join(configDir, "appKitTypes.d.ts");
+  const filePath = path.join(clientSrcDir, "appKitTypes.d.ts");
   fs.writeFileSync(filePath, content, "utf-8");
-
   console.log(`[AppKit] Plugin types generated: ${filePath}`);
 }

packages/app-kit/src/server/index.ts

@@ -1,36 +1,29 @@
 import fs from "node:fs";
 import type { Server as HTTPServer } from "node:http";
 import path from "node:path";
-import { Plugin, toPlugin } from "../plugin";
-import type { BasePluginConfig, PluginPhase } from "shared";
-import { databricksClientMiddleware, isRemoteServerEnabled } from "../utils";
 import dotenv from "dotenv";
 import express from "express";
+import type { PluginPhase } from "shared";
+import { Plugin, toPlugin } from "../plugin";
+import { instrumentations } from "../telemetry";
+import { databricksClientMiddleware, isRemoteServerEnabled } from "../utils";
 import { DevModeManager } from "./dev-mode";
+import type { ServerConfig } from "./types";
 import { getQueries, getRoutes } from "./utils";
-import { instrumentations } from "../telemetry";
 
 dotenv.config({ path: path.resolve(process.cwd(), "./server/.env") });
 
-export interface ServerConfig extends BasePluginConfig {
-  port?: number;
-  plugins?: Record<string, Plugin>;
-  staticPath?: string;
-  autoStart?: boolean;
-  host?: string;
-  watch?: boolean;
-}
-
 export class ServerPlugin extends Plugin {
-  public name = "server" as const;
-  public envVars = ["DATABRICKS_APP_PORT", "FLASK_RUN_HOST"];
   public static DEFAULT_CONFIG = {
     autoStart: true,
     staticPath: path.resolve(process.cwd(), "client", "dist"),
     host: process.env.FLASK_RUN_HOST || "0.0.0.0",
     port: Number(process.env.DATABRICKS_APP_PORT) || 8000,
     watch: process.env.NODE_ENV === "development",
   };
+
+  public name = "server" as const;
+  public envVars = ["DATABRICKS_APP_PORT", "FLASK_RUN_HOST"];
   private serverApplication: express.Application;
   private server: HTTPServer | null;
   private devModeManager?: DevModeManager;
@@ -63,7 +56,7 @@ export class ServerPlugin extends Plugin {
   }
 
   shouldAutoStart() {
-    return this.config.autoStart ?? ServerPlugin.DEFAULT_CONFIG.autoStart;
+    return this.config.autoStart;
   }
 
   isRemoteServingEnabled() {
@@ -99,14 +92,10 @@ export class ServerPlugin extends Plugin {
     }
 
     const server = this.serverApplication.listen(
-      this.config.port || ServerPlugin.DEFAULT_CONFIG.port,
-      this.config.host || ServerPlugin.DEFAULT_CONFIG.host,
+      this.config.port ?? ServerPlugin.DEFAULT_CONFIG.port,
+      this.config.host ?? ServerPlugin.DEFAULT_CONFIG.host,
       () => {
-        console.log(
-          `Server is running on port ${
-            this.config.port || ServerPlugin.DEFAULT_CONFIG.port
-          }`,
-        );
+        console.log(`Server is running on port ${this.config.port}`);
         if (this.config.staticPath && !this.config.watch) {
           console.log(`Serving static files from: ${this.config.staticPath}`);
         }

packages/app-kit/src/server/types.ts

@@ -0,0 +1,11 @@
+import type { BasePluginConfig } from "shared";
+import type { Plugin } from "@/plugin";
+
+export interface ServerConfig extends BasePluginConfig {
+  port?: number;
+  plugins?: Record<string, Plugin>;
+  staticPath?: string;
+  autoStart?: boolean;
+  host?: string;
+  watch?: boolean;
+}