fix(generation): expand output budget

packages/diagram/generation/src/lib/cloudflare-google-ai-studio.ts

@@ -15,6 +15,7 @@ import {
 } from "./errors.js";
 import {
   buildGeminiGenerateContentBody,
+  extractGeminiFinishReason,
   extractGeminiText,
   extractGeminiUsage,
   stripGoogleModelPrefix,
@@ -183,8 +184,15 @@ const runGatewayAttempt = Effect.fn(
       }),
   });
   const usage = extractGeminiUsage(raw);
+  const finishReason = extractGeminiFinishReason(raw);
 
   return candidateFromText({
+    diagnostics:
+      finishReason === "MAX_TOKENS"
+        ? [
+            "output_truncated: Gemini stopped at the maximum output-token budget; regenerate the complete diagram.",
+          ]
+        : [],
     model,
     provider: "cloudflare-google-ai-studio",
     raw,

packages/diagram/generation/src/lib/gemini.ts

@@ -36,7 +36,7 @@ export class GeminiGenerateContentBody extends Schema.Class<GeminiGenerateConten
   system_instruction: GeminiSystemInstruction,
 }) {}
 
-const DEFAULT_MAX_OUTPUT_TOKENS = 2_048;
+const DEFAULT_MAX_OUTPUT_TOKENS = 16_384;
 const DEFAULT_TEMPERATURE = 0.1;
 
 function messageContent(
@@ -141,3 +141,12 @@ export function extractGeminiUsage(
 
   return Object.keys(usage).length > 0 ? usage : undefined;
 }
+
+export function extractGeminiFinishReason(
+  response: unknown,
+): string | undefined {
+  const candidates = objectValue(response, "candidates");
+  const firstCandidate = Array.isArray(candidates) ? candidates[0] : undefined;
+  const finishReason = objectValue(firstCandidate, "finishReason");
+  return typeof finishReason === "string" ? finishReason : undefined;
+}

packages/diagram/generation/src/lib/messages.ts

@@ -36,7 +36,7 @@ export class DiagramGenerationMessages extends Schema.Class<DiagramGenerationMes
 }) {}
 
 const FLOWCHART_IR_INSTRUCTIONS = [
-  "Return only JSON. Do not wrap the JSON in markdown.",
+  "Return only compact, minified JSON on one line. Do not use markdown.",
   'Use type "flowchart".',
   'Every node must have id, label, and kind: "start", "process", "decision", or "end".',
   "Use exactly one start node and at least one end node.",
@@ -49,7 +49,7 @@ const FLOWCHART_IR_INSTRUCTIONS = [
 ];
 
 const MINDMAP_IR_INSTRUCTIONS = [
-  "Return only JSON. Do not wrap the JSON in markdown.",
+  "Return only compact, minified JSON on one line. Do not use markdown.",
   'Use type "mindmap".',
   'Every node must have id, label, kind ("root" or "topic"), and metadata with depth and siblingIndex.',
   "Use exactly one root node at depth 0 with siblingIndex 0.",
@@ -62,62 +62,54 @@ const MINDMAP_IR_INSTRUCTIONS = [
 
 function expectedJsonShape(prompt: DiagramGenerationPrompt): string {
   if (prompt.type === "mindmap") {
-    return JSON.stringify(
-      {
-        id: "short-kebab-case-id",
-        title: prompt.title,
-        type: "mindmap",
-        nodes: [
-          {
-            id: "topic-0",
-            label: "Root topic",
-            kind: "root",
-            metadata: { depth: 0, siblingIndex: 0 },
-          },
-          {
-            id: "topic-0-0",
-            label: "Child topic",
-            kind: "topic",
-            metadata: { depth: 1, siblingIndex: 0 },
-          },
-        ],
-        edges: [
-          {
-            id: "branch-0-0",
-            source: "topic-0",
-            target: "topic-0-0",
-            metadata: { depth: 1, siblingIndex: 0 },
-          },
-        ],
-        layout: { direction: "LR", edgeRouting: "curved" },
-      },
-      null,
-      2,
-    );
-  }
-
-  return JSON.stringify(
-    {
+    return JSON.stringify({
       id: "short-kebab-case-id",
       title: prompt.title,
-      type: "flowchart",
+      type: "mindmap",
       nodes: [
-        { id: "start-id", label: "Human label", kind: "start" },
-        { id: "decision-id", label: "Question?", kind: "decision" },
+        {
+          id: "topic-0",
+          label: "Root topic",
+          kind: "root",
+          metadata: { depth: 0, siblingIndex: 0 },
+        },
+        {
+          id: "topic-0-0",
+          label: "Child topic",
+          kind: "topic",
+          metadata: { depth: 1, siblingIndex: 0 },
+        },
       ],
       edges: [
         {
-          id: "edge-id",
-          source: "decision-id",
-          target: "target-id",
-          label: "yes",
+          id: "branch-0-0",
+          source: "topic-0",
+          target: "topic-0-0",
+          metadata: { depth: 1, siblingIndex: 0 },
         },
       ],
-      layout: { direction: "TB", edgeRouting: "orthogonal" },
-    },
-    null,
-    2,
-  );
+      layout: { direction: "LR", edgeRouting: "curved" },
+    });
+  }
+
+  return JSON.stringify({
+    id: "short-kebab-case-id",
+    title: prompt.title,
+    type: "flowchart",
+    nodes: [
+      { id: "start-id", label: "Human label", kind: "start" },
+      { id: "decision-id", label: "Question?", kind: "decision" },
+    ],
+    edges: [
+      {
+        id: "edge-id",
+        source: "decision-id",
+        target: "target-id",
+        label: "yes",
+      },
+    ],
+    layout: { direction: "TB", edgeRouting: "orthogonal" },
+  });
 }
 
 function requiredList(title: string, values: readonly string[]): string[] {

packages/diagram/generation/src/lib/public-api.test.ts

@@ -73,7 +73,7 @@ const expectedSystem = [
   "You are creating a Sketchi typed intermediate diagram.",
   "",
   "Flowchart IR rules:",
-  "- Return only JSON. Do not wrap the JSON in markdown.",
+  "- Return only compact, minified JSON on one line. Do not use markdown.",
   '- Use type "flowchart".',
   '- Every node must have id, label, and kind: "start", "process", "decision", or "end".',
   "- Use exactly one start node and at least one end node.",
@@ -97,28 +97,24 @@ const expectedUser = [
   "Use these required labels exactly unless the scenario explicitly asks for a clearer synonym.",
   "",
   "Expected JSON shape:",
-  JSON.stringify(
-    {
-      id: "short-kebab-case-id",
-      title: "Pharma batch disposition",
-      type: "flowchart",
-      nodes: [
-        { id: "start-id", label: "Human label", kind: "start" },
-        { id: "decision-id", label: "Question?", kind: "decision" },
-      ],
-      edges: [
-        {
-          id: "edge-id",
-          source: "decision-id",
-          target: "target-id",
-          label: "yes",
-        },
-      ],
-      layout: { direction: "TB", edgeRouting: "orthogonal" },
-    },
-    null,
-    2,
-  ),
+  JSON.stringify({
+    id: "short-kebab-case-id",
+    title: "Pharma batch disposition",
+    type: "flowchart",
+    nodes: [
+      { id: "start-id", label: "Human label", kind: "start" },
+      { id: "decision-id", label: "Question?", kind: "decision" },
+    ],
+    edges: [
+      {
+        id: "edge-id",
+        source: "decision-id",
+        target: "target-id",
+        label: "yes",
+      },
+    ],
+    layout: { direction: "TB", edgeRouting: "orthogonal" },
+  }),
 ].join("\n");
 const expectedGeminiBody = {
   contents: [{ role: "user", parts: [{ text: expectedUser }] }],
@@ -202,7 +198,7 @@ describe("diagram generation prompt mapping", () => {
 
     expect(messages.system).toContain('Use type "mindmap".');
     expect(messages.system).toContain("exactly one root node");
-    expect(messages.user).toContain('"type": "mindmap"');
+    expect(messages.user).toContain('"type":"mindmap"');
   });
 
   it("keeps the Gemini REST body byte-for-byte stable", () => {
@@ -216,6 +212,15 @@ describe("diagram generation prompt mapping", () => {
     ).toEqual(expectedGeminiBody);
   });
 
+  it("uses the expanded output budget unless the request overrides it", () => {
+    expect(
+      buildGeminiGenerateContentBody({
+        model: "google/gemini-3.1-flash-lite",
+        prompt,
+      }).generationConfig.maxOutputTokens,
+    ).toBe(16_384);
+  });
+
   it("normalizes Cloudflare Google model ids for provider-native calls", () => {
     expect(stripGoogleModelPrefix("google/gemini-3.1-flash-lite")).toBe(
       "gemini-3.1-flash-lite",
@@ -774,6 +779,45 @@ layer(malformedClientLayer)("malformed model responses", (it) => {
   );
 });
 
+const truncatedRun = vi.fn(async () =>
+  jsonResponse({
+    candidates: [
+      {
+        content: { role: "model", parts: [{ text: '{"type":"flowchart"' }] },
+        finishReason: "MAX_TOKENS",
+      },
+    ],
+  }),
+);
+const truncatedClientLayer = CloudflareGoogleAiStudioClientLive.pipe(
+  Layer.provide(
+    Layer.mergeAll(
+      Layer.succeed(CloudflareAiGatewayBinding, {
+        gateway: () => ({ getUrl: vi.fn(), run: truncatedRun }),
+      }),
+      configLayer,
+      retryPolicyLayer,
+    ),
+  ),
+);
+
+layer(truncatedClientLayer)("token-budget exhaustion", (it) => {
+  it.effect("returns an explicit output-truncated diagnostic", () =>
+    Effect.gen(function* () {
+      const client = yield* DiagramGenerationClient;
+      const candidate = yield* client.generate({
+        model: "google/gemini-3.1-flash-lite",
+        prompt,
+      });
+
+      assert.isUndefined(candidate.diagram);
+      expect(candidate.diagnostics).toContain(
+        "output_truncated: Gemini stopped at the maximum output-token budget; regenerate the complete diagram.",
+      );
+    }),
+  );
+});
+
 const substitutedClientLayer = Layer.succeed(DiagramGenerationClient, {
   provider: "fixture",
   generate: Effect.fn("diagramGeneration.test.generate")(function* () {

packages/diagram/scenarios/src/lib/prompt.test.ts

@@ -45,15 +45,17 @@ describe("scenario prompts", () => {
     expect(prompt.user).toContain("- QA Manager final review");
     expect(prompt.user).toContain("Required decision branch labels:");
     expect(prompt.user).toContain("- retest");
-    expect(prompt.user).toContain('"title": "Pharma batch disposition"');
+    expect(prompt.user).toContain('"title":"Pharma batch disposition"');
   });
 
   it("keeps a flattened prompt for stdin-based generator commands", () => {
     const flattened = buildScenarioPrompt(scenario);
 
     expect(flattened).toContain("System message:");
     expect(flattened).toContain("User message:");
-    expect(flattened).toContain("Return only JSON.");
+    expect(flattened).toContain(
+      "Return only compact, minified JSON on one line.",
+    );
     expect(flattened).toContain(scenario.prompt);
   });
 });