From 9f64991c008e36b480f5d8ee081f390ac6f89459 Mon Sep 17 00:00:00 2001 From: anandpant <109482096+anandpant@users.noreply.github.com> Date: Tue, 1 Sep 2026 21:25:46 +0000 Subject: [PATCH] fix(generation): repair invalid model output (#313) ## Scope - Add a configurable semantic repair budget with `maxRepairAttempts: 2` by default, allowing one more semantic repair when the first repair introduces a different violation class. - Repair invalid JSON/IR with the prior output plus all validator diagnostics; regenerate truncated output fresh. - Promote `flowchart.start_has_incoming` and `flowchart.self_loop`, including their validator hints, at the top of repair requests. Start-node failures reroute to the first process after start; self-loops reroute to an earlier distinct process or decision node. - Preserve the original malformed candidate and its text/raw/error/usage when a repair call fails or returns invalid output; enrich it with both attempts' diagnostics. - Catch typed repair-call failures while allowing interruption to propagate, and always record `sketchi_generation_repairs{outcome="failed"}` for failed repair calls. - Classify truncation from only the latest model response so a truncated regeneration that is still invalid proceeds to semantic repair. - Preserve every existing node and label except the minimal edit required by reported diagnostics, and restate the original diagram type, scenario, parsed minimum node count, required node labels, and required branch labels on every repair call. - Keep transport retries and timeouts around each model call, normalize empty optional edge labels, and extract the first complete JSON object when a model appends trailing output. - Bypass Cloudflare AI Gateway cache in fresh mode with `cf-aig-skip-cache: true`. ## Expected latency tradeoff - Each provider attempt has a 30-second timeout and up to two transient retries with 250 ms and 500 ms backoff, so one exhausted model call can take up to 90.75 seconds. - A malformed initial response can trigger two complete repair calls, making the configured generation-policy ceiling 272.25 seconds before outer request overhead. Requests above 120 seconds are therefore expected in the worst repair/retry cases; the added latency is the tradeoff for recovering when a repair introduces a new violation such as dropping below a requested minimum node count. - The deployed probe permits 302.25 seconds: the full 272.25-second policy budget plus a 30-second margin. ## Proof - `pnpm nx test diagram-generation --skip-nx-cache` - 37/37 passed at `237076a`. - `pnpm nx run-many -t typecheck,test,build --skip-nx-cache` - passed at `237076a`. - `pnpm nx build-storybook diagram-ui --skip-nx-cache` - passed at `237076a`. - Added typed repair-failure preservation, invalid-repair original-candidate preservation, interruption propagation, failed repair metric, latest-response truncation classification, priority start-node and self-loop correction, preserve-content/hard-requirement repair instructions, repair-introduced minimum violation recovery, cache bypass, requested-cycle, minimum-count, diagnostic-composition, and repair-success fixtures. This is PR 5 of 6 in the generation-reliability stack. --- .../src/lib/generate-scenario.test.ts | 1 + .../diagram/generation/src/lib/candidates.ts | 179 +++- packages/diagram/generation/src/lib/client.ts | 2 + .../src/lib/cloudflare-google-ai-studio.ts | 73 +- packages/diagram/generation/src/lib/policy.ts | 262 +++++- .../generation/src/lib/public-api.test.ts | 871 +++++++++++++++++- 6 files changed, 1284 insertions(+), 104 deletions(-) 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:"); + }), ); });