diff --git a/packages/diagram/generation/src/lib/cloudflare-google-ai-studio.ts b/packages/diagram/generation/src/lib/cloudflare-google-ai-studio.ts index 319eefe8..635db551 100644 --- a/packages/diagram/generation/src/lib/cloudflare-google-ai-studio.ts +++ b/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, diff --git a/packages/diagram/generation/src/lib/gemini.ts b/packages/diagram/generation/src/lib/gemini.ts index 20df09a2..10904a51 100644 --- a/packages/diagram/generation/src/lib/gemini.ts +++ b/packages/diagram/generation/src/lib/gemini.ts @@ -36,7 +36,7 @@ export class GeminiGenerateContentBody extends Schema.Class 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; +} diff --git a/packages/diagram/generation/src/lib/messages.ts b/packages/diagram/generation/src/lib/messages.ts index ac46cc41..cc048461 100644 --- a/packages/diagram/generation/src/lib/messages.ts +++ b/packages/diagram/generation/src/lib/messages.ts @@ -36,7 +36,7 @@ export class DiagramGenerationMessages extends Schema.Class { 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* () { diff --git a/packages/diagram/scenarios/src/lib/prompt.test.ts b/packages/diagram/scenarios/src/lib/prompt.test.ts index 79fc2704..040a0721 100644 --- a/packages/diagram/scenarios/src/lib/prompt.test.ts +++ b/packages/diagram/scenarios/src/lib/prompt.test.ts @@ -45,7 +45,7 @@ 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", () => { @@ -53,7 +53,9 @@ describe("scenario prompts", () => { 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); }); });