fix: add AST extraction to serving type generator and move types to shared/

apps/dev-playground/.gitignore

@@ -1,3 +1,6 @@
 # Playwright
 test-results/
-playwright-report/
+playwright-report/
+
+# Auto-generated types (endpoint-specific, varies per developer)
+shared/appkit-types/serving.d.ts

apps/dev-playground/client/.gitignore

@@ -12,9 +12,6 @@ dist
 dist-ssr
 *.local
 
-# Auto-generated types (endpoint-specific, varies per developer)
-src/appkit-types/serving.d.ts
-
 # Editor directories and files
 .vscode/*
 !.vscode/extensions.json

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

@@ -29,5 +29,5 @@
       "@/*": ["./src/*"]
     }
   },
-  "include": ["src"]
+  "include": ["src", "../shared/appkit-types"]
 }

docs/docs/api/appkit/TypeAlias.ServingFactory.md

@@ -1,19 +1,18 @@
 # Type Alias: ServingFactory
 
 ```ts
-type ServingFactory = keyof ServingEndpointRegistry extends never ? (alias?: string) => ServingEndpointHandle : <K>(alias: K) => ServingEndpointHandle<ServingEndpointRegistry[K]["request"], ServingEndpointRegistry[K]["response"]>;
+type ServingFactory = keyof ServingEndpointRegistry extends never ? (alias?: string) => ServingEndpointHandle : true extends IsUnion<keyof ServingEndpointRegistry> ? <K>(alias: K) => ServingEndpointHandle<ServingEndpointRegistry[K]["request"], ServingEndpointRegistry[K]["response"]> : {
+<K>  (alias: K): ServingEndpointHandle<ServingEndpointRegistry[K]["request"], ServingEndpointRegistry[K]["response"]>;
+  (): ServingEndpointHandle<never, never>;
+};
 ```
 
 Factory function returned by `AppKit.serving`.
 
-This is a conditional type that adapts based on whether `ServingEndpointRegistry`
-has been populated via module augmentation (generated by `appKitServingTypesPlugin()`):
+Adapts based on the `ServingEndpointRegistry` state:
 
-- **Registry empty (default):** `(alias?: string) => ServingEndpointHandle` —
-  accepts any alias string with untyped request/response.
-- **Registry populated:** `<K>(alias: K) => ServingEndpointHandle<...>` —
-  restricts `alias` to known endpoint keys and infers typed request/response
-  from the registry entry.
+- **Empty (default):** `(alias?: string) => ServingEndpointHandle` — any string, untyped.
+- **Single key:** alias optional — `serving()` returns the typed handle for the only endpoint.
+- **Multiple keys:** alias required — must specify which endpoint.
 
-Run `appKitServingTypesPlugin()` in your Vite config to generate the registry
-augmentation and enable full type safety.
+Run `appKitServingTypesPlugin()` in your Vite config to generate the registry.

packages/appkit/src/plugins/serving/serving.ts

@@ -151,24 +151,46 @@ export class ServingPlugin extends Plugin {
         },
       });
     } else {
+      // Unnamed mode: register both /invoke and /:alias/invoke patterns.
+      // The type generator creates a "default" alias, so clients may use either URL.
+      const invokeHandler = async (
+        req: express.Request,
+        res: express.Response,
+      ) => {
+        req.params.alias ??= "default";
+        await this.asUser(req)._handleInvoke(req, res);
+      };
+      const streamHandler = async (
+        req: express.Request,
+        res: express.Response,
+      ) => {
+        req.params.alias ??= "default";
+        await this.asUser(req)._handleStream(req, res);
+      };
+
       this.route(router, {
         name: "invoke",
         method: "post",
         path: "/invoke",
-        handler: async (req: express.Request, res: express.Response) => {
-          req.params.alias = "default";
-          await this.asUser(req)._handleInvoke(req, res);
-        },
+        handler: invokeHandler,
+      });
+      this.route(router, {
+        name: "invoke-named",
+        method: "post",
+        path: "/:alias/invoke",
+        handler: invokeHandler,
       });
-
       this.route(router, {
         name: "stream",
         method: "post",
         path: "/stream",
-        handler: async (req: express.Request, res: express.Response) => {
-          req.params.alias = "default";
-          await this.asUser(req)._handleStream(req, res);
-        },
+        handler: streamHandler,
+      });
+      this.route(router, {
+        name: "stream-named",
+        method: "post",
+        path: "/:alias/stream",
+        handler: streamHandler,
       });
     }
   }

packages/appkit/src/plugins/serving/tests/serving.test.ts

@@ -90,6 +90,16 @@ describe("Serving Plugin", () => {
       expect(handlers["POST:/stream"]).toBeDefined();
     });
 
+    test("also registers /:alias/invoke and /:alias/stream for type-generated clients", () => {
+      const plugin = new ServingPlugin({});
+      const { router, handlers } = createMockRouter();
+
+      plugin.injectRoutes(router);
+
+      expect(handlers["POST:/:alias/invoke"]).toBeDefined();
+      expect(handlers["POST:/:alias/stream"]).toBeDefined();
+    });
+
     test("exports returns a factory that provides invoke", () => {
       const plugin = new ServingPlugin({});
       const factory = plugin.exports() as any;

packages/appkit/src/plugins/serving/types.ts

@@ -49,26 +49,41 @@ export type ServingEndpointHandle<
   ) => ServingEndpointMethods<TRequest, TResponse>;
 };
 
+/** True when T is a union of 2+ members; false for a single literal type. */
+type IsUnion<T, C = T> = T extends C ? ([C] extends [T] ? false : true) : never;
+
 /**
  * Factory function returned by `AppKit.serving`.
  *
- * This is a conditional type that adapts based on whether `ServingEndpointRegistry`
- * has been populated via module augmentation (generated by `appKitServingTypesPlugin()`):
+ * Adapts based on the `ServingEndpointRegistry` state:
  *
- * - **Registry empty (default):** `(alias?: string) => ServingEndpointHandle` —
- *   accepts any alias string with untyped request/response.
- * - **Registry populated:** `<K>(alias: K) => ServingEndpointHandle<...>` —
- *   restricts `alias` to known endpoint keys and infers typed request/response
- *   from the registry entry.
+ * - **Empty (default):** `(alias?: string) => ServingEndpointHandle` — any string, untyped.
+ * - **Single key:** alias optional — `serving()` returns the typed handle for the only endpoint.
+ * - **Multiple keys:** alias required — must specify which endpoint.
  *
- * Run `appKitServingTypesPlugin()` in your Vite config to generate the registry
- * augmentation and enable full type safety.
+ * Run `appKitServingTypesPlugin()` in your Vite config to generate the registry.
  */
 export type ServingFactory = keyof ServingEndpointRegistry extends never
-  ? (alias?: string) => ServingEndpointHandle
-  : <K extends keyof ServingEndpointRegistry>(
-      alias: K,
-    ) => ServingEndpointHandle<
-      ServingEndpointRegistry[K]["request"],
-      ServingEndpointRegistry[K]["response"]
-    >;
+  ? // Empty registry: accept any string, alias optional
+    (alias?: string) => ServingEndpointHandle
+  : true extends IsUnion<keyof ServingEndpointRegistry>
+    ? // Multiple keys: alias REQUIRED for disambiguation
+      <K extends keyof ServingEndpointRegistry>(
+        alias: K,
+      ) => ServingEndpointHandle<
+        ServingEndpointRegistry[K]["request"],
+        ServingEndpointRegistry[K]["response"]
+      >
+    : // Single key: alias optional (runtime defaults to "default")
+      {
+        <K extends keyof ServingEndpointRegistry>(
+          alias: K,
+        ): ServingEndpointHandle<
+          ServingEndpointRegistry[K]["request"],
+          ServingEndpointRegistry[K]["response"]
+        >;
+        (): ServingEndpointHandle<
+          ServingEndpointRegistry[keyof ServingEndpointRegistry]["request"],
+          ServingEndpointRegistry[keyof ServingEndpointRegistry]["response"]
+        >;
+      };

packages/appkit/src/type-generator/index.ts

@@ -2,6 +2,7 @@ import fs from "node:fs/promises";
 import path from "node:path";
 import dotenv from "dotenv";
 import { createLogger } from "../logging/logger";
+import { removeOldGeneratedTypes } from "./migration";
 import { generateQueriesFromDescribe } from "./query-registry";
 import { generateServingTypes as generateServingTypesImpl } from "./serving/generator";
 import type { QuerySchema } from "./types";
@@ -87,6 +88,9 @@ export async function generateFromEntryPoint(options: {
   await fs.mkdir(path.dirname(outFile), { recursive: true });
   await fs.writeFile(outFile, typeDeclarations, "utf-8");
 
+  // One-time migration: remove old generated file from client/src/ to avoid duplicate module augmentation
+  await removeOldGeneratedTypes(outFile, "analytics.d.ts");
+
   logger.debug("Type generation complete!");
 }
 

packages/appkit/src/type-generator/migration.ts

@@ -0,0 +1,29 @@
+import fs from "node:fs/promises";
+import path from "node:path";
+import { createLogger } from "../logging/logger";
+
+const logger = createLogger("type-generator:migration");
+
+/**
+ * Remove old generated types from client/src/appkit-types/ (pre-shared/ location).
+ * Best-effort: silently ignores missing files.
+ */
+export async function removeOldGeneratedTypes(
+  newOutFile: string,
+  filename: string,
+): Promise<void> {
+  const projectRoot = path.resolve(path.dirname(newOutFile), "..", "..");
+  const oldFile = path.join(
+    projectRoot,
+    "client",
+    "src",
+    "appkit-types",
+    filename,
+  );
+  try {
+    await fs.unlink(oldFile);
+    logger.debug("Removed old types at %s", oldFile);
+  } catch {
+    // File doesn't exist — nothing to clean up
+  }
+}

packages/appkit/src/type-generator/serving/generator.ts

@@ -4,6 +4,7 @@ import { WorkspaceClient } from "@databricks/sdk-experimental";
 import pc from "picocolors";
 import { createLogger } from "../../logging/logger";
 import type { EndpointConfig } from "../../plugins/serving/types";
+import { removeOldGeneratedTypes } from "../migration";
 import {
   CACHE_VERSION,
   hashSchema,
@@ -18,6 +19,10 @@ import {
   extractRequestKeys,
 } from "./converter";
 import { fetchOpenApiSchema } from "./fetcher";
+import {
+  extractServingEndpoints,
+  findServerFile,
+} from "./server-file-extractor";
 
 const logger = createLogger("type-generator:serving");
 
@@ -34,14 +39,21 @@ interface GenerateServingTypesOptions {
 /**
  * Generates TypeScript type declarations for serving endpoints
  * by fetching their OpenAPI schemas and converting to TypeScript.
+ *
+ * Endpoint discovery order (when `endpoints` is not provided):
+ * 1. AST extraction from server file (server/index.ts or server/server.ts)
+ * 2. DATABRICKS_SERVING_ENDPOINT_NAME env var (single default endpoint)
  */
 export async function generateServingTypes(
   options: GenerateServingTypesOptions,
 ): Promise<void> {
   const { outFile, noCache } = options;
 
-  // Resolve endpoints from config or env
-  const endpoints = options.endpoints ?? resolveDefaultEndpoints();
+  // Resolve endpoints: explicit > AST extraction from server file > env var fallback
+  const endpoints =
+    options.endpoints ??
+    resolveEndpointsFromServerFile() ??
+    resolveDefaultEndpoints();
   if (Object.keys(endpoints).length === 0) {
     logger.debug("No serving endpoints configured, skipping type generation");
     return;
@@ -77,6 +89,9 @@ export async function generateServingTypes(
   await fs.mkdir(path.dirname(outFile), { recursive: true });
   await fs.writeFile(outFile, output, "utf-8");
 
+  // One-time migration: remove old generated file from client/src/ to avoid duplicate module augmentation
+  await removeOldGeneratedTypes(outFile, "serving.d.ts");
+
   if (registryEntries.length === 0) {
     logger.debug(
       "Wrote empty serving types to %s (no endpoints resolved)",
@@ -227,6 +242,22 @@ function printLogTable(
   console.log("");
 }
 
+function resolveEndpointsFromServerFile():
+  | Record<string, EndpointConfig>
+  | undefined {
+  try {
+    const serverFile = findServerFile(process.cwd());
+    if (!serverFile) return undefined;
+    return extractServingEndpoints(serverFile) ?? undefined;
+  } catch (error) {
+    logger.debug(
+      "Failed to extract endpoints from server file: %s",
+      (error as Error).message,
+    );
+    return undefined;
+  }
+}
+
 function resolveDefaultEndpoints(): Record<string, EndpointConfig> {
   if (process.env.DATABRICKS_SERVING_ENDPOINT_NAME) {
     return { default: { env: "DATABRICKS_SERVING_ENDPOINT_NAME" } };

packages/appkit/src/type-generator/serving/tests/vite-plugin.test.ts

@@ -69,7 +69,7 @@ describe("appKitServingTypesPlugin", () => {
   });
 
   describe("configResolved()", () => {
-    test("resolves outFile relative to config.root", async () => {
+    test("resolves outFile relative to project root", async () => {
       const plugin = appKitServingTypesPlugin({
         endpoints: { llm: { env: "LLM" } },
       });
@@ -79,7 +79,7 @@ describe("appKitServingTypesPlugin", () => {
       expect(mockGenerateServingTypes).toHaveBeenCalledWith(
         expect.objectContaining({
           outFile: expect.stringContaining(
-            "/app/client/src/appkit-types/serving.d.ts",
+            "/app/shared/appkit-types/serving.d.ts",
           ),
         }),
       );

packages/appkit/src/type-generator/serving/vite-plugin.ts

@@ -95,8 +95,8 @@ export function appKitServingTypesPlugin(
       // - pnpm build: process.cwd() is client/ (cd client && vite build), config.root is client/
       projectRoot = path.resolve(config.root, "..");
       outFile = path.resolve(
-        config.root,
-        options?.outFile ?? `src/${TYPES_DIR}/${SERVING_TYPES_FILE}`,
+        projectRoot,
+        options?.outFile ?? `shared/${TYPES_DIR}/${SERVING_TYPES_FILE}`,
       );
     },
 

packages/appkit/src/type-generator/vite-plugin.ts

@@ -27,7 +27,6 @@ interface AppKitTypesPluginOptions {
  * @returns Vite plugin to generate types for AppKit queries.
  */
 export function appKitTypesPlugin(options?: AppKitTypesPluginOptions): Plugin {
-  let root: string;
   let outFile: string;
   let watchFolders: string[];
 
@@ -74,10 +73,10 @@ export function appKitTypesPlugin(options?: AppKitTypesPluginOptions): Plugin {
     },
 
     configResolved(config) {
-      root = config.root;
+      const projectRoot = path.resolve(config.root, "..");
       outFile = path.resolve(
-        root,
-        options?.outFile ?? `src/${TYPES_DIR}/${ANALYTICS_TYPES_FILE}`,
+        projectRoot,
+        options?.outFile ?? `shared/${TYPES_DIR}/${ANALYTICS_TYPES_FILE}`,
       );
       watchFolders = options?.watchFolders ?? [
         path.join(process.cwd(), "config", "queries"),

packages/shared/src/cli/commands/generate-types.ts

@@ -24,7 +24,7 @@ async function runGenerateTypes(
     if (resolvedWarehouseId) {
       const resolvedOutFile =
         outFile ||
-        path.join(process.cwd(), "client/src/appkit-types/analytics.d.ts");
+        path.join(process.cwd(), "shared/appkit-types/analytics.d.ts");
 
       const queryFolder = path.join(resolvedRootDir, "config/queries");
       if (fs.existsSync(queryFolder)) {
@@ -45,7 +45,7 @@ async function runGenerateTypes(
     // Generate serving endpoint types (no warehouse required)
     const servingOutFile = path.join(
       process.cwd(),
-      "client/src/appkit-types/serving.d.ts",
+      "shared/appkit-types/serving.d.ts",
     );
     await typeGen.generateServingTypes({
       outFile: servingOutFile,
@@ -73,7 +73,7 @@ export const generateTypesCommand = new Command("generate-types")
   .argument(
     "[outFile]",
     "Output file path",
-    path.join(process.cwd(), "client/src/appkit-types/analytics.d.ts"),
+    path.join(process.cwd(), "shared/appkit-types/analytics.d.ts"),
   )
   .argument("[warehouseId]", "Databricks warehouse ID")
   .option("--no-cache", "Disable caching for type generation")

template/_gitignore

@@ -8,3 +8,6 @@ build/
 .smoke-test/
 test-results/
 playwright-report/
+
+# Auto-generated types (endpoint-specific, varies per developer)
+shared/appkit-types/serving.d.ts