fix(sdk): resolve 0.4.0 premortem and reconcile PRs #363 and #368

RELEASING.md

@@ -29,8 +29,8 @@ bun run check:bundle         # root-entry size budget + catalog-leak probe
 
 ## Dist-tag policy (semver safety)
 
-- **Pre-1.0 / rc:** version is `1.0.0-rc.N` and `publishConfig.tag` is `next`. An rc does NOT satisfy a consumer's `^0.3`/`~0.3` range and `next` is not installed by default, so a breaking pre-release can never auto-upgrade existing `0.x` consumers. This is the current state.
-- **GA:** flip the version to `1.0.0` and `publishConfig.tag` to `latest` ONLY after the open post-review majors land (see `V1_REVIEW_FIXES.md` Tranches 2–4) — several are wrong public types/behavior that are themselves breaking to fix after GA.
+- **Stable / Bridge releases (`0.4.0`, `1.0.0`):** `publishConfig.tag` is `latest` so `npm install @google/stitch-sdk` installs the active release.
+- **Pre-releases (`1.0.0-rc.N`):** version includes a pre-release suffix (`-rc.N`) and `publishConfig.tag` is `next` (`publish:readiness` enforces `prerelease → next, GA → latest`).
 
 ## Invariants the gates enforce
 

packages/sdk/generated/src/designsystem.ts

@@ -63,6 +63,12 @@ export class DesignSystem {
     private client: StitchToolClientSpec,
     data: any,
   ) {
+    Object.defineProperty(this, "client", {
+      value: client,
+      enumerable: false,
+      writable: true,
+      configurable: true,
+    });
     this.data = typeof data === "object" && data !== null ? data : undefined;
     const _rawName =
       typeof data === "string"

packages/sdk/generated/src/project.ts

@@ -68,6 +68,12 @@ export class Project {
     protected client: StitchToolClientSpec,
     data: any,
   ) {
+    Object.defineProperty(this, "client", {
+      value: client,
+      enumerable: false,
+      writable: true,
+      configurable: true,
+    });
     this.data = typeof data === "object" && data !== null ? data : undefined;
     const _rawName =
       typeof data === "string"
@@ -128,12 +134,18 @@ export class Project {
           modelId?:
             | "MODEL_ID_UNSPECIFIED"
             | "GEMINI_3_8_FLASH"
-            | "GEMINI_3_5_FLASH_LITE";
+            | "GEMINI_3_5_FLASH_LITE"
+            | "GEMINI_3_PRO"
+            | "GEMINI_3_FLASH"
+            | "GEMINI_3_1_PRO";
         },
     modelId?:
       | "MODEL_ID_UNSPECIFIED"
       | "GEMINI_3_8_FLASH"
-      | "GEMINI_3_5_FLASH_LITE",
+      | "GEMINI_3_5_FLASH_LITE"
+      | "GEMINI_3_PRO"
+      | "GEMINI_3_FLASH"
+      | "GEMINI_3_1_PRO",
   ): Promise<Generation<Screen, GenerateScreenFromTextResponse>> {
     const options =
       typeof deviceTypeOrOptions === "object" &&

packages/sdk/generated/src/screen.ts

@@ -64,6 +64,12 @@ export class Screen {
     protected client: StitchToolClientSpec,
     data: any,
   ) {
+    Object.defineProperty(this, "client", {
+      value: client,
+      enumerable: false,
+      writable: true,
+      configurable: true,
+    });
     this.data = typeof data === "object" && data !== null ? data : undefined;
     const _rawName =
       typeof data === "string"
@@ -124,12 +130,18 @@ export class Screen {
           modelId?:
             | "MODEL_ID_UNSPECIFIED"
             | "GEMINI_3_8_FLASH"
-            | "GEMINI_3_5_FLASH_LITE";
+            | "GEMINI_3_5_FLASH_LITE"
+            | "GEMINI_3_PRO"
+            | "GEMINI_3_FLASH"
+            | "GEMINI_3_1_PRO";
         },
     modelId?:
       | "MODEL_ID_UNSPECIFIED"
       | "GEMINI_3_8_FLASH"
-      | "GEMINI_3_5_FLASH_LITE",
+      | "GEMINI_3_5_FLASH_LITE"
+      | "GEMINI_3_PRO"
+      | "GEMINI_3_FLASH"
+      | "GEMINI_3_1_PRO",
   ): Promise<Generation<Screen, EditScreensResponse>> {
     const options =
       typeof deviceTypeOrOptions === "object" &&
@@ -194,12 +206,18 @@ export class Screen {
           modelId?:
             | "MODEL_ID_UNSPECIFIED"
             | "GEMINI_3_8_FLASH"
-            | "GEMINI_3_5_FLASH_LITE";
+            | "GEMINI_3_5_FLASH_LITE"
+            | "GEMINI_3_PRO"
+            | "GEMINI_3_FLASH"
+            | "GEMINI_3_1_PRO";
         },
     modelId?:
       | "MODEL_ID_UNSPECIFIED"
       | "GEMINI_3_8_FLASH"
-      | "GEMINI_3_5_FLASH_LITE",
+      | "GEMINI_3_5_FLASH_LITE"
+      | "GEMINI_3_PRO"
+      | "GEMINI_3_FLASH"
+      | "GEMINI_3_1_PRO",
   ): Promise<Generation<Screen, GenerateVariantsResponse>> {
     const options =
       typeof deviceTypeOrOptions === "object" &&

packages/sdk/generated/src/tool-definitions.ts

@@ -186,6 +186,9 @@ export const toolDefinitions: ToolDefinition[] = [
             "MODEL_ID_UNSPECIFIED",
             "GEMINI_3_8_FLASH",
             "GEMINI_3_5_FLASH_LITE",
+            "GEMINI_3_PRO",
+            "GEMINI_3_FLASH",
+            "GEMINI_3_1_PRO",
           ],
           type: "string",
           "x-google-enum-descriptions": [
@@ -240,6 +243,9 @@ export const toolDefinitions: ToolDefinition[] = [
             "MODEL_ID_UNSPECIFIED",
             "GEMINI_3_8_FLASH",
             "GEMINI_3_5_FLASH_LITE",
+            "GEMINI_3_PRO",
+            "GEMINI_3_FLASH",
+            "GEMINI_3_1_PRO",
           ],
           type: "string",
           "x-google-enum-descriptions": [
@@ -302,6 +308,9 @@ export const toolDefinitions: ToolDefinition[] = [
             "MODEL_ID_UNSPECIFIED",
             "GEMINI_3_8_FLASH",
             "GEMINI_3_5_FLASH_LITE",
+            "GEMINI_3_PRO",
+            "GEMINI_3_FLASH",
+            "GEMINI_3_1_PRO",
           ],
           type: "string",
           "x-google-enum-descriptions": [

packages/sdk/generated/stitch-sdk.lock

@@ -1,12 +1,16 @@
 {
   "schemaVersion": 1,
   "generated": {
-    "generatedAt": "2026-09-23T18:39:30.614Z",
-    "sourceHash": "sha256:a126fbb6fcc1161fd4de0d05e23cbaceacccddeb61f60fa54433d47aba54ff15",
+    "generatedAt": "2026-09-25T04:45:34.219Z",
+    "sourceHash": "sha256:aad79acad078a250f95d5e58767a9082c41be37d4f5565bf307c96f1cd654b38",
     "manifestHash": "sha256:88ec3dfa066ce5060be236b0a5c0b5fb51671a7dd60d05556c028e6c4c8f0e1c",
     "domainMapHash": "sha256:a6177cc7e2f44d2d536928d4e0ae7bd51ed2dd3393235787216b976498c3b1b2",
     "fileCount": 8,
-    "repairedTools": []
+    "repairedTools": [
+      "generate_screen_from_text",
+      "edit_screens",
+      "generate_variants"
+    ]
   },
   "domainMap": {
     "generatedAt": "2026-09-23T18:36:09.460Z",

packages/sdk/package.json

@@ -58,7 +58,7 @@
   "publishConfig": {
     "registry": "https://wombat-dressing-room.appspot.com",
     "access": "public",
-    "tag": "next"
+    "tag": "latest"
   },
   "scripts": {
     "build": "bun scripts/inject-version.ts && tsc",

packages/sdk/src/client.ts

@@ -209,6 +209,23 @@ function normalizeTransportError(err: unknown, toolName: string): unknown {
       retryAfter: parseRetryAfter(retryAfterVal),
     });
   }
+  if (err instanceof Error) {
+    const causeMsg =
+      err.cause instanceof Error ? err.cause.message : String(err.cause ?? "");
+    const combined = `${err.message} ${causeMsg}`;
+    if (
+      /fetch failed|econnreset|socket|other side closed|econnrefused|etimedout|network/i.test(
+        combined,
+      )
+    ) {
+      return new StitchError({
+        code: "NETWORK_ERROR",
+        message: `Tool Call Failed [${toolName}]: ${err.message}`,
+        recoverable: true,
+        toolName,
+      });
+    }
+  }
   return err;
 }
 
@@ -236,6 +253,11 @@ export class StitchToolClient implements StitchToolClientSpec {
   private localVirtualTools: VirtualToolDefinition[] = [];
   public entities: EntityManager;
 
+  /** Whether close() has been called on this client. */
+  public get closed(): boolean {
+    return this.isClosed;
+  }
+
   constructor(
     inputConfig?: Partial<StitchConfig> & {
       localVirtualTools?: VirtualToolDefinition[];
@@ -387,6 +409,7 @@ export class StitchToolClient implements StitchToolClientSpec {
     const maxAttempts = retry ? retry.attempts : 1;
 
     for (let attempt = 0; ; attempt++) {
+      if (!this.isConnected) await this.connect();
       try {
         const result = await this.client.callTool(
           { name, arguments: args },
@@ -395,13 +418,19 @@ export class StitchToolClient implements StitchToolClientSpec {
         );
         return this.parseToolResponse<T>(result, name);
       } catch (rawErr) {
-        // Normalize transport HTTP errors first, so a real 429 is both
-        // classified and retry-eligible (it never reaches parseToolResult).
+        // Normalize transport HTTP / network errors first, so a real 429 or
+        // transient socket reset is classified and retry-eligible.
         const err = normalizeTransportError(rawErr, name);
+        if (err instanceof StitchError && err.code === "NETWORK_ERROR") {
+          this.isConnected = false;
+          this.connectPromise = null;
+        }
         const isRetryable =
           retry !== null &&
           err instanceof StitchError &&
-          (err.code === "RATE_LIMITED" || err.code === "SERVICE_UNAVAILABLE");
+          (err.code === "RATE_LIMITED" ||
+            err.code === "SERVICE_UNAVAILABLE" ||
+            err.code === "NETWORK_ERROR");
         if (!isRetryable || attempt >= maxAttempts - 1) throw err;
         debugLog("retry", `${err.code} on ${name}; backing off`, {
           attempt: attempt + 1,

packages/sdk/src/entity-manager.ts

@@ -82,7 +82,12 @@ export class EntityManager {
   }
 
   constructor(client: any, opts?: { enabled?: boolean }) {
-    this.client = client;
+    Object.defineProperty(this, "client", {
+      value: client,
+      enumerable: false,
+      writable: true,
+      configurable: true,
+    });
     this.enabled = opts?.enabled ?? true;
   }
 

packages/sdk/src/index.ts

@@ -36,6 +36,13 @@ export { StitchError, StitchErrorCode } from "./spec/errors.js";
 // Resource name utilities
 export { parseResourceName } from "./utils.js";
 
+// Schema repair utilities
+export {
+  repairToolSchemas,
+  repairSchema,
+  collectDefPool,
+} from "./schema-repair.js";
+
 // Tool catalog (re-exported with deprecation for 0.4.0 bridge; preferred subpath is @google/stitch-sdk/tools)
 /** @deprecated Import from "@google/stitch-sdk/tools" to optimize bundle size. */
 export { toolDefinitions, toolMap } from "./tools.js";

packages/sdk/src/schema-repair.ts

@@ -103,46 +103,113 @@ function collectRefTargets(
 }
 
 /**
- * Repair a single JSON Schema by injecting any missing well-known $defs
- * that are referenced via $ref but not present.
+ * Legacy 0.3.5 modelId enum literals preserved during the 0.4.0 bridge so
+ * existing callers and Zod tool definitions accept both 0.3.5 and 0.4.0 IDs.
+ */
+const LEGACY_MODEL_IDS = [
+  "GEMINI_3_PRO",
+  "GEMINI_3_FLASH",
+  "GEMINI_3_1_PRO",
+] as const;
+
+/**
+ * Collect every `$defs` entry present across all tools' `inputSchema` and
+ * `outputSchema`, merged on top of `WELL_KNOWN_DEFS`.
+ *
+ * Server-provided `$defs` take precedence over `WELL_KNOWN_DEFS` stubs so
+ * sibling tools share the richest live schema definitions available.
+ */
+export function collectDefPool(tools: Tool[]): Record<string, object> {
+  const pool: Record<string, object> = { ...WELL_KNOWN_DEFS };
+
+  for (const tool of tools) {
+    const inputDefs = (tool.inputSchema as Record<string, any> | undefined)
+      ?.$defs;
+    if (inputDefs && typeof inputDefs === "object") {
+      for (const [name, def] of Object.entries(inputDefs)) {
+        if (def && typeof def === "object") {
+          pool[name] = def as object;
+        }
+      }
+    }
+
+    const outputDefs = (tool as any).outputSchema?.$defs;
+    if (outputDefs && typeof outputDefs === "object") {
+      for (const [name, def] of Object.entries(outputDefs)) {
+        if (def && typeof def === "object") {
+          pool[name] = def as object;
+        }
+      }
+    }
+  }
+
+  return pool;
+}
+
+/**
+ * Repair a single JSON Schema by injecting any missing $defs that are
+ * referenced via $ref (including transitive references inside injected $defs)
+ * and preserving legacy 0.3.5 `modelId` enum literals when present.
  *
  * Mutates the schema in place and returns it for convenience.
  */
-export function repairSchema(schema: Record<string, any>): Record<string, any> {
+export function repairSchema(
+  schema: Record<string, any>,
+  defPool: Record<string, object> = WELL_KNOWN_DEFS,
+): Record<string, any> {
   if (!schema || typeof schema !== "object") return schema;
 
-  const referencedDefs = collectRefTargets(schema);
-  if (referencedDefs.size === 0) return schema;
+  // Preserve 0.3.5 modelId enum values alongside new server modelId literals
+  const modelIdEnum = schema.properties?.modelId?.enum;
+  if (Array.isArray(modelIdEnum)) {
+    for (const legacyId of LEGACY_MODEL_IDS) {
+      if (!modelIdEnum.includes(legacyId)) {
+        modelIdEnum.push(legacyId);
+      }
+    }
+  }
+
+  const initialRefs = collectRefTargets(schema);
+  if (initialRefs.size === 0) return schema;
 
   // Ensure $defs block exists
   schema.$defs = schema.$defs || {};
 
-  for (const defName of referencedDefs) {
-    // Only inject if: (a) the def is missing, and (b) we have a well-known stub
-    if (!schema.$defs[defName] && WELL_KNOWN_DEFS[defName]) {
-      schema.$defs[defName] = { ...WELL_KNOWN_DEFS[defName] };
+  // Multi-pass resolution for transitive $ref dependencies inside injected $defs
+  const MAX_PASSES = 10;
+  for (let pass = 0; pass < MAX_PASSES; pass++) {
+    const referencedDefs = collectRefTargets(schema);
+    let injected = false;
+    for (const defName of referencedDefs) {
+      if (!schema.$defs[defName] && defPool[defName]) {
+        schema.$defs[defName] = structuredClone(defPool[defName]);
+        injected = true;
+      }
     }
+    if (!injected) break;
   }
 
   return schema;
 }
 
 /**
- * Apply schema repair to every tool's inputSchema and outputSchema.
+ * Apply schema repair to every tool's inputSchema and outputSchema using a
+ * shared cross-tool `$defs` pool.
  *
  * This MUST run before the MCP SDK's AJV validator sees the schemas.
  * Mutates tools in place.
  */
 export function repairToolSchemas(tools: Tool[]): void {
+  const pool = collectDefPool(tools);
   for (const tool of tools) {
     if (tool.inputSchema && typeof tool.inputSchema === "object") {
-      repairSchema(tool.inputSchema as Record<string, any>);
+      repairSchema(tool.inputSchema as Record<string, any>, pool);
     }
     // outputSchema was added in MCP SDK ≥1.27 and is the primary crash vector:
     // Client.cacheToolMetadata() eagerly compiles outputSchema with AJV.
     const anyTool = tool as any;
     if (anyTool.outputSchema && typeof anyTool.outputSchema === "object") {
-      repairSchema(anyTool.outputSchema);
+      repairSchema(anyTool.outputSchema, pool);
     }
   }
 }