schemaTokenEstimate: size string budgets from .max(), raise ceiling to 32k

Author: jddunn

src/api/runtime/__tests__/schemaTokenEstimate.test.ts

@@ -44,15 +44,35 @@ describe('estimateMaxTokensForZodSchema', () => {
     expect(nestedBudget).toBeGreaterThanOrEqual(1024);
   });
 
-  it('clamps to the maximum (8192) for very large schemas', () => {
+  it('sizes the budget from a string field\'s .max() so large-output schemas do not truncate', () => {
+    // Regression: z.string() was treated as a flat ~30-token leaf regardless
+    // of .max(), so a `source: z.string().max(80_000)` field (designed for
+    // ~20K tokens of generated code) auto-estimated to the 512 floor and the
+    // provider truncated the JSON mid-string. The estimate must scale with the
+    // declared max length.
+    const big = z.object({
+      source: z.string().max(80_000),
+      rationale: z.string().max(500),
+    });
+    const small = z.object({
+      source: z.string().max(2_000),
+      rationale: z.string().max(500),
+    });
+    const bigBudget = estimateMaxTokensForZodSchema(big);
+    const smallBudget = estimateMaxTokensForZodSchema(small);
+    expect(bigBudget).toBeGreaterThan(smallBudget);
+    expect(bigBudget).toBeGreaterThanOrEqual(16_000);
+  });
+
+  it('clamps to the maximum (32000) for very large schemas', () => {
     // Build a wide schema with many fields and deeply nested arrays so the
     // raw walk well exceeds the 8192 cap.
     const wide = z.object({
       a: z.array(z.array(z.array(z.object({ x: z.string(), y: z.string(), z: z.string() })))),
       b: z.array(z.array(z.array(z.object({ x: z.string(), y: z.string(), z: z.string() })))),
       c: z.array(z.array(z.array(z.object({ x: z.string(), y: z.string(), z: z.string() })))),
     });
-    expect(estimateMaxTokensForZodSchema(wide)).toBe(8192);
+    expect(estimateMaxTokensForZodSchema(wide)).toBe(32000);
   });
 
   it('unwraps optional / nullable / default wrappers', () => {
@@ -107,6 +127,6 @@ describe('estimateMaxTokensForZodSchema', () => {
     expect(() => estimateMaxTokensForZodSchema(cyclic)).not.toThrow();
     const tokens = estimateMaxTokensForZodSchema(cyclic);
     expect(tokens).toBeGreaterThanOrEqual(512);
-    expect(tokens).toBeLessThanOrEqual(8192);
+    expect(tokens).toBeLessThanOrEqual(32000);
   });
 });

src/api/runtime/schemaTokenEstimate.ts

@@ -16,7 +16,12 @@
 const TOKENS_PER_LEAF = 30;       // average tokens per primitive field
 const TOKENS_PER_ARRAY_ITEM = 60; // assumed per-element budget for typical strings
 const MIN_BUDGET = 512;
-const MAX_BUDGET = 8192;
+// 2026-05-30: raised 8192 -> 32000. The old 8192 ceiling truncated large-output
+// schemas (e.g. a `source: z.string().max(80_000)` codegen field) even when the
+// model supports far more, which forced every such caller to pass maxTokens by
+// hand. Sonnet 4.x supports ~64K output and Opus 4.8 128K, so 32000 is a safe
+// auto-estimate ceiling; callers needing more still pass maxTokens explicitly.
+const MAX_BUDGET = 32000;
 const HEADROOM = 1.5;             // 50% headroom for prose-heavy fields
 const MAX_DEPTH = 8;              // recursion safety on cyclic-ish schemas
 
@@ -112,7 +117,46 @@ function walk(node: unknown, depth: number): number {
         : TOKENS_PER_LEAF;
     }
 
+    case 'ZodString': {
+      const maxChars = stringMaxChars(def);
+      // A bounded string contributes ~maxChars/4 tokens (~4 chars/token); the
+      // global HEADROOM multiplier covers denser code / JSON-escaped text.
+      // Unbounded strings keep the small leaf default — the caller should add
+      // .max() or pass an explicit maxTokens for large free-form output.
+      return maxChars && maxChars > 0 ? Math.ceil(maxChars / 4) : TOKENS_PER_LEAF;
+    }
+
     default:
       return TOKENS_PER_LEAF;
   }
 }
+
+/**
+ * Extracts a string's `.max(N)` character constraint across Zod v3 + v4
+ * internal shapes. Returns null when the string is unbounded.
+ *
+ * - Zod v3: `_def.checks` is `[{ kind: 'max', value: N }]`.
+ * - Zod v4: the constraint nests under `_zod.def` (or `def`) with
+ *   `check: 'max_length'` and a numeric `maximum`.
+ */
+function stringMaxChars(def: Record<string, unknown>): number | null {
+  const checks = def.checks as unknown;
+  if (!Array.isArray(checks)) return null;
+  for (const c of checks) {
+    const check = c as Record<string, unknown> | null;
+    if (!check) continue;
+    if (check.kind === 'max' && typeof check.value === 'number') {
+      return check.value;
+    }
+    const inner = ((check._zod as Record<string, unknown> | undefined)?.def
+      ?? (check.def as Record<string, unknown> | undefined)
+      ?? check) as Record<string, unknown>;
+    if (
+      (inner.check === 'max_length' || inner.check === 'max')
+      && typeof inner.maximum === 'number'
+    ) {
+      return inner.maximum;
+    }
+  }
+  return null;
+}