diff --git a/apps/eval-harness/src/lib/generate-scenario.test.ts b/apps/eval-harness/src/lib/generate-scenario.test.ts index a687fc13..e1e32fdf 100644 --- a/apps/eval-harness/src/lib/generate-scenario.test.ts +++ b/apps/eval-harness/src/lib/generate-scenario.test.ts @@ -307,6 +307,7 @@ const concurrencyLayer = Layer.mergeAll( }), Layer.succeed(DiagramGenerationPolicy, { concurrency: 1, + maxRepairAttempts: 1, maxRetries: 0, requestTimeoutMs: 1_000, retryDelayMs: 0, diff --git a/packages/diagram/generation/src/lib/candidates.ts b/packages/diagram/generation/src/lib/candidates.ts index e697abe8..c30a24c0 100644 --- a/packages/diagram/generation/src/lib/candidates.ts +++ b/packages/diagram/generation/src/lib/candidates.ts @@ -103,13 +103,28 @@ export function extractJsonObject(text: string): unknown { return JSON.parse(text); } catch { const firstBrace = text.indexOf("{"); - const lastBrace = text.lastIndexOf("}"); - - if (firstBrace === -1 || lastBrace === -1 || lastBrace <= firstBrace) { + if (firstBrace === -1) { throw new Error("Model output did not contain a JSON object."); } - - return JSON.parse(text.slice(firstBrace, lastBrace + 1)); + let depth = 0; + let escaped = false; + let inString = false; + for (let index = firstBrace; index < text.length; index += 1) { + const character = text[index]; + if (inString) { + if (escaped) escaped = false; + else if (character === "\\") escaped = true; + else if (character === '"') inString = false; + continue; + } + if (character === '"') inString = true; + else if (character === "{") depth += 1; + else if (character === "}") { + depth -= 1; + if (depth === 0) return JSON.parse(text.slice(firstBrace, index + 1)); + } + } + throw new Error("Model output did not contain one complete JSON object."); } } @@ -131,6 +146,25 @@ function withSketchiDiagramStyle(input: unknown): unknown { : input; } +function normalizeGeneratedFlowchartInput(input: unknown): unknown { + const styled = withSketchiDiagramStyle(input); + if (!isUnknownRecord(styled) || !Array.isArray(styled["edges"])) { + return styled; + } + return { + ...styled, + edges: styled["edges"].map((edge) => + isUnknownRecord(edge) && + typeof edge["label"] === "string" && + edge["label"].trim().length === 0 + ? Object.fromEntries( + Object.entries(edge).filter(([key]) => key !== "label"), + ) + : edge, + ), + }; +} + function firstString(values: readonly unknown[]): string | undefined { return values.find((value): value is string => typeof value === "string"); } @@ -160,7 +194,7 @@ export function responseErrorDiagnostic(raw: unknown): string | undefined { export function parseGeneratedFlowchart(text: string): FlowchartDiagram { return parseFlowchartDiagram( - withSketchiDiagramStyle(extractJsonObject(text)), + normalizeGeneratedFlowchartInput(extractJsonObject(text)), ); } @@ -175,7 +209,7 @@ export function parseGeneratedDiagram( } return parseMindmapDiagram(withSketchiDiagramStyle(extracted)); } - return parseFlowchartDiagram(withSketchiDiagramStyle(extracted)); + return parseFlowchartDiagram(normalizeGeneratedFlowchartInput(extracted)); } interface CandidateParseFailure { @@ -281,7 +315,7 @@ function parseCandidateDiagram(text: string): CandidateParseResult { const decoded = safeParseDiagramSchema( FlowchartDiagramSchema, - withSketchiDiagramStyle(extracted), + normalizeGeneratedFlowchartInput(extracted), ); if (!decoded.success) { const diagnostics = decoded.error.issues.map(schemaIssueDiagnostic); @@ -330,6 +364,135 @@ export function candidateFromText( }; } +const EXPLICIT_MINIMUM_PATTERN = + /\bat least\s+(\d+|one|two|three|four|five|six|seven|eight|nine|ten)\s+(?:(?:distinct|labeled)\s+)?(steps?|nodes?|topics?|decisions?(?:\s+nodes?)?)\b/giu; +const LOOP_REQUIREMENT_PATTERN = + /\b(?:feedback|review|resubmission|retry|revision|remediation|investigation)\s+loop\b|\bloop(?:s|ed|ing)?\s+(?:back|through|to)\b|\breturns?\s+to\b/iu; +const NUMBER_WORD_COUNTS: Readonly> = { + eight: 8, + five: 5, + four: 4, + nine: 9, + one: 1, + seven: 7, + six: 6, + ten: 10, + three: 3, + two: 2, +}; + +export interface ExplicitRequestMinimum { + readonly expectedCount: number; + readonly expectedUnit: "decision nodes" | "nodes" | "topics"; + readonly requestedCount: number; + readonly requestedUnit: string; +} + +function requestedCount(value: string): number | undefined { + const numeric = Number.parseInt(value, 10); + return Number.isInteger(numeric) + ? numeric + : NUMBER_WORD_COUNTS[value.toLowerCase()]; +} + +export function explicitRequestMinimums( + request: string, + diagramType: DiagramGenerationPrompt["type"], +): readonly ExplicitRequestMinimum[] { + return Array.from(request.matchAll(EXPLICIT_MINIMUM_PATTERN), (match) => { + const requested = match[1] ? requestedCount(match[1]) : undefined; + const unit = match[2]?.toLowerCase(); + if (!requested || !unit) return undefined; + + const decisionMinimum = /^decisions?(?:\s+nodes?)?$/.test(unit); + const applies = + (diagramType === "flowchart" && + (/^(?:steps?|nodes?)$/.test(unit) || decisionMinimum)) || + (diagramType === "mindmap" && /^topics?$/.test(unit)); + if (!applies) return undefined; + + const expectedUnit: ExplicitRequestMinimum["expectedUnit"] = decisionMinimum + ? "decision nodes" + : diagramType === "flowchart" + ? "nodes" + : "topics"; + return { + expectedCount: + diagramType === "flowchart" && !decisionMinimum + ? Math.min(requested, 24) + : requested, + expectedUnit, + requestedCount: requested, + requestedUnit: unit, + }; + }).filter( + (minimum): minimum is ExplicitRequestMinimum => minimum !== undefined, + ); +} + +function flowchartHasDirectedCycle(diagram: FlowchartDiagram): boolean { + const adjacency = new Map(); + for (const edge of diagram.edges) { + adjacency.set(edge.source, [ + ...(adjacency.get(edge.source) ?? []), + edge.target, + ]); + } + const pathExists = (start: string, destination: string): boolean => { + const pending = [start]; + const visited = new Set(); + while (pending.length > 0) { + const current = pending.pop(); + if (!current) continue; + if (current === destination) return true; + if (visited.has(current)) continue; + visited.add(current); + pending.push(...(adjacency.get(current) ?? [])); + } + return false; + }; + return diagram.edges.some((edge) => pathExists(edge.target, edge.source)); +} + +/** Turn a structurally valid result that misses explicit requirements into repair input. */ +export function enforceCandidateRequestRequirements( + candidate: DiagramGenerationCandidate, + request: DiagramGenerationRequest, +): DiagramGenerationCandidate { + const diagram = candidate.diagram; + if (!diagram || diagram.type !== request.prompt.type) { + return candidate; + } + const requestDiagnostics = [ + ...explicitRequestMinimums(request.prompt.request, request.prompt.type) + .map((minimum) => { + const actual = + minimum.expectedUnit === "decision nodes" + ? diagram.nodes.filter((node) => node.kind === "decision").length + : diagram.nodes.length; + if (actual >= minimum.expectedCount) return undefined; + + return `request_minimum_not_met: requested at least ${minimum.requestedCount} ${minimum.requestedUnit}, but the generated ${diagram.type} contained ${actual}. Hint: return a complete diagram with at least ${minimum.expectedCount} ${minimum.expectedUnit}.`; + }) + .filter((diagnostic): diagnostic is string => diagnostic !== undefined), + ...(diagram.type === "flowchart" && + LOOP_REQUIREMENT_PATTERN.test(request.prompt.request) && + !flowchartHasDirectedCycle(diagram) + ? [ + "request_loop_not_met: prompt requires a retry or loop, but the generated flowchart contains no directed cycle. Hint: add a real back-edge from the loop path to the intended process or decision node, never the start node.", + ] + : []), + ]; + if (requestDiagnostics.length === 0) return candidate; + + const { diagram: _diagram, ...withoutDiagram } = candidate; + return { + ...withoutDiagram, + diagnostics: [...candidate.diagnostics, ...requestDiagnostics], + error: "Generated diagram did not satisfy explicit request requirements.", + }; +} + export function summarizeGenerationCandidate( candidate: DiagramGenerationCandidate, ): DiagramGenerationCandidateSummary { diff --git a/packages/diagram/generation/src/lib/client.ts b/packages/diagram/generation/src/lib/client.ts index 5015f6e8..c624f795 100644 --- a/packages/diagram/generation/src/lib/client.ts +++ b/packages/diagram/generation/src/lib/client.ts @@ -21,6 +21,7 @@ export class DiagramGenerationPolicyConfig extends Schema.Class ai.gateway(config.gatewayId), - catch: (cause) => - DiagramGenerationTransportError.make({ - cause, - message: errorMessage( + (attemptRequest) => + Effect.gen(function* () { + const gateway = yield* Effect.try({ + try: () => ai.gateway(config.gatewayId), + catch: (cause) => + DiagramGenerationTransportError.make({ cause, - "AI Gateway could not be initialized.", - ), - operation: "ai.gateway", - provider: "cloudflare-google-ai-studio", - retryable: false, - }), - }); - return runGatewayAttempt(gateway, config.collectLog, request); - }), + message: errorMessage( + cause, + "AI Gateway could not be initialized.", + ), + operation: "ai.gateway", + provider: "cloudflare-google-ai-studio", + retryable: false, + }), + }); + return runGatewayAttempt( + gateway, + config.collectLog, + attemptRequest, + ); + }), request, "cloudflare-google-ai-studio", policy, diff --git a/packages/diagram/generation/src/lib/policy.ts b/packages/diagram/generation/src/lib/policy.ts index 128d5c6d..12a52645 100644 --- a/packages/diagram/generation/src/lib/policy.ts +++ b/packages/diagram/generation/src/lib/policy.ts @@ -1,15 +1,17 @@ import { recordMetric, withTelemetryCorrelation } from "@sketchi/observability"; -import { Clock, Effect, Metric, Ref, Schedule } from "effect"; +import { Clock, Effect, Metric, Ref, Result, Schedule } from "effect"; -import type { - DiagramGenerationCandidate, - DiagramGenerationProviderId, - DiagramGenerationRequest, +import { + type DiagramGenerationCandidate, + type DiagramGenerationProviderId, + type DiagramGenerationRequest, + explicitRequestMinimums, } from "./candidates.js"; import type { DiagramGenerationPolicyConfig } from "./client.js"; import { type DiagramGenerationError, DiagramGenerationTimeoutError, + generationErrorToCandidate, isRetryableGenerationError, } from "./errors.js"; @@ -25,6 +27,10 @@ const generationRetries = Metric.counter("sketchi_generation_retries", { description: "Diagram generation retry attempts", incremental: true, }); +const generationRepairs = Metric.counter("sketchi_generation_repairs", { + description: "Diagram generation semantic repair attempts by outcome", + incremental: true, +}); const generationFailures = Metric.counter("sketchi_generation_failures", { description: "Diagram generation terminal failures", incremental: true, @@ -43,7 +49,9 @@ const generationDuration = Metric.histogram("sketchi_generation_duration_ms", { export const runDiagramGenerationWithPolicy = Effect.fn( "diagramGeneration.runWithPolicy", )(function* ( - prepareAttempt: Effect.Effect< + prepareAttempt: ( + request: DiagramGenerationRequest, + ) => Effect.Effect< Effect.Effect, DiagramGenerationError >, @@ -52,57 +60,133 @@ export const runDiagramGenerationWithPolicy = Effect.fn( policy: DiagramGenerationPolicyConfig, ) { const startedAt = yield* Clock.currentTimeMillis; + const executeModelCall = Effect.fn("diagramGeneration.executeModelCall")( + function* (callRequest: DiagramGenerationRequest) { + const attempt = yield* prepareAttempt(callRequest); + const attemptRef = yield* Ref.make(0); + const previousErrorTagRef = yield* Ref.make("initial"); + const measuredAttempt = Effect.gen(function* () { + const attemptNumber = yield* Ref.updateAndGet( + attemptRef, + (value) => value + 1, + ); + yield* recordMetric(generationAttempts, 1, { + operation: "generate", + provider, + }); + if (attemptNumber > 1) { + const previousErrorTag = yield* Ref.get(previousErrorTagRef); + yield* recordMetric(generationRetries, 1, { + operation: "generate", + provider, + retryKind: "transient", + }); + yield* Effect.logWarning("Retrying diagram generation", { + attempt: attemptNumber, + error_tag: previousErrorTag, + operation: "generate", + provider, + retry_kind: "transient", + }); + } + return yield* attempt.pipe( + Effect.annotateSpans({ attempt: attemptNumber }), + Effect.timeoutOrElse({ + duration: policy.requestTimeoutMs, + orElse: () => + Effect.fail( + DiagramGenerationTimeoutError.make({ + message: `Generation timed out after ${policy.requestTimeoutMs} ms.`, + provider, + timeoutMs: policy.requestTimeoutMs, + }), + ), + }), + Effect.tapError((error) => Ref.set(previousErrorTagRef, error._tag)), + ); + }); + return yield* measuredAttempt.pipe( + Effect.retry({ + schedule: Schedule.exponential(policy.retryDelayMs), + times: policy.maxRetries, + while: isRetryableGenerationError, + }), + ); + }, + ); + const operation = Effect.gen(function* () { - const attempt = yield* prepareAttempt; - const attemptRef = yield* Ref.make(0); - const previousErrorTagRef = yield* Ref.make("initial"); - const measuredAttempt = Effect.gen(function* () { - const attemptNumber = yield* Ref.updateAndGet( - attemptRef, - (value) => value + 1, + const originalCandidate = yield* executeModelCall(request); + let latestCandidate = originalCandidate; + let candidate = originalCandidate; + let diagnostics = [...originalCandidate.diagnostics]; + for ( + let repairAttempt = 1; + repairAttempt <= policy.maxRepairAttempts && !latestCandidate.diagram; + repairAttempt += 1 + ) { + const truncated = latestCandidate.diagnostics.some((diagnostic) => + diagnostic.startsWith("output_truncated:"), ); - yield* recordMetric(generationAttempts, 1, { + yield* recordMetric(generationRepairs, 1, { operation: "generate", + outcome: "attempted", provider, }); - if (attemptNumber > 1) { - const previousErrorTag = yield* Ref.get(previousErrorTagRef); - yield* recordMetric(generationRetries, 1, { - operation: "generate", - provider, - retryKind: "transient", - }); - yield* Effect.logWarning("Retrying diagram generation", { - attempt: attemptNumber, - error_tag: previousErrorTag, + yield* Effect.logWarning("Repairing invalid diagram generation", { + operation: "generate", + provider, + repair_attempt: repairAttempt, + repair_kind: truncated ? "regenerate_truncated" : "repair_invalid", + }); + const repairRequest = truncated + ? { ...request, cacheMode: "fresh" as const } + : buildRepairRequest( + request, + { ...latestCandidate, diagnostics }, + repairAttempt, + ); + const attemptedDiagnostic = `repair_attempted: ${truncated ? "regenerated a truncated response" : "requested a corrected response"} (attempt ${repairAttempt}).`; + const repairResult = yield* Effect.result( + executeModelCall(repairRequest), + ); + if (Result.isFailure(repairResult)) { + const failedRepairCandidate = generationErrorToCandidate( + repairResult.failure, + repairRequest, + ); + diagnostics = [ + ...diagnostics, + attemptedDiagnostic, + ...failedRepairCandidate.diagnostics, + `repair_failed: semantic repair attempt ${repairAttempt} failed.`, + ]; + yield* recordMetric(generationRepairs, 1, { operation: "generate", + outcome: "failed", provider, - retry_kind: "transient", }); + candidate = { ...originalCandidate, diagnostics }; + break; } - return yield* attempt.pipe( - Effect.annotateSpans({ attempt: attemptNumber }), - Effect.timeoutOrElse({ - duration: policy.requestTimeoutMs, - orElse: () => - Effect.fail( - DiagramGenerationTimeoutError.make({ - message: `Generation timed out after ${policy.requestTimeoutMs} ms.`, - provider, - timeoutMs: policy.requestTimeoutMs, - }), - ), - }), - Effect.tapError((error) => Ref.set(previousErrorTagRef, error._tag)), - ); - }); - const candidate = yield* measuredAttempt.pipe( - Effect.retry({ - schedule: Schedule.exponential(policy.retryDelayMs), - times: policy.maxRetries, - while: isRetryableGenerationError, - }), - ); + const repaired = repairResult.success; + const outcome = repaired.diagram ? "succeeded" : "failed"; + yield* recordMetric(generationRepairs, 1, { + operation: "generate", + outcome, + provider, + }); + diagnostics = [ + ...diagnostics, + attemptedDiagnostic, + ...repaired.diagnostics, + `repair_${outcome}: semantic repair attempt ${repairAttempt} ${outcome}.`, + ]; + latestCandidate = repaired; + candidate = repaired.diagram + ? { ...repaired, diagnostics } + : { ...originalCandidate, diagnostics }; + } const finishedAt = yield* Clock.currentTimeMillis; return { ...candidate, @@ -113,12 +197,12 @@ export const runDiagramGenerationWithPolicy = Effect.fn( Effect.all([ recordMetric(generationRequests, 1, { operation: "generate", - outcome: "success", + outcome: candidate.diagram ? "success" : "invalid", provider, }), recordMetric(generationDuration, candidate.durationMs, { operation: "generate", - outcome: "success", + outcome: candidate.diagram ? "success" : "invalid", provider, }), ]), @@ -159,3 +243,81 @@ export const runDiagramGenerationWithPolicy = Effect.fn( scenarioId: request.prompt.id, }); }); + +function buildRepairRequest( + request: DiagramGenerationRequest, + candidate: DiagramGenerationCandidate, + repairAttempt: number, +): DiagramGenerationRequest { + const priorityDiagnostics = candidate.diagnostics.filter( + (diagnostic) => + diagnostic.startsWith("flowchart.start_has_incoming:") || + diagnostic.startsWith("flowchart.self_loop:"), + ); + const priorityCorrections = [ + ...(priorityDiagnostics.some((diagnostic) => + diagnostic.startsWith("flowchart.start_has_incoming:"), + ) + ? [ + "Required correction for start-node incoming edges: reroute each offending loop-back edge to the first process node after start. Never target the start node; start nodes have no incoming edges.", + ] + : []), + ...(priorityDiagnostics.some((diagnostic) => + diagnostic.startsWith("flowchart.self_loop:"), + ) + ? [ + "Required correction for self-loops: reroute each offending edge to an earlier distinct process or decision node. Never keep the same node as both source and target; model every retry or re-check as a decision branch returning to that earlier distinct node.", + ] + : []), + ]; + const priorityGuidance = + priorityDiagnostics.length > 0 + ? [ + "Priority validator issue and hint:", + ...priorityDiagnostics.map((diagnostic) => `- ${diagnostic}`), + ...priorityCorrections, + "", + ] + : []; + const parsedMinimums = explicitRequestMinimums( + request.prompt.request, + request.prompt.type, + ); + const originalHardRequirements = [ + "Original hard requirements (all remain mandatory):", + `- Diagram type: ${request.prompt.type}.`, + `- Original scenario: ${request.prompt.request}`, + ...parsedMinimums.map( + (minimum) => + `- Parsed minimum: at least ${minimum.expectedCount} ${minimum.expectedUnit} (from the original request for at least ${minimum.requestedCount} ${minimum.requestedUnit}).`, + ), + ...request.prompt.requiredNodeLabels.map( + (label) => `- Required node label: ${label}`, + ), + ...request.prompt.requiredBranchLabels.map( + (label) => `- Required decision branch label: ${label}`, + ), + "", + ]; + return { + ...request, + cacheMode: "fresh", + prompt: { + ...request.prompt, + request: [ + request.prompt.request, + "", + `Repair attempt ${repairAttempt}: return a complete corrected diagram that satisfies every validator diagnostic below.`, + "PRESERVE all existing nodes and labels except the minimal edit needed to fix the listed issues. Do not compact, summarize, remove, combine, rename, or relabel unaffected content.", + "", + ...originalHardRequirements, + ...priorityGuidance, + "Invalid model output:", + candidate.text, + "", + "Validator diagnostics:", + ...candidate.diagnostics.map((diagnostic) => `- ${diagnostic}`), + ].join("\n"), + }, + }; +} diff --git a/packages/diagram/generation/src/lib/public-api.test.ts b/packages/diagram/generation/src/lib/public-api.test.ts index b7b12898..d7e2472a 100644 --- a/packages/diagram/generation/src/lib/public-api.test.ts +++ b/packages/diagram/generation/src/lib/public-api.test.ts @@ -9,8 +9,17 @@ import { import { Cause, Effect, Exit, Fiber, Layer, Schema } from "effect"; import { TestClock } from "effect/testing"; -import { candidateFromText, responseErrorDiagnostic } from "./candidates.js"; -import { DiagramGenerationClient, DiagramGenerationPolicy } from "./client.js"; +import { + candidateFromText, + enforceCandidateRequestRequirements, + extractJsonObject, + responseErrorDiagnostic, +} from "./candidates.js"; +import { + diagramGenerationPolicyDefaults, + DiagramGenerationClient, + DiagramGenerationPolicy, +} from "./client.js"; import { type CloudflareAiGateway, CloudflareAiGatewayBinding, @@ -53,13 +62,40 @@ const expectedDiagram = { type: "flowchart", nodes: [ { id: "received", label: "Batch received", kind: "start" }, + { + id: "final-review", + label: "QA Manager final review", + kind: "decision", + }, + { + id: "investigation", + label: "Investigate retesting", + kind: "process", + }, { id: "packaging", label: "Send to packaging", kind: "end" }, ], edges: [ { - id: "received-to-packaging", + id: "received-to-final-review", source: "received", + target: "final-review", + }, + { + id: "final-review-to-packaging", + source: "final-review", target: "packaging", + label: "yes", + }, + { + id: "final-review-to-investigation", + source: "final-review", + target: "investigation", + label: "retest", + }, + { + id: "investigation-to-final-review", + source: "investigation", + target: "final-review", }, ], layout: { direction: "TB", edgeRouting: "orthogonal" }, @@ -161,12 +197,14 @@ const configLayer = Layer.succeed(CloudflareGoogleAiStudioConfig, { }); const retryPolicyLayer = Layer.succeed(DiagramGenerationPolicy, { concurrency: 2, + maxRepairAttempts: 1, maxRetries: 2, requestTimeoutMs: 1_000, retryDelayMs: 100, }); const cancellationPolicyLayer = Layer.succeed(DiagramGenerationPolicy, { concurrency: 2, + maxRepairAttempts: 1, maxRetries: 1, requestTimeoutMs: 500, retryDelayMs: 100, @@ -253,6 +291,7 @@ describe("diagram generation prompt mapping", () => { prompt, }).generationConfig.maxOutputTokens, ).toBe(16_384); + expect(diagramGenerationPolicyDefaults.maxRepairAttempts).toBe(2); }); it("normalizes Cloudflare Google model ids for provider-native calls", () => { @@ -266,6 +305,14 @@ describe("diagram generation prompt mapping", () => { }); describe("pure candidate behavior", () => { + it("extracts the first complete JSON object without swallowing trailing output", () => { + expect( + extractJsonObject( + 'Result: {"label":"Review {draft}","nested":{"ok":true}}\n{"duplicate":true}', + ), + ).toEqual({ label: "Review {draft}", nested: { ok: true } }); + }); + it.effect.prop( "derives a valid flat mindmap from every generated nested hierarchy", { tree: GeneratedMindmapTree }, @@ -318,6 +365,100 @@ describe("pure candidate behavior", () => { expect(candidate.diagram?.style).toEqual(expectedDiagram.style); }); + it("normalizes empty optional edge labels before schema validation", () => { + const candidate = candidateFromText({ + model: "fixture", + provider: "fixture", + text: JSON.stringify({ + ...expectedDiagram, + edges: expectedDiagram.edges.map((edge, index) => + index === 0 ? { ...edge, label: "" } : edge, + ), + }), + }); + + expect(candidate.error).toBeUndefined(); + expect(candidate.diagram?.edges[0]?.label).toBeUndefined(); + }); + + it("turns an explicitly undersized valid result into semantic repair input", () => { + const candidate = enforceCandidateRequestRequirements( + candidateFromText({ + model: "fixture", + provider: "fixture", + text: expectedText, + }), + { + model: "fixture", + prompt: { + ...prompt, + request: "Create a return flow with at least 18 distinct steps.", + }, + }, + ); + + expect(candidate.diagram).toBeUndefined(); + expect(candidate.diagnostics).toContain( + "request_minimum_not_met: requested at least 18 steps, but the generated flowchart contained 4. Hint: return a complete diagram with at least 18 nodes.", + ); + }); + + it("turns missing explicit decision counts into semantic repair input", () => { + const candidate = enforceCandidateRequestRequirements( + candidateFromText({ + model: "fixture", + provider: "fixture", + text: expectedText, + }), + { + model: "fixture", + prompt: { + ...prompt, + request: "Use at least five labeled decision nodes.", + }, + }, + ); + + expect(candidate.diagram).toBeUndefined(); + expect(candidate.diagnostics).toContain( + "request_minimum_not_met: requested at least 5 decision nodes, but the generated flowchart contained 1. Hint: return a complete diagram with at least 5 decision nodes.", + ); + }); + + it("turns a named loop without a directed cycle into semantic repair input", () => { + const candidate = enforceCandidateRequestRequirements( + candidateFromText({ + model: "fixture", + provider: "fixture", + text: JSON.stringify({ + ...expectedDiagram, + nodes: expectedDiagram.nodes.filter( + (node) => node.id === "received" || node.id === "packaging", + ), + edges: [ + { + id: "received-to-packaging", + source: "received", + target: "packaging", + }, + ], + }), + }), + { + model: "fixture", + prompt: { + ...prompt, + request: "Add a fraud review loop before closure.", + }, + }, + ); + + expect(candidate.diagram).toBeUndefined(); + expect(candidate.diagnostics).toContain( + "request_loop_not_met: prompt requires a retry or loop, but the generated flowchart contains no directed cycle. Hint: add a real back-edge from the loop path to the intended process or decision node, never the start node.", + ); + }); + it("parses validated mindmap candidates through diagram-core", () => { const candidate = candidateFromText({ model: "fixture", @@ -491,6 +632,7 @@ layer(successfulClientLayer)("Cloudflare Google AI Studio live layer", (it) => { expect.objectContaining({ headers: expect.objectContaining({ "Cache-Control": "no-store", + "cf-aig-skip-cache": "true", Pragma: "no-cache", }), }), @@ -505,6 +647,685 @@ layer(successfulClientLayer)("Cloudflare Google AI Studio live layer", (it) => { ); }); +const invalidDiagramText = JSON.stringify({ + id: "invalid-retry", + title: "Invalid retry", + type: "flowchart", + nodes: [ + { id: "start", kind: "start", label: "Start" }, + { id: "retry", kind: "decision", label: "Retry?" }, + { id: "end", kind: "end", label: "Done" }, + ], + edges: [ + { id: "start-retry", source: "start", target: "retry" }, + { id: "retry-end", source: "retry", target: "end" }, + ], + layout: { direction: "TB", edgeRouting: "orthogonal" }, +}); + +const startIncomingDiagramText = JSON.stringify({ + id: "expense-resubmission", + title: "Expense resubmission", + type: "flowchart", + nodes: [ + { id: "start", kind: "start", label: "Start" }, + { id: "submission", kind: "process", label: "Submission" }, + { id: "manager", kind: "decision", label: "Manager approves?" }, + { id: "finance", kind: "decision", label: "Finance approves?" }, + { id: "rejection", kind: "process", label: "Rejection" }, + { id: "resubmit", kind: "process", label: "Resubmit" }, + { id: "reimbursement", kind: "end", label: "Reimbursement" }, + ], + edges: [ + { id: "start-submission", source: "start", target: "submission" }, + { id: "submission-manager", source: "submission", target: "manager" }, + { + id: "manager-finance", + source: "manager", + target: "finance", + label: "approved", + }, + { + id: "manager-rejection", + source: "manager", + target: "rejection", + label: "rejected", + }, + { + id: "finance-reimbursement", + source: "finance", + target: "reimbursement", + label: "approved", + }, + { + id: "finance-rejection", + source: "finance", + target: "rejection", + label: "rejected", + }, + { id: "rejection-resubmit", source: "rejection", target: "resubmit" }, + { id: "resubmit-start", source: "resubmit", target: "start" }, + ], + layout: { direction: "TB", edgeRouting: "orthogonal" }, +}); + +const selfLoopDiagramText = JSON.stringify({ + id: "returns-fraud-review", + title: "Returns fraud review", + type: "flowchart", + nodes: [ + { id: "start", kind: "start", label: "Start" }, + { id: "return", kind: "process", label: "Receive return" }, + { id: "fraud", kind: "decision", label: "Fraud check passes?" }, + { id: "done", kind: "end", label: "Refund complete" }, + ], + edges: [ + { id: "start-return", source: "start", target: "return" }, + { id: "return-fraud", source: "return", target: "fraud" }, + { + id: "fraud-self", + source: "fraud", + target: "fraud", + label: "re-check", + }, + { id: "fraud-done", source: "fraud", target: "done", label: "clear" }, + ], + layout: { direction: "TB", edgeRouting: "orthogonal" }, +}); + +function returnsDiagramText(nodeCount: number, selfLoop: boolean): string { + const processCount = nodeCount - 3; + const processNodes = Array.from({ length: processCount }, (_, index) => ({ + id: `step-${index + 1}`, + kind: "process", + label: `Return step ${index + 1}`, + })); + const processEdges = processNodes.map((node, index) => ({ + id: `step-edge-${index + 1}`, + source: node.id, + target: processNodes[index + 1]?.id ?? "fraud-check", + })); + const lastProcessId = processNodes.at(-1)?.id ?? "fraud-review"; + return JSON.stringify({ + id: "returns-minimum-repair", + title: "Ecommerce returns", + type: "flowchart", + nodes: [ + { id: "start", kind: "start", label: "Return initiated" }, + ...processNodes, + { id: "fraud-check", kind: "decision", label: "Fraud check?" }, + { id: "done", kind: "end", label: "Return closed" }, + ], + edges: [ + { + id: "start-step", + source: "start", + target: processNodes[0]?.id ?? "fraud-check", + }, + ...processEdges, + { + id: "fraud-recheck", + source: "fraud-check", + target: selfLoop ? "fraud-check" : lastProcessId, + label: "re-check", + }, + { + id: "fraud-clear", + source: "fraud-check", + target: "done", + label: "clear", + }, + ], + layout: { direction: "TB", edgeRouting: "orthogonal" }, + }); +} + +function geminiTextResponse(text: string): Response { + return jsonResponse({ + candidates: [ + { content: { role: "model", parts: [{ text }] }, finishReason: "STOP" }, + ], + }); +} + +let repairedRunCalls = 0; +const repairedRun = vi.fn(async () => { + repairedRunCalls += 1; + return repairedRunCalls === 1 + ? geminiTextResponse(invalidDiagramText) + : jsonResponse(geminiResponse); +}); +const repairedClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: repairedRun }), + }), + configLayer, + retryPolicyLayer, + ), + ), +); + +layer(repairedClientLayer)("semantic repair policy", (it) => { + it.effect("repairs one invalid response and records repair outcomes", () => { + const { probe, sink } = makeTelemetryTestSink(); + const telemetryLayer = makeWorkersTelemetryLayer({ + resource: { serviceName: "sketchi-generation-repair-test" }, + sink, + }); + return Effect.gen(function* () { + repairedRunCalls = 0; + repairedRun.mockClear(); + const client = yield* DiagramGenerationClient; + const candidate = yield* client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt, + }); + + assert.strictEqual(repairedRun.mock.calls.length, 2); + assert.strictEqual(candidate.diagram?.id, expectedDiagram.id); + assert.isTrue( + candidate.diagnostics.some((diagnostic) => + diagnostic.startsWith("flowchart.underbranched_decision:"), + ), + ); + assert.isTrue( + candidate.diagnostics.includes( + "repair_succeeded: semantic repair attempt 1 succeeded.", + ), + ); + expect(repairedRun.mock.calls[1]?.[0]).toEqual( + expect.objectContaining({ + headers: expect.objectContaining({ + "Cache-Control": "no-store", + "cf-aig-skip-cache": "true", + }), + query: expect.objectContaining({ + contents: [ + expect.objectContaining({ + parts: [ + expect.objectContaining({ + text: expect.stringContaining("Invalid model output:"), + }), + ], + }), + ], + }), + }), + ); + const repairs = probe.events.filter( + (event): event is TelemetryMetricEvent => + event.event === "effect.metric" && + event.metric === "sketchi_generation_repairs", + ); + assert.deepStrictEqual( + repairs.map((event) => event.attributes["outcome"]), + ["attempted", "succeeded"], + ); + }).pipe(Effect.provide(telemetryLayer)); + }); +}); + +let startIncomingRepairCalls = 0; +const startIncomingRepairRun = vi.fn(async () => { + startIncomingRepairCalls += 1; + return startIncomingRepairCalls === 1 + ? geminiTextResponse(startIncomingDiagramText) + : jsonResponse(geminiResponse); +}); +const startIncomingRepairClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: startIncomingRepairRun }), + }), + configLayer, + retryPolicyLayer, + ), + ), +); + +layer(startIncomingRepairClientLayer)("start edge semantic repair", (it) => { + it.effect("prioritizes rerouting loop-back edges away from start", () => + Effect.gen(function* () { + startIncomingRepairCalls = 0; + startIncomingRepairRun.mockClear(); + const client = yield* DiagramGenerationClient; + const candidate = yield* client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt, + }); + + assert.strictEqual(startIncomingRepairRun.mock.calls.length, 2); + assert.strictEqual(candidate.diagram?.id, expectedDiagram.id); + const repairBody = startIncomingRepairRun.mock.calls[1]?.[0]; + expect(repairBody).toEqual( + expect.objectContaining({ + query: expect.objectContaining({ + contents: [ + expect.objectContaining({ + parts: [ + expect.objectContaining({ + text: expect.stringMatching( + /Priority validator issue and hint:\n- flowchart\.start_has_incoming: Start node "start" cannot have incoming edges\. Hint: Route the start node only to later nodes\.[\s\S]*reroute each offending loop-back edge to the first process node after start\. Never target the start node; start nodes have no incoming edges\./u, + ), + }), + ], + }), + ], + }), + }), + ); + }), + ); +}); + +let selfLoopRepairCalls = 0; +const selfLoopRepairRun = vi.fn(async () => { + selfLoopRepairCalls += 1; + return selfLoopRepairCalls === 1 + ? geminiTextResponse(selfLoopDiagramText) + : jsonResponse(geminiResponse); +}); +const selfLoopRepairClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: selfLoopRepairRun }), + }), + configLayer, + retryPolicyLayer, + ), + ), +); + +layer(selfLoopRepairClientLayer)("self-loop semantic repair", (it) => { + it.effect( + "prioritizes rerouting self-loops to an earlier distinct node", + () => + Effect.gen(function* () { + selfLoopRepairCalls = 0; + selfLoopRepairRun.mockClear(); + const client = yield* DiagramGenerationClient; + const candidate = yield* client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt, + }); + + assert.strictEqual(selfLoopRepairRun.mock.calls.length, 2); + assert.strictEqual(candidate.diagram?.id, expectedDiagram.id); + const repairBody = selfLoopRepairRun.mock.calls[1]?.[0]; + expect(repairBody).toEqual( + expect.objectContaining({ + query: expect.objectContaining({ + contents: [ + expect.objectContaining({ + parts: [ + expect.objectContaining({ + text: expect.stringMatching( + /Priority validator issue and hint:\n- flowchart\.self_loop: Edge "fraud-self" connects node "fraud" to itself\. Hint: Connect the edge to a different target node\.[\s\S]*reroute each offending edge to an earlier distinct process or decision node\.[\s\S]*model every retry or re-check as a decision branch returning to that earlier distinct node\./u, + ), + }), + ], + }), + ], + }), + }), + ); + }), + ); +}); + +const returnsMinimumPrompt: DiagramGenerationPrompt = { + id: "returns-minimum-repair", + request: + "Create an ecommerce returns flowchart with at least 18 distinct steps and a fraud review loop.", + requiredBranchLabels: ["re-check", "clear"], + requiredNodeLabels: ["Fraud check?", "Return closed"], + title: "Ecommerce returns", + type: "flowchart", +}; +let compactingRepairCalls = 0; +const compactingRepairRun = vi.fn(async () => { + compactingRepairCalls += 1; + if (compactingRepairCalls === 1) { + return geminiTextResponse(returnsDiagramText(18, true)); + } + return geminiTextResponse( + compactingRepairCalls === 2 + ? returnsDiagramText(16, false) + : returnsDiagramText(18, false), + ); +}); +const compactingRepairClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: compactingRepairRun }), + }), + configLayer, + Layer.succeed(DiagramGenerationPolicy, diagramGenerationPolicyDefaults), + ), + ), +); + +layer(compactingRepairClientLayer)("repair-introduced violations", (it) => { + it.effect( + "preserves original hard requirements and repairs a new violation class", + () => + Effect.gen(function* () { + compactingRepairCalls = 0; + compactingRepairRun.mockClear(); + const client = yield* DiagramGenerationClient; + const candidate = yield* client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt: returnsMinimumPrompt, + }); + + assert.strictEqual(compactingRepairRun.mock.calls.length, 3); + assert.strictEqual(candidate.diagram?.nodes.length, 18); + expect(candidate.diagnostics).toEqual( + expect.arrayContaining([ + expect.stringContaining("flowchart.self_loop:"), + expect.stringContaining( + "request_minimum_not_met: requested at least 18 steps", + ), + "repair_failed: semantic repair attempt 1 failed.", + "repair_succeeded: semantic repair attempt 2 succeeded.", + ]), + ); + + const firstRepairBody = JSON.stringify( + compactingRepairRun.mock.calls[1]?.[0].query, + ); + expect(firstRepairBody).toContain( + "PRESERVE all existing nodes and labels except the minimal edit needed", + ); + expect(firstRepairBody).toContain( + "Original hard requirements (all remain mandatory):", + ); + expect(firstRepairBody).toContain( + "Parsed minimum: at least 18 nodes (from the original request for at least 18 steps).", + ); + expect(firstRepairBody).toContain( + "Required decision branch label: re-check", + ); + expect(firstRepairBody).toContain("Required node label: Fraud check?"); + + const secondRepairBody = JSON.stringify( + compactingRepairRun.mock.calls[2]?.[0].query, + ); + expect(secondRepairBody).toContain( + "request_minimum_not_met: requested at least 18 steps", + ); + expect(secondRepairBody).toContain("Parsed minimum: at least 18 nodes"); + }), + ); +}); + +let failedRepairCalls = 0; +const failedRepairRun = vi.fn(async () => { + failedRepairCalls += 1; + return geminiTextResponse( + failedRepairCalls === 1 ? invalidDiagramText : '{"type":"flowchart"', + ); +}); +const failedRepairClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: failedRepairRun }), + }), + configLayer, + retryPolicyLayer, + ), + ), +); + +layer(failedRepairClientLayer)("failed semantic repair", (it) => { + it.effect("returns composed diagnostics after the bounded repair fails", () => + Effect.gen(function* () { + failedRepairCalls = 0; + failedRepairRun.mockClear(); + const client = yield* DiagramGenerationClient; + const candidate = yield* client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt, + }); + + assert.strictEqual(failedRepairRun.mock.calls.length, 2); + assert.isUndefined(candidate.diagram); + assert.strictEqual(candidate.text, invalidDiagramText); + assert.notStrictEqual(candidate.text, '{"type":"flowchart"'); + assert.match(candidate.error ?? "", /Decision node/u); + assert.isTrue( + candidate.diagnostics.some((diagnostic) => + diagnostic.startsWith("flowchart.underbranched_decision:"), + ), + ); + assert.isTrue( + candidate.diagnostics.some((diagnostic) => + diagnostic.startsWith("json_parse_error:"), + ), + ); + assert.isTrue( + candidate.diagnostics.includes( + "repair_failed: semantic repair attempt 1 failed.", + ), + ); + }), + ); +}); + +let exhaustedRepairCalls = 0; +const exhaustedRepairRun = vi.fn(async () => { + exhaustedRepairCalls += 1; + return exhaustedRepairCalls === 1 + ? geminiTextResponse(invalidDiagramText) + : jsonResponse( + { error: { message: "repair provider temporarily unavailable" } }, + { status: 503 }, + ); +}); +const exhaustedRepairClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: exhaustedRepairRun }), + }), + configLayer, + retryPolicyLayer, + ), + ), +); + +layer(exhaustedRepairClientLayer)("failed repair transport", (it) => { + it.effect( + "returns the original malformed candidate and records a failed repair", + () => { + const { probe, sink } = makeTelemetryTestSink(); + const telemetryLayer = makeWorkersTelemetryLayer({ + resource: { serviceName: "sketchi-generation-repair-failure-test" }, + sink, + }); + return Effect.gen(function* () { + exhaustedRepairCalls = 0; + exhaustedRepairRun.mockClear(); + const client = yield* DiagramGenerationClient; + const fiber = yield* Effect.forkChild( + client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt, + }), + ); + yield* TestClock.adjust("1 second"); + const candidate = yield* Fiber.join(fiber); + + assert.strictEqual(exhaustedRepairRun.mock.calls.length, 4); + assert.isUndefined(candidate.diagram); + assert.strictEqual(candidate.text, invalidDiagramText); + assert.match(candidate.error ?? "", /Decision node/u); + assert.isTrue( + candidate.diagnostics.some((diagnostic) => + diagnostic.includes("repair provider temporarily unavailable"), + ), + ); + assert.isTrue( + candidate.diagnostics.includes( + "repair_failed: semantic repair attempt 1 failed.", + ), + ); + const repairs = probe.events.filter( + (event): event is TelemetryMetricEvent => + event.event === "effect.metric" && + event.metric === "sketchi_generation_repairs", + ); + assert.deepStrictEqual( + repairs.map((event) => event.attributes["outcome"]), + ["attempted", "failed"], + ); + assert.strictEqual( + probe.events.filter( + (event): event is TelemetryMetricEvent => + event.event === "effect.metric" && + event.metric === "sketchi_generation_failures", + ).length, + 0, + ); + }).pipe(Effect.provide(telemetryLayer)); + }, + ); +}); + +let interruptedRepairCalls = 0; +const interruptedRepairSignals: AbortSignal[] = []; +const interruptedRepairRun = vi.fn( + async (_data, options) => { + interruptedRepairCalls += 1; + if (interruptedRepairCalls === 1) { + return geminiTextResponse(invalidDiagramText); + } + const signal = options?.signal; + if (!signal) { + throw new Error("Repair request omitted AbortSignal."); + } + interruptedRepairSignals.push(signal); + return new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => reject(signal.reason), { + once: true, + }); + }); + }, +); +const interruptedRepairClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: interruptedRepairRun }), + }), + configLayer, + retryPolicyLayer, + ), + ), +); + +layer(interruptedRepairClientLayer)("repair interruption", (it) => { + it.effect("does not turn repair interruption into malformed output", () => + Effect.gen(function* () { + interruptedRepairCalls = 0; + interruptedRepairSignals.length = 0; + interruptedRepairRun.mockClear(); + const client = yield* DiagramGenerationClient; + const fiber = yield* Effect.forkChild( + client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt, + }), + ); + while (interruptedRepairSignals.length === 0) { + yield* Effect.yieldNow; + } + yield* Fiber.interrupt(fiber); + const exit = yield* Fiber.await(fiber); + + if (Exit.isSuccess(exit)) { + return assert.fail("Interrupted repair unexpectedly succeeded."); + } + assert.isTrue(Cause.hasInterrupts(exit.cause)); + assert.strictEqual(interruptedRepairRun.mock.calls.length, 2); + assert.isTrue(interruptedRepairSignals[0]?.aborted); + }), + ); +}); + +let multiRepairCalls = 0; +const multiRepairRun = vi.fn(async () => { + multiRepairCalls += 1; + if (multiRepairCalls === 1) { + return jsonResponse({ + candidates: [ + { + content: { + role: "model", + parts: [{ text: '{"type":"flowchart"' }], + }, + finishReason: "MAX_TOKENS", + }, + ], + }); + } + return multiRepairCalls === 2 + ? geminiTextResponse(invalidDiagramText) + : jsonResponse(geminiResponse); +}); +const multiRepairClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: multiRepairRun }), + }), + configLayer, + Layer.succeed(DiagramGenerationPolicy, { + concurrency: 2, + maxRepairAttempts: 2, + maxRetries: 0, + requestTimeoutMs: 1_000, + retryDelayMs: 100, + }), + ), + ), +); + +layer(multiRepairClientLayer)("multi-attempt semantic repair", (it) => { + it.effect("classifies truncation from only the latest response", () => + Effect.gen(function* () { + multiRepairCalls = 0; + multiRepairRun.mockClear(); + const client = yield* DiagramGenerationClient; + const candidate = yield* client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt, + }); + + assert.strictEqual(candidate.diagram?.id, expectedDiagram.id); + assert.strictEqual(multiRepairRun.mock.calls.length, 3); + expect( + JSON.stringify(multiRepairRun.mock.calls[1]?.[0].query), + ).not.toContain("Invalid model output:"); + expect(JSON.stringify(multiRepairRun.mock.calls[2]?.[0].query)).toContain( + "Invalid model output:", + ); + expect(candidate.diagnostics).toEqual( + expect.arrayContaining([ + "repair_failed: semantic repair attempt 1 failed.", + "repair_succeeded: semantic repair attempt 2 succeeded.", + ]), + ); + }), + ); +}); + const gatewayConstructionFailure = new Error("gateway construction failed"); const gatewayConstructionFailureLayer = CloudflareGoogleAiStudioClientLive.pipe( Layer.provide( @@ -864,7 +1685,7 @@ layer(malformedClientLayer)("malformed model responses", (it) => { ); }); -const truncatedRun = vi.fn(async () => +const truncatedRun = vi.fn(async () => jsonResponse({ candidates: [ { @@ -887,19 +1708,37 @@ const truncatedClientLayer = CloudflareGoogleAiStudioClientLive.pipe( ); 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, - }); + it.effect( + "regenerates a truncated response without repairing partial JSON", + () => + Effect.gen(function* () { + truncatedRun.mockClear(); + 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.", - ); - }), + assert.isUndefined(candidate.diagram); + expect(candidate.diagnostics).toContain( + "output_truncated: Gemini stopped at the maximum output-token budget; regenerate the complete diagram.", + ); + expect(candidate.diagnostics).toEqual( + expect.arrayContaining([ + "repair_attempted: regenerated a truncated response (attempt 1).", + "repair_failed: semantic repair attempt 1 failed.", + ]), + ); + assert.strictEqual(truncatedRun.mock.calls.length, 2); + expect(truncatedRun.mock.calls[1]?.[0]).toEqual( + expect.objectContaining({ + headers: expect.objectContaining({ "Cache-Control": "no-store" }), + }), + ); + expect( + JSON.stringify(truncatedRun.mock.calls[1]?.[0].query), + ).not.toContain("Invalid model output:"); + }), ); });