feat(appkit): embed typegen cache in generated .d.ts files

apps/dev-playground/.gitignore

@@ -1,6 +1,3 @@
 # Playwright
 test-results/
-playwright-report/
-
-# Auto-generated types (regenerated on `pnpm dev` by appKitTypesPlugin)
-shared/appkit-types/serving.d.ts
+playwright-report/

packages/appkit/src/plugins/serving/schema-filter.ts

@@ -1,54 +1,50 @@
 import fs from "node:fs/promises";
 import { createLogger } from "../../logging/logger";
 import {
-  CACHE_VERSION,
-  type ServingCache,
-} from "../../type-generator/serving/cache";
+  objectMemberValue,
+  objectTopLevelKeys,
+  splitEntryBlocks,
+} from "../../type-generator/embedded-cache";
 
 const logger = createLogger("serving:schema-filter");
 
-function isValidCache(data: unknown): data is ServingCache {
-  return (
-    typeof data === "object" &&
-    data !== null &&
-    "version" in data &&
-    (data as ServingCache).version === CACHE_VERSION &&
-    "endpoints" in data &&
-    typeof (data as ServingCache).endpoints === "object"
-  );
-}
+/** Registry interface name whose members hold the serving type blocks. */
+const SERVING_INTERFACE = "ServingEndpointRegistry";
 
 /**
- * Loads endpoint schemas from the type generation cache file.
- * Returns a map of alias → allowed parameter keys.
+ * Load per-endpoint request-parameter allowlists from the committed generated
+ * `serving.d.ts`. The allowlist for an alias is the set of top-level keys of
+ * its rendered `request` object type (`stream` is already excluded at
+ * generation time). A generic `Record<string, unknown>` request has no keys, so
+ * it yields no allowlist (passthrough). A missing file → no filtering
+ * (passthrough).
  */
 export async function loadEndpointSchemas(
-  cacheFile: string,
+  typesFile: string,
 ): Promise<Map<string, Set<string>>> {
   const allowlists = new Map<string, Set<string>>();
 
   try {
-    const raw = await fs.readFile(cacheFile, "utf8");
-    const parsed: unknown = JSON.parse(raw);
-    if (!isValidCache(parsed)) {
-      logger.warn("Serving types cache has invalid structure, skipping");
-      return allowlists;
-    }
-    const cache = parsed;
-
-    for (const [alias, entry] of Object.entries(cache.endpoints)) {
-      if (entry.requestKeys && entry.requestKeys.length > 0) {
-        allowlists.set(alias, new Set(entry.requestKeys));
+    const source = await fs.readFile(typesFile, "utf8");
+    const blocks = splitEntryBlocks(source, SERVING_INTERFACE);
+    for (const [alias, block] of Object.entries(blocks)) {
+      // Each member is `{ request: <type>; response: ...; chunk: ...; }`.
+      // Extract the `request` value, then its top-level object keys.
+      const requestType = objectMemberValue(block, "request");
+      if (!requestType) continue;
+      const keys = objectTopLevelKeys(requestType);
+      if (keys.length > 0) {
+        allowlists.set(alias, new Set(keys));
       }
     }
   } catch (err) {
     if ((err as NodeJS.ErrnoException).code !== "ENOENT") {
       logger.warn(
-        "Failed to load serving types cache: %s",
+        "Failed to load serving types for request filtering: %s",
         (err as Error).message,
       );
     }
-    // No cache → no filtering, passthrough mode
+    // No file → no filtering, passthrough mode
   }
 
   return allowlists;

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

@@ -66,14 +66,15 @@ export class ServingPlugin extends Plugin {
   }
 
   async setup(): Promise<void> {
-    const cacheFile = path.join(
+    // The request-parameter allowlist is derived from the committed generated
+    // serving types.
+    const typesFile = path.join(
       process.cwd(),
-      "node_modules",
-      ".databricks",
-      "appkit",
-      ".appkit-serving-types-cache.json",
+      "shared",
+      "appkit-types",
+      "serving.d.ts",
     );
-    this.schemaAllowlists = await loadEndpointSchemas(cacheFile);
+    this.schemaAllowlists = await loadEndpointSchemas(typesFile);
     if (this.schemaAllowlists.size > 0) {
       logger.debug(
         "Loaded schema allowlists for %d endpoint(s)",

packages/appkit/src/plugins/serving/tests/schema-filter.test.ts

@@ -109,50 +109,53 @@ describe("schema-filter", () => {
       expect(result.size).toBe(0);
     });
 
-    test("reads requestKeys from cache entries", async () => {
+    test("derives the allowlist from the request object's top-level keys", async () => {
       const fs = (await import("node:fs/promises")).default;
-      vi.mocked(fs.readFile).mockResolvedValue(
-        JSON.stringify({
-          version: "1",
-          endpoints: {
-            default: {
-              hash: "abc",
-              requestType: "{}",
-              responseType: "{}",
-              chunkType: null,
-              requestKeys: ["messages", "temperature", "max_tokens"],
-            },
-          },
-        }),
-      );
-
-      const result = await loadEndpointSchemas("/some/path");
+      // A committed serving.d.ts: the `request` type's top-level keys are the
+      // allowlist for the alias.
+      vi.mocked(fs.readFile).mockResolvedValue(`import "@databricks/appkit";
+declare module "@databricks/appkit" {
+  interface ServingEndpointRegistry {
+    default: {
+      request: {
+        messages: Array<{ role: string; content: string }>;
+        temperature: number;
+        max_tokens: number;
+      };
+      response: { model: string };
+      chunk: unknown;
+    };
+  }
+}
+`);
+
+      const result = await loadEndpointSchemas("/some/path/serving.d.ts");
       expect(result.size).toBe(1);
       const keys = result.get("default");
       expect(keys).toBeDefined();
       expect(keys?.has("messages")).toBe(true);
       expect(keys?.has("temperature")).toBe(true);
       expect(keys?.has("max_tokens")).toBe(true);
+      // Nested keys (e.g. `role`, `content`) are NOT part of the top-level allowlist.
+      expect(keys?.has("role")).toBe(false);
     });
 
-    test("skips entries without requestKeys (backwards compat)", async () => {
+    test("a generic Record request yields no allowlist (passthrough mode)", async () => {
       const fs = (await import("node:fs/promises")).default;
-      vi.mocked(fs.readFile).mockResolvedValue(
-        JSON.stringify({
-          version: "1",
-          endpoints: {
-            default: {
-              hash: "abc",
-              requestType: "{ messages: string[] }",
-              responseType: "{}",
-              chunkType: null,
-            },
-          },
-        }),
-      );
-
-      const result = await loadEndpointSchemas("/some/path");
-      // No requestKeys → passthrough mode (no allowlist)
+      vi.mocked(fs.readFile).mockResolvedValue(`import "@databricks/appkit";
+declare module "@databricks/appkit" {
+  interface ServingEndpointRegistry {
+    default: {
+      request: Record<string, unknown>;
+      response: unknown;
+      chunk: unknown;
+    };
+  }
+}
+`);
+
+      const result = await loadEndpointSchemas("/some/path/serving.d.ts");
+      // Record<string, unknown> has no top-level keys → passthrough (no allowlist).
       expect(result.size).toBe(0);
     });
   });