diff --git a/containers/liteparse/server.mjs b/containers/liteparse/server.mjs index d8de117fe..aa36cb98e 100644 --- a/containers/liteparse/server.mjs +++ b/containers/liteparse/server.mjs @@ -11,9 +11,19 @@ import { Transform } from "node:stream"; import { promisify } from "node:util"; const port = 8080; +// LiteParse defaults to 1000 pages and drops everything past that without reporting +// it, which publishes a "ready" projection silently missing the tail of the document. +// Parse one page past the supported ceiling so a document of exactly the ceiling is +// distinguishable from a truncated longer one. Measured at roughly 0.57 MB resident +// per page, so a full parse sits near 3 GB on the 8 GiB standard-2 instance — which +// is why concurrent parses are also capped below. +const supportedPages = 5000; +const maxConcurrentParses = 2; +let activeParses = 0; const parser = new LiteParse({ extractLinks: true, imageMode: "placeholder", + maxPages: supportedPages + 1, ocrEnabled: false, outputFormat: "markdown", quiet: true, @@ -29,6 +39,7 @@ createServer(async (request, response) => { let status = 500; let errorType = null; let errorMessage = null; + let holdsParseSlot = false; try { if ( @@ -53,9 +64,36 @@ createServer(async (request, response) => { return sendJson(response, status, { error: "Not found." }); } + // Each parse can hold gigabytes resident, and one over-committed container + // dies taking every in-flight request with it. Shed load instead: 503 is + // retryable by the caller, an OOM crash is not. + if (activeParses >= maxConcurrentParses) { + status = 503; + return sendJson(response, status, { + code: "EXTRACTOR_BUSY", + error: "Extractor is at capacity, retry shortly.", + }); + } + + // Held until the whole request finishes, including streaming the result out: + // the parsed pages stay resident while the response drains, so releasing the + // slot any earlier would let admissions outrun actual memory use. A parse that + // outlives its timeout still runs to completion in the background holding + // memory — that zombie cannot be cancelled, only kept rare by the timeout. + activeParses += 1; + holdsParseSlot = true; const bytes = await readPdfRequestBytes(request); inputBytes = bytes.byteLength; const result = await withTimeout(parser.parse(bytes), parseTimeoutMs); + + if (result.pages.length > supportedPages) { + throw new PdfValidationError( + 422, + "TOO_MANY_PAGES", + `PDFs longer than ${supportedPages} pages are not supported.`, + ); + } + pageCount = result.pages.length; status = 200; response.writeHead(status, { "content-type": "application/x-ndjson; charset=utf-8" }); @@ -78,6 +116,9 @@ createServer(async (request, response) => { } return sendJson(response, status, { error: "PDF parsing failed." }); } finally { + if (holdsParseSlot) { + activeParses -= 1; + } console.info( JSON.stringify({ duration_ms: Date.now() - startedAt, diff --git a/src/features/workspaces/content/workspace-read-observability.ts b/src/features/workspaces/content/workspace-read-observability.ts new file mode 100644 index 000000000..c08ce3ce1 --- /dev/null +++ b/src/features/workspaces/content/workspace-read-observability.ts @@ -0,0 +1,88 @@ +import type { WorkspaceContentReadResult } from "#/features/workspaces/content/workspace-content-contract"; +import { capturePostHogServerEvent } from "#/integrations/posthog/server"; + +/** + * Records the readiness state every file read was served in, across every surface + * that reads — assistant tools and MCP both route through the same operation. + * + * This is the consumption side of the extraction pipeline's telemetry: the + * extraction event says how long the fast-pass window lasted, and this says whether + * anyone was actually in it — reads served `provisional`, reads that hit the + * pending spinner, and reads that found a stalled or failed document. + */ +export function recordWorkspaceFileReadOutcomes(input: { + operationId: string; + results: readonly WorkspaceContentReadResult[]; + userId: string; + workspaceId: string; +}) { + for (const result of input.results) { + if (result.status === "pending") { + capture(input, { + elapsed_seconds: result.elapsedSeconds, + empty_page_count: null, + failure_code: null, + item_id: null, + phase: result.phase, + provisional: null, + returned_page_count: null, + status: "pending", + }); + continue; + } + + if (result.status === "failed" && result.type === "file") { + capture(input, { + elapsed_seconds: null, + empty_page_count: null, + failure_code: result.code, + item_id: null, + phase: null, + provisional: null, + returned_page_count: null, + status: "failed", + }); + continue; + } + + if (result.status === "ready" && result.type === "file") { + capture(input, { + elapsed_seconds: null, + empty_page_count: result.emptyPages?.length ?? 0, + failure_code: null, + item_id: result.itemId, + phase: null, + provisional: result.provisional ?? false, + returned_page_count: result.location.returned.length, + status: "ready", + }); + } + } +} + +function capture( + input: { operationId: string; userId: string; workspaceId: string }, + properties: { + elapsed_seconds: number | null; + empty_page_count: number | null; + failure_code: string | null; + item_id: string | null; + phase: "queued" | "extracting" | null; + provisional: boolean | null; + returned_page_count: number | null; + status: "ready" | "pending" | "failed"; + }, +) { + capturePostHogServerEvent({ + distinctId: input.userId, + event: "workspace_file_read_completed", + // Legitimate interest: operational readiness telemetry — ids and states only, + // no file names or content. + consentExempt: true, + properties: { + ...properties, + operation_id: input.operationId, + workspace_id: input.workspaceId, + }, + }); +} diff --git a/src/features/workspaces/conversion/container-file-conversion.ts b/src/features/workspaces/conversion/container-file-conversion.ts index 0e41177f1..f3fd92bbf 100644 --- a/src/features/workspaces/conversion/container-file-conversion.ts +++ b/src/features/workspaces/conversion/container-file-conversion.ts @@ -57,7 +57,7 @@ export async function convertFileStreamWithContainer(input: { sizeBytes: input.sizeBytes, }); - const [response] = await Promise.all([ + const response = await multipart.awaitResponse( input.container.fetch( new Request(input.url, { body: multipart.body, @@ -66,8 +66,7 @@ export async function convertFileStreamWithContainer(input: { method: "POST", } as RequestInit & { duplex: "half" }), ), - multipart.done, - ]); + ); if (!response.ok) { throw input.error(await getConversionErrorMessage(response)); diff --git a/src/features/workspaces/extraction/liteparse-projection.ts b/src/features/workspaces/extraction/liteparse-projection.ts index 7d6a22306..a8f8ef61a 100644 --- a/src/features/workspaces/extraction/liteparse-projection.ts +++ b/src/features/workspaces/extraction/liteparse-projection.ts @@ -1,6 +1,10 @@ import type { WorkflowStep } from "cloudflare:workers"; import { extractPdfWithLiteParse } from "#/features/workspaces/extraction/providers/liteparse"; +import { + getWorkspaceExtractionStepConfig, + workspaceExtractionStepBudgets, +} from "#/features/workspaces/extraction/workspace-extraction-budgets"; import type { LiteParseStageOutcome, WorkspaceFileExtractionWorkflowParams, @@ -28,10 +32,7 @@ export async function publishLiteParseProjection( try { return await step.do( "publish fast LiteParse projection", - { - retries: { limit: 1, delay: "5 seconds", backoff: "constant" }, - timeout: "2 minutes", - }, + getWorkspaceExtractionStepConfig(workspaceExtractionStepBudgets.liteParse), async () => { const kernel = await getWorkspaceKernelFromEnv(env, params.workspaceId); const { object, source } = await getWorkspaceFileSourceObject({ @@ -70,6 +71,9 @@ export async function publishLiteParseProjection( markdownLength: projection.manifest.markdownLength, pageCount: projection.manifest.pageCount, provisional: true, + // Brand healing runs so the reconciler never picks this row up again: + // one upgrade attempt per document, bounded structurally. + ...(params.healing ? { healed: true } : {}), }, actorUserId: params.actorUserId, clientMutationId: `${runId}:projection:liteparse-ready`, @@ -104,6 +108,7 @@ export async function publishLiteParseProjection( }); return { durationMs: Date.now() - startedAt, + errorMessage: error instanceof Error ? error.message : String(error), errorType: error instanceof Error ? error.name : "UnknownError", outcome: "error", }; diff --git a/src/features/workspaces/extraction/providers/liteparse.test.ts b/src/features/workspaces/extraction/providers/liteparse.test.ts index 7e4ee1e0c..264d58e61 100644 --- a/src/features/workspaces/extraction/providers/liteparse.test.ts +++ b/src/features/workspaces/extraction/providers/liteparse.test.ts @@ -33,6 +33,39 @@ describe("LiteParse response parsing", () => { expect(() => parseLiteParsePage(payload)).toThrow("LiteParse returned an invalid"); }); + // The workflow reads this error's name to decide whether to skip the paid tier, so + // mislabelling a refusal here means paying a provider to reach the same verdict on + // every reconciler sweep, forever. + it("reports a rejected document as unsupported rather than a retryable failure", async () => { + vi.mocked(requestWorkspaceFileProcessor).mockResolvedValue( + Response.json( + { code: "TOO_MANY_PAGES", error: "PDFs longer than 4999 pages." }, + { status: 422 }, + ), + ); + const pages = extractPdfWithLiteParse({} as Cloudflare.Env, { + body: new ReadableStream(), + fileName: "document.pdf", + sizeBytes: 1, + }); + + await expect(pages.next()).rejects.toMatchObject({ + name: "WorkspaceDocumentUnsupportedError", + message: "PDFs longer than 4999 pages.", + }); + }); + + it("treats any other processor failure as retryable", async () => { + vi.mocked(requestWorkspaceFileProcessor).mockResolvedValue(new Response("", { status: 500 })); + const pages = extractPdfWithLiteParse({} as Cloudflare.Env, { + body: new ReadableStream(), + fileName: "document.pdf", + sizeBytes: 1, + }); + + await expect(pages.next()).rejects.toMatchObject({ name: "Error" }); + }); + it("rejects an oversized NDJSON record and cancels the processor response", async () => { const cancel = vi.fn(); const body = new ReadableStream({ diff --git a/src/features/workspaces/extraction/providers/liteparse.ts b/src/features/workspaces/extraction/providers/liteparse.ts index 88b1e91f1..122c622d3 100644 --- a/src/features/workspaces/extraction/providers/liteparse.ts +++ b/src/features/workspaces/extraction/providers/liteparse.ts @@ -1,5 +1,6 @@ import type { MarkdownProjectionPage } from "#/features/workspaces/extraction/page-markdown-projection"; import { parseLiteParsePage } from "#/features/workspaces/extraction/providers/liteparse-response"; +import { WorkspaceDocumentUnsupportedError } from "#/features/workspaces/extraction/types"; import { requestWorkspaceFileProcessor } from "#/features/workspaces/files/workspace-file-processor"; const maxNdjsonLineBytes = 8 * 1024 * 1024; @@ -21,7 +22,12 @@ export async function* extractPdfWithLiteParse( }); if (!response.ok) { - throw new Error(`LiteParse failed with status ${response.status}.`); + // The processor answers 422 only when it has read the file and found it + // unusable — too long, encrypted, or damaged. Every other status is an + // extraction that went wrong and may work next time. + throw response.status === 422 + ? new WorkspaceDocumentUnsupportedError(await getLiteParseErrorMessage(response)) + : new Error(`LiteParse failed with status ${response.status}.`); } if (!response.body) { @@ -39,6 +45,17 @@ export async function* extractPdfWithLiteParse( } } +async function getLiteParseErrorMessage(response: Response) { + const body: unknown = await response.json().catch(() => null); + + return typeof body === "object" && + body !== null && + "error" in body && + typeof body.error === "string" + ? body.error + : "This document cannot be read."; +} + async function* readNdjsonLines(body: ReadableStream): AsyncGenerator { const reader = body.getReader(); const decoder = new TextDecoder(); diff --git a/src/features/workspaces/extraction/providers/llama-parse-credits.test.ts b/src/features/workspaces/extraction/providers/llama-parse-credits.test.ts new file mode 100644 index 000000000..c1164c2ca --- /dev/null +++ b/src/features/workspaces/extraction/providers/llama-parse-credits.test.ts @@ -0,0 +1,35 @@ +import { describe, expect, it } from "vitest"; + +import { getLlamaParseDerivedCredits } from "#/features/workspaces/extraction/providers/llama-parse-credits"; + +function createPages(total: number, optimized: number) { + return Array.from({ length: total }, (_, index) => ({ cost_optimized: index < optimized })); +} + +describe("getLlamaParseDerivedCredits", () => { + it("bills every page at the requested tier when nothing was optimized", () => { + expect(getLlamaParseDerivedCredits({ pages: createPages(9, 0) }, "agentic")).toBe(90); + }); + + it("bills optimizer-downgraded pages at the cost_effective rate", () => { + // The production 1,527-page run: 717 pages downgraded, 810 left on agentic. + expect(getLlamaParseDerivedCredits({ pages: createPages(1527, 717) }, "agentic")).toBe( + 717 * 3 + 810 * 10, + ); + }); + + it("prices each tier from its own rate", () => { + expect(getLlamaParseDerivedCredits({ pages: createPages(2, 0) }, "cost_effective")).toBe(6); + expect(getLlamaParseDerivedCredits({ pages: createPages(2, 0) }, "agentic_plus")).toBe(90); + }); + + it("treats a missing cost_optimized flag as not downgraded", () => { + expect(getLlamaParseDerivedCredits({ pages: [{}, {}] }, "agentic")).toBe(20); + }); + + it("returns null when there is no per-page breakdown to price", () => { + expect(getLlamaParseDerivedCredits({}, "agentic")).toBeNull(); + expect(getLlamaParseDerivedCredits({ pages: [] }, "agentic")).toBeNull(); + expect(getLlamaParseDerivedCredits(null, "agentic")).toBeNull(); + }); +}); diff --git a/src/features/workspaces/extraction/providers/llama-parse-credits.ts b/src/features/workspaces/extraction/providers/llama-parse-credits.ts new file mode 100644 index 000000000..d14db6006 --- /dev/null +++ b/src/features/workspaces/extraction/providers/llama-parse-credits.ts @@ -0,0 +1,36 @@ +import type { LlamaParseTier } from "#/features/workspaces/extraction/types"; +import { getBooleanValue, getRecordArrayValue } from "#/integrations/llamaparse/client"; + +// Published per-page rates, at $1.25 per 1,000 credits. The cost optimizer bills the +// pages it downgrades at the cost_effective rate, so a job's real cost is a blend and +// cannot be read off the requested tier alone. +const llamaParseCreditsPerPage: Record = { + cost_effective: 3, + agentic: 10, + agentic_plus: 45, +}; + +/** + * Blended credit cost read off the per-page breakdown: pages the optimizer downgraded + * bill at the cost_effective rate, the rest at the requested tier. + * + * This is derived from published rates, not an invoice. LlamaParse v2 has never + * populated a billed figure on a parse response, and reporting null there understates + * spend on exactly the runs worth investigating — a wrong-by-a-few-percent number + * still catches a job that cost ten times what it should. + */ +export function getLlamaParseDerivedCredits(metadata: unknown, tier: LlamaParseTier) { + const pages = getRecordArrayValue(metadata, "pages"); + if (pages.length === 0) { + return null; + } + + return pages.reduce( + (total, page) => + total + + (getBooleanValue(page, "cost_optimized") + ? llamaParseCreditsPerPage.cost_effective + : llamaParseCreditsPerPage[tier]), + 0, + ); +} diff --git a/src/features/workspaces/extraction/providers/llama-parse.test.ts b/src/features/workspaces/extraction/providers/llama-parse.test.ts new file mode 100644 index 000000000..ae07a7bc3 --- /dev/null +++ b/src/features/workspaces/extraction/providers/llama-parse.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, it } from "vitest"; + +import { buildLlamaParseJobRequest } from "#/features/workspaces/extraction/providers/llama-parse"; + +describe("buildLlamaParseJobRequest", () => { + it.each(["agentic", "agentic_plus"] as const)("enables the cost optimizer on %s", (tier) => { + expect(buildLlamaParseJobRequest({ fileId: "f", tier }).processing_options).toEqual({ + cost_optimizer: { enable: true }, + }); + }); + + // LlamaParse rejects cost_optimizer + cost_effective with a 422 — sending it + // unconditionally once made every cost_effective parse fail outright. + it("omits the cost optimizer on cost_effective", () => { + expect( + buildLlamaParseJobRequest({ fileId: "f", tier: "cost_effective" }).processing_options, + ).toEqual({}); + }); +}); diff --git a/src/features/workspaces/extraction/providers/llama-parse.ts b/src/features/workspaces/extraction/providers/llama-parse.ts index 0771a7179..1941843d7 100644 --- a/src/features/workspaces/extraction/providers/llama-parse.ts +++ b/src/features/workspaces/extraction/providers/llama-parse.ts @@ -8,6 +8,7 @@ import type { MarkdownExtractionProvider, MarkdownExtractionResult, } from "#/features/workspaces/extraction/types"; +import { getLlamaParseDerivedCredits } from "#/features/workspaces/extraction/providers/llama-parse-credits"; import { getNumberValue, getRecordArrayValue, @@ -18,7 +19,13 @@ import { import { createStreamingMultipartFile } from "#/lib/http/streaming-multipart"; const llamaParsePollIntervalMs = 2_000; -const llamaParseMaxPollMs = 300_000; +// Parse time scales with page count, so this has to clear the largest document we +// accept, not the typical one. Measured: 1,162 agentic pages in ~4m48s, about 0.25s +// per page, so the 5,000-page ceiling projects to roughly 21 minutes — and +// LlamaParse's own job timeout is 30 minutes of parsing excluding queue time. This +// exists to turn a wedged job into a clear error, not to cap normal work; Workflows +// imposes no wall-clock limit on a step. +const llamaParseMaxPollMs = 35 * 60_000; const llamaParseVersion = "latest"; export function createLlamaParseExtractionProvider(env: Env): MarkdownExtractionProvider { @@ -54,6 +61,30 @@ export function createLlamaParseExtractionProvider(env: Env): MarkdownExtraction }; } +/** Exported for tests: sending the wrong shape here once broke an entire tier. */ +export function buildLlamaParseJobRequest(input: { fileId: string; tier: LlamaParseTier }) { + return { + file_id: input.fileId, + tier: input.tier, + version: llamaParseVersion, + output_options: { + markdown: { + tables: { + output_tables_as_markdown: true, + }, + }, + }, + // The optimizer downgrades individual simple pages to the cost_effective + // tier, so LlamaParse rejects the combination with a 422 when that is + // already the requested tier — there is nothing left to downgrade to. + // Sending it unconditionally made every cost_effective parse fail outright. + processing_options: + input.tier === "agentic" || input.tier === "agentic_plus" + ? { cost_optimizer: { enable: true } } + : {}, + }; +} + function normalizeLlamaParseTier(mode: MarkdownExtractionInput["mode"]): LlamaParseTier { if (mode === "cost_effective" || mode === "agentic_plus") { return mode; @@ -71,7 +102,7 @@ async function uploadLlamaParseFile(env: Env, input: MarkdownExtractionInput) { formFieldName: "file", sizeBytes: input.sizeBytes, }); - const [responseJson] = await Promise.all([ + const responseJson = await multipart.awaitResponse( llamaCloudJsonRequest({ env, path: "/api/v1/beta/files", @@ -80,8 +111,7 @@ async function uploadLlamaParseFile(env: Env, input: MarkdownExtractionInput) { headers: { "content-type": multipart.contentType }, body: multipart.body, }), - multipart.done, - ]); + ); const fileId = getStringValue(responseJson, "id"); if (!fileId) { @@ -106,23 +136,7 @@ async function startLlamaParseJob( headers: { "Content-Type": "application/json", }, - body: JSON.stringify({ - file_id: input.fileId, - tier: input.tier, - version: llamaParseVersion, - output_options: { - markdown: { - tables: { - output_tables_as_markdown: true, - }, - }, - }, - processing_options: { - cost_optimizer: { - enable: true, - }, - }, - }), + body: JSON.stringify(buildLlamaParseJobRequest(input)), }); const jobId = getStringValue(responseJson, "id"); @@ -160,7 +174,9 @@ async function pollLlamaParseJob(env: Env, jobId: string) { await wait(llamaParsePollIntervalMs); } - throw new Error("LlamaParse job timed out."); + throw new Error( + `LlamaParse job ${jobId} did not finish within ${llamaParseMaxPollMs / 60_000} minutes.`, + ); } function getLlamaParseMarkdownPages(value: unknown): MarkdownProjectionPage[] { @@ -211,6 +227,13 @@ function getLlamaParseMetadata( const metadata = getRecordValue(value, "metadata"); const usage = getRecordValue(value, "usage"); const job = getRecordValue(value, "job"); + // Splits a slow enhancement into "LlamaParse was busy" versus "LlamaParse was + // slow", which decide different responses: waiting out a queue versus changing + // tier or budgets. Verified shape: job_metadata.state_transitions carries + // pending_at / running_at / completed_at ISO timestamps. + const transitions = getRecordValue(getRecordValue(value, "job_metadata"), "state_transitions"); + const queuedMs = getStateTransitionMs(transitions, "pending_at", "running_at"); + const parseMs = getStateTransitionMs(transitions, "running_at", "completed_at"); const result: Record = { fileId: input.fileId, jobId: input.jobId, @@ -223,11 +246,16 @@ function getLlamaParseMetadata( getNumberValue(metadata, "pageCount") || getNumberValue(value, "page_count") || getNumberValue(value, "pageCount"); + // v2 has never actually populated a billed figure on any of these — they are kept + // only so a real one takes precedence if LlamaParse starts reporting it. Until + // then the per-page breakdown is the only signal, so derive from that instead of + // reporting null and understating spend on exactly the runs worth investigating. const creditsUsed = getNumberValue(usage, "credits") ?? getNumberValue(usage, "credits_used") ?? getNumberValue(metadata, "credits") ?? - getNumberValue(metadata, "credits_used"); + getNumberValue(metadata, "credits_used") ?? + getLlamaParseDerivedCredits(metadata, input.tier); const status = getStringValue(job, "status") ?? getStringValue(value, "status"); if (pageCount !== null) { @@ -243,9 +271,28 @@ function getLlamaParseMetadata( result.status = status; } + if (queuedMs !== null) { + result.queuedMs = queuedMs; + } + + if (parseMs !== null) { + result.parseMs = parseMs; + } + return result; } +function getStateTransitionMs(transitions: unknown, fromKey: string, toKey: string) { + const from = getStringValue(transitions, fromKey); + const to = getStringValue(transitions, toKey); + if (!from || !to) { + return null; + } + + const elapsedMs = Date.parse(to) - Date.parse(from); + return Number.isFinite(elapsedMs) && elapsedMs >= 0 ? Math.round(elapsedMs) : null; +} + function wait(ms: number) { return new Promise((resolve) => setTimeout(resolve, ms)); } diff --git a/src/features/workspaces/extraction/types.ts b/src/features/workspaces/extraction/types.ts index fcb52fb4c..a2c499933 100644 --- a/src/features/workspaces/extraction/types.ts +++ b/src/features/workspaces/extraction/types.ts @@ -13,11 +13,19 @@ export interface WorkspaceFileExtractionWorkflowParams { actorUserId: string | null; assetKind: WorkspaceFileAssetKind; requestId: string | null; + /** + * True only on runs the reconciler starts to upgrade a projection stuck on the + * fast tier. The run brands the fast projection it republishes with + * `healed: true` and the reconciler only heals unbranded rows, which bounds + * healing to one attempt per document structurally. A dedicated field rather + * than a sentinel request id, so nothing a client influences can brand a row. + */ + healing?: boolean; } export type LiteParseStageOutcome = | { durationMs: number; outcome: "skipped" } - | { durationMs: number; errorType: string; outcome: "error" } + | { durationMs: number; errorMessage: string; errorType: string; outcome: "error" } | { durationMs: number; markdownLength: number; @@ -47,3 +55,17 @@ export interface MarkdownExtractionProvider { id: WorkspaceFileExtractionProviderId; extract(input: MarkdownExtractionInput): Promise; } + +export const workspaceDocumentUnsupportedErrorName = "WorkspaceDocumentUnsupportedError"; + +/** + * The document itself cannot be read — too long, encrypted, or damaged — as opposed to + * an extraction that merely failed. Terminal for every tier: the fast pass reaches + * this verdict in seconds and for free, and no paid provider will reach a different + * one, so paying for a second opinion only buys the same answer. + */ +export class WorkspaceDocumentUnsupportedError extends Error { + // Spelled out rather than read off the class, which a minifier may rename while the + // stage outcome carrying it between steps stays a plain string. + override name = workspaceDocumentUnsupportedErrorName; +} diff --git a/src/features/workspaces/extraction/workspace-extraction-budgets.ts b/src/features/workspaces/extraction/workspace-extraction-budgets.ts new file mode 100644 index 000000000..8eb948bf2 --- /dev/null +++ b/src/features/workspaces/extraction/workspace-extraction-budgets.ts @@ -0,0 +1,73 @@ +const minuteMs = 60_000; + +interface WorkspaceExtractionStepBudget { + attempts: number; + retryDelayMs: number; + timeoutMs: number; +} + +/** + * Time budget for every step of an extraction, in one place. + * + * These are deliberately generous: parse time, upload time and projection writes all + * scale with page count, and a budget sized for a typical document silently kills a + * long one that is working correctly. + * + * They live together because the stall threshold below is their sum. Kept apart, that + * relationship survives only as a comment, and a comment does not fail the build when + * someone raises a timeout — which is how a healthy long document ends up classified + * as abandoned and re-queued as a second billable workflow. + */ +export const workspaceExtractionStepBudgets = { + /** + * Fast pass: container cold start (up to 60s), streaming the source in, the + * container's own 90s parse cap, and a single projection write since schema v2 + * packed pages into one object. Must stay above the processor request abort so + * the inner timeout fires first and yields a named error. + */ + liteParse: { attempts: 2, retryDelayMs: 15_000, timeoutMs: 6 * minuteMs }, + /** + * Upload, the provider's own poll ceiling, fetching the result, and writing the + * page projection. Single attempt: a retry cannot resume the job it lost, only + * start and pay for another. + */ + extract: { attempts: 1, retryDelayMs: 30_000, timeoutMs: 45 * minuteMs }, + /** Publishing the finished projection — a single kernel call. */ + publish: { attempts: 4, retryDelayMs: 10_000, timeoutMs: 2 * minuteMs }, +} as const satisfies Record; + +function getWorstCaseMs(budget: WorkspaceExtractionStepBudget) { + return budget.timeoutMs * budget.attempts + budget.retryDelayMs * (budget.attempts - 1); +} + +/** + * How long a projection may sit in `processing` before it is treated as stalled. + * + * The row is written once when extraction starts and not touched again until it + * finishes, so this has to clear every step's budget back to back, with headroom. + * + * Both the read path and the reconciler use it, so setting it too low does not merely + * mislabel a slow document — it queues a duplicate workflow, and a duplicate bill, + * against one that is still parsing. + */ +export const workspaceExtractionStallThresholdMs = + Math.ceil( + (Object.values(workspaceExtractionStepBudgets).reduce( + (total, budget) => total + getWorstCaseMs(budget), + 0, + ) * + 1.25) / + minuteMs, + ) * minuteMs; + +/** Shapes a budget into the retry and timeout options a workflow step expects. */ +export function getWorkspaceExtractionStepConfig(budget: WorkspaceExtractionStepBudget) { + return { + retries: { + backoff: "constant", + delay: budget.retryDelayMs, + limit: budget.attempts - 1, + }, + timeout: budget.timeoutMs, + } as const; +} diff --git a/src/features/workspaces/extraction/workspace-file-extraction-observability.ts b/src/features/workspaces/extraction/workspace-file-extraction-observability.ts index b7b1b283b..ff2d0521a 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-observability.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-observability.ts @@ -14,66 +14,95 @@ import { capturePostHogServerEvent } from "#/integrations/posthog/server"; import { getTelemetryRuntimeContext } from "#/integrations/posthog/server-context"; import type { PostHogTelemetryScheduler } from "#/integrations/posthog/scheduler"; -interface WorkspaceFileExtractionOutcomeBase { +/** + * What the enhanced pass produced, as a value rather than control flow, so the + * workflow can hand both stage outcomes to one telemetry call instead of shaping a + * different record at every exit. + */ +export type WorkspaceFileEnhancementOutcome = + | { + creditsUsed: number | null; + durationMs: number; + outcome: "success"; + pageCount: number; + /** Provider-reported time the job sat queued before running, when known. */ + queuedMs: number | null; + /** Provider-reported time spent actually parsing, when known. */ + parseMs: number | null; + provider: WorkspaceFileExtractionProviderId; + providerMode: WorkspaceFileExtractionMode; + routeReason: string; + } + | { + // Non-null when the provider finished and billed but a later step failed. + // Still null in one narrow case: the projection write that shares the + // extract step fails after the provider returned, because the billed + // metadata cannot escape a failed step. Rare, and preferable to widening + // the step contract just to carry it out. + creditsUsed: number | null; + durationMs: number; + error: unknown; + outcome: "error"; + }; + +export function recordWorkspaceFileExtractionOutcome(input: { durationMs: number; + enhancement: WorkspaceFileEnhancementOutcome; instanceId: string; liteParse: LiteParseStageOutcome; params: WorkspaceFileExtractionWorkflowParams; schedule: PostHogTelemetryScheduler; - enhancement: - | { durationMs: number; outcome: "success" } - | { durationMs: number; error: unknown; outcome: "error" }; -} - -type WorkspaceFileExtractionOutcome = WorkspaceFileExtractionOutcomeBase & - ( - | { - error: unknown; - outcome: "error"; - } - | { - creditsUsed: number | null; - outcome: "partial" | "success"; - pageCount: number; - provider: WorkspaceFileExtractionProviderId | "liteparse"; - providerMode: WorkspaceFileExtractionMode; - routeReason: string; - } - ); - -export function recordWorkspaceFileExtractionOutcome(input: WorkspaceFileExtractionOutcome) { +}) { const requestContext = getTelemetryRuntimeContext(); if (input.params.requestId) { requestContext.properties.request_id = input.params.requestId; } + // The run's outcome follows from the two stages: the enhanced pass succeeding is + // success, failing with a fast projection still published is partial, and failing + // with nothing readable is an error. + const outcome = + input.enhancement.outcome === "success" + ? ("success" as const) + : input.liteParse.outcome === "success" + ? ("partial" as const) + : ("error" as const); const outcomeFields = - input.outcome !== "error" + input.enhancement.outcome === "success" ? { // provider_mode is the tier we asked for; credits_used is what the cost // optimizer actually billed, so the two disagree on mixed-complexity files. - credits_used: input.creditsUsed, + credits_used: input.enhancement.creditsUsed, error_type: null, - page_count: input.pageCount, - provider: input.provider, - provider_mode: input.providerMode, - route_reason: input.routeReason, + page_count: input.enhancement.pageCount, + provider: input.enhancement.provider, + provider_mode: input.enhancement.providerMode, + route_reason: input.enhancement.routeReason, } - : { - credits_used: null, - error_type: input.error instanceof Error ? input.error.name : "UnknownError", - page_count: null, - provider: null, - provider_mode: null, - route_reason: null, - }; + : input.liteParse.outcome === "success" + ? { + credits_used: input.enhancement.creditsUsed, + error_type: null, + page_count: input.liteParse.pageCount, + provider: "liteparse", + provider_mode: "fast", + route_reason: "LiteParse projection retained after enhancement failed.", + } + : { + credits_used: input.enhancement.creditsUsed, + error_type: getErrorType(input.enhancement.error), + page_count: null, + provider: null, + provider_mode: null, + route_reason: null, + }; const fields = { actor_user_id: input.params.actorUserId, asset_kind: input.params.assetKind, duration_ms: input.durationMs, item_id: input.params.itemId, - outcome: input.outcome, + outcome, request_id: input.params.requestId, workflow_id: input.instanceId, workspace_id: input.params.workspaceId, @@ -82,6 +111,10 @@ export function recordWorkspaceFileExtractionOutcome(input: WorkspaceFileExtract enhancement_error_message: null, enhancement_outcome: input.enhancement.outcome, enhancement_duration_ms: input.enhancement.durationMs, + enhancement_queued_ms: + input.enhancement.outcome === "success" ? input.enhancement.queuedMs : null, + enhancement_parse_ms: + input.enhancement.outcome === "success" ? input.enhancement.parseMs : null, liteparse_duration_ms: input.liteParse.durationMs, liteparse_error_type: input.liteParse.outcome === "error" ? input.liteParse.errorType : null, liteparse_markdown_length: @@ -91,10 +124,10 @@ export function recordWorkspaceFileExtractionOutcome(input: WorkspaceFileExtract ...outcomeFields, }; - if (input.outcome === "error") { + if (input.enhancement.outcome === "error" && outcome === "error") { recordOperationalFailure({ distinctId: input.params.actorUserId ?? undefined, - error: input.error, + error: input.enhancement.error, event: "workspace_file_extraction", fields, requestContext, @@ -104,7 +137,7 @@ export function recordWorkspaceFileExtractionOutcome(input: WorkspaceFileExtract logOperationalEvent({ event: "workspace_file_extraction", fields, - outcome: input.outcome, + outcome, requestContext, }); } diff --git a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts index 351781884..8044c01fb 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts @@ -1,6 +1,7 @@ import type { WorkspaceFileExtractionWorkflowParams } from "#/features/workspaces/extraction/types"; +import { capturePostHogServerEvent } from "#/integrations/posthog/server"; import { getWorkspaceFileExtractionWorkflowId } from "#/features/workspaces/extraction/workspace-file-extraction-workflow-id"; -import { workspaceExtractionStallThresholdMs } from "#/features/workspaces/extraction/workspace-projection-readiness"; +import { workspaceExtractionStallThresholdMs } from "#/features/workspaces/extraction/workspace-extraction-budgets"; import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; import { workspaceFileAssetKindSchema } from "#/features/workspaces/model/workspace-file"; @@ -9,6 +10,7 @@ const failedExtractionCooldownMs = 15 * 60_000; const workflowBatchSize = 100; export async function reconcileWorkspaceFileExtractions(input: { + schedule?: (task: Promise) => void; sql: WorkspaceKernelSql; workflow: Workflow; workspaceId: string; @@ -19,12 +21,20 @@ export async function reconcileWorkspaceFileExtractions(input: { id: string; object_key: string; projection_updated_at: number; + reason: string; }>` SELECT json_extract(i.metadata_json, '$.assetKind') AS asset_kind, i.id, i.object_key, - COALESCE(p.updated_at, i.created_at) AS projection_updated_at + COALESCE(p.updated_at, i.created_at) AS projection_updated_at, + CASE + WHEN p.item_id IS NULL THEN 'missing' + WHEN p.status = 'failed' THEN 'failed' + WHEN p.status = 'processing' THEN 'stalled' + WHEN p.object_key IS NULL OR p.source_hash IS NULL THEN 'unreadable' + ELSE 'provisional' + END AS reason FROM kernel_items i LEFT JOIN kernel_item_projections p ON p.item_id = i.id AND p.format = 'pages' @@ -46,9 +56,44 @@ export async function reconcileWorkspaceFileExtractions(input: { AND (p.object_key IS NULL OR p.source_hash IS NULL) AND p.updated_at <= ${now - workspaceExtractionStallThresholdMs} ) + -- A provisional row means the fast pass published and the enhanced pass + -- failed: readable, but stuck at fast-tier quality with nothing else + -- revisiting it. Heal it once — a healing run brands the row it + -- republishes, and branded rows are never picked up again, so a + -- persistently failing document cannot become a paid retry per sweep. + OR ( + p.status = 'ready' + AND json_extract(p.metadata_json, '$.provisional') = 1 + AND json_extract(p.metadata_json, '$.healed') IS NULL + AND p.updated_at <= ${now - workspaceExtractionStallThresholdMs} + ) ) ORDER BY i.created_at ASC `; + // Healing is supposed to be rare and bounded; a workspace that shows up here + // every sweep is a loop this event exists to catch. + if (candidates.length > 0) { + const countByReason = (reason: string) => + candidates.filter((candidate) => candidate.reason === reason).length; + capturePostHogServerEvent({ + distinctId: input.workspaceId, + event: "workspace_file_extraction_healing_enqueued", + // Legitimate interest: operational pipeline telemetry, no user identity. + consentExempt: true, + processPerson: false, + properties: { + failed: countByReason("failed"), + missing: countByReason("missing"), + provisional: countByReason("provisional"), + stalled: countByReason("stalled"), + total: candidates.length, + unreadable: countByReason("unreadable"), + workspace_id: input.workspaceId, + }, + schedule: input.schedule, + }); + } + for (let offset = 0; offset < candidates.length; offset += workflowBatchSize) { const workflows = ( await Promise.all( @@ -62,6 +107,7 @@ export async function reconcileWorkspaceFileExtractions(input: { const params = { actorUserId: null, assetKind: assetKind.data, + healing: true, itemId: candidate.id, requestId: extractionHealingVersion, workspaceId: input.workspaceId, diff --git a/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts b/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts index 1bfaa966b..78a8b4600 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts @@ -1,9 +1,21 @@ import { WorkflowEntrypoint, type WorkflowEvent, type WorkflowStep } from "cloudflare:workers"; import { publishLiteParseProjection } from "#/features/workspaces/extraction/liteparse-projection"; -import { recordWorkspaceFileExtractionOutcome } from "#/features/workspaces/extraction/workspace-file-extraction-observability"; +import { + recordWorkspaceFileExtractionOutcome, + type WorkspaceFileEnhancementOutcome, +} from "#/features/workspaces/extraction/workspace-file-extraction-observability"; import { createMarkdownExtractionProvider } from "#/features/workspaces/extraction/providers/index"; -import type { WorkspaceFileExtractionWorkflowParams } from "#/features/workspaces/extraction/types"; +import { + getWorkspaceExtractionStepConfig, + workspaceExtractionStepBudgets, +} from "#/features/workspaces/extraction/workspace-extraction-budgets"; +import { + WorkspaceDocumentUnsupportedError, + workspaceDocumentUnsupportedErrorName, + type LiteParseStageOutcome, + type WorkspaceFileExtractionWorkflowParams, +} from "#/features/workspaces/extraction/types"; import type { WorkspaceFileExtractionMode, WorkspaceFileExtractionProviderId, @@ -16,6 +28,12 @@ import { import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; import { getWorkspaceUploadFamily } from "#/features/workspaces/model/workspace-file"; +/** + * Extracts an uploaded file into page markdown in two passes: a fast local one so the + * document is readable within seconds, then an enhanced provider pass that replaces + * it. The run settles into exactly one end state — enhanced ready, fast retained, or + * failed — and records one telemetry event describing both passes. + */ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< Cloudflare.Env, WorkspaceFileExtractionWorkflowParams @@ -25,7 +43,6 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< step: WorkflowStep, ) { const params = assertWorkflowParams(event.payload); - const schedule = (task: Promise) => this.ctx.waitUntil(task); const processing = await step.do("mark extraction processing", async () => { const kernel = await getWorkspaceKernelFromEnv(this.env, params.workspaceId); @@ -45,32 +62,98 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< if (liteParse.outcome === "discarded") { return { status: "discarded" as const }; } - const enhancementStartedAt = Date.now(); - let extraction: StagedPageExtractionResult; - // Captured as soon as extraction returns, because the steps after it can still - // fail. LlamaParse has already billed by then, and reporting null there would - // quietly understate spend on exactly the runs worth investigating. - let extractionCreditsUsed: number | null = null; - let result: - | { - pageCount: number; - provider: WorkspaceFileExtractionProviderId; - providerMode: WorkspaceFileExtractionMode; - status: "ready"; - } - | { status: "discarded" }; + + const enhancement = await this.enhance(step, event, params, liteParse); + if (enhancement.outcome === "discarded") { + return { status: "discarded" as const }; + } + + // Nothing readable was published: the fast pass did not produce a projection + // and the enhanced pass failed, so the item must leave `processing` or readers + // would wait on an extraction that is no longer running. + if (enhancement.outcome === "error" && liteParse.outcome !== "success") { + const failed = await step.do("mark extraction failed", async () => { + const kernel = await getWorkspaceKernelFromEnv(this.env, params.workspaceId); + return kernel.upsertFileProjection({ + itemId: params.itemId, + format: "pages", + status: "failed", + errorMessage: getErrorMessage(enhancement.error), + actorUserId: params.actorUserId, + clientMutationId: `${event.instanceId}:projection:failed`, + }); + }); + if (failed === "discarded") { + return { status: "discarded" as const }; + } + } + + await step.do("record extraction outcome", async () => { + recordWorkspaceFileExtractionOutcome({ + durationMs: Date.now() - event.timestamp.getTime(), + enhancement, + instanceId: event.instanceId, + liteParse, + params, + schedule: (task) => this.ctx.waitUntil(task), + }); + + return { recorded: true }; + }); + + if (enhancement.outcome === "success") { + return { + pageCount: enhancement.pageCount, + provider: enhancement.provider, + providerMode: enhancement.providerMode, + status: "ready" as const, + }; + } + + if (liteParse.outcome === "success") { + return { + pageCount: liteParse.pageCount, + provider: "liteparse" as const, + providerMode: "fast" as const, + status: "ready" as const, + }; + } + + throw enhancement.error; + } + + /** + * The enhanced pass, returned as a value: failure here is an expected outcome the + * run settles on — retained fast projection or a failed item — not an exception + * that abandons the workflow. + */ + private async enhance( + step: WorkflowStep, + event: Readonly>, + params: WorkspaceFileExtractionWorkflowParams, + liteParse: LiteParseStageOutcome, + ): Promise { + const startedAt = Date.now(); + // Captured as soon as the provider returns, because the publish step after it + // can still fail and the provider has already billed by then. + let creditsUsed: number | null = null; try { - extraction = await step.do( + // A document the free pass has already read and rejected will not become + // readable by paying for a slower one. Failing without calling the provider + // matters because the reconciler re-runs failures on a cooldown — letting + // this through would buy an identical verdict from a paid provider on every + // sweep. + if ( + liteParse.outcome === "error" && + liteParse.errorType === workspaceDocumentUnsupportedErrorName + ) { + throw new WorkspaceDocumentUnsupportedError(liteParse.errorMessage); + } + + const extraction = await step.do( "extract page markdown with provider", - { - retries: { - limit: 2, - delay: "30 seconds", - backoff: "exponential", - }, - timeout: "10 minutes", - }, + getWorkspaceExtractionStepConfig(workspaceExtractionStepBudgets.extract), async (): Promise => { const kernel = await getWorkspaceKernelFromEnv(this.env, params.workspaceId); const { object, source } = await getWorkspaceFileSourceObject({ @@ -117,28 +200,15 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< }, ); - extractionCreditsUsed = getExtractionCreditsUsed(extraction.metadata); + creditsUsed = getExtractionCreditsUsed(extraction.metadata); - result = await step.do( + const published = await step.do( "write extracted projections", - { - retries: { - limit: 3, - delay: "10 seconds", - backoff: "exponential", - }, - timeout: "5 minutes", - }, + getWorkspaceExtractionStepConfig(workspaceExtractionStepBudgets.publish), async () => { const kernel = await getWorkspaceKernelFromEnv(this.env, params.workspaceId); - const metadataJson = { - ...extraction.metadata, - routeReason: extraction.routeReason, - pageCount: extraction.pageCount, - markdownLength: extraction.markdownLength, - }; - const status = await publishWorkspacePageProjection({ + return publishWorkspacePageProjection({ bucket: this.env.WORKSPACE_KERNEL_FILES, kernel, projection: { @@ -149,121 +219,41 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< provider: extraction.provider, providerMode: extraction.providerMode, sourceHash: extraction.sourceHash, - metadataJson, + metadataJson: { + ...extraction.metadata, + routeReason: extraction.routeReason, + pageCount: extraction.pageCount, + markdownLength: extraction.markdownLength, + }, actorUserId: params.actorUserId, clientMutationId: `${event.instanceId}:projection:enhanced-ready`, }, }); - if (status === "discarded") { - return { status: "discarded" as const }; - } - - return { - status: "ready" as const, - provider: extraction.provider, - providerMode: extraction.providerMode, - pageCount: extraction.pageCount, - }; }, ); - if (result.status === "discarded") { - return result; - } - } catch (error) { - if (liteParse.outcome === "success") { - await step.do("record partial extraction outcome", async () => { - recordWorkspaceFileExtractionOutcome({ - // Null only when extraction itself never completed; a failure in the - // steps after it still owes whatever LlamaParse already charged. - creditsUsed: extractionCreditsUsed, - durationMs: Date.now() - event.timestamp.getTime(), - enhancement: { - durationMs: Date.now() - enhancementStartedAt, - error, - outcome: "error", - }, - instanceId: event.instanceId, - liteParse, - outcome: "partial", - pageCount: liteParse.pageCount, - params, - provider: "liteparse", - providerMode: "fast", - routeReason: "LiteParse projection retained after enhancement failed.", - schedule, - }); - - return { outcome: "partial" }; - }); - - return { - pageCount: liteParse.pageCount, - provider: "liteparse", - providerMode: "fast", - status: "ready", - }; - } - - const failed = await step.do("mark extraction failed", async () => { - const kernel = await getWorkspaceKernelFromEnv(this.env, params.workspaceId); - return kernel.upsertFileProjection({ - itemId: params.itemId, - format: "pages", - status: "failed", - errorMessage: getErrorMessage(error), - actorUserId: params.actorUserId, - clientMutationId: `${event.instanceId}:projection:failed`, - }); - }); - if (failed === "discarded") { - return { status: "discarded" as const }; + if (published === "discarded") { + return { outcome: "discarded" as const }; } - await step.do("record extraction failure", async () => { - recordWorkspaceFileExtractionOutcome({ - durationMs: Date.now() - event.timestamp.getTime(), - enhancement: { - durationMs: Date.now() - enhancementStartedAt, - error, - outcome: "error", - }, - error, - instanceId: event.instanceId, - liteParse, - outcome: "error", - params, - schedule, - }); - - return { outcome: "error" }; - }); - - throw error; - } - - await step.do("record extraction outcome", async () => { - recordWorkspaceFileExtractionOutcome({ - creditsUsed: getExtractionCreditsUsed(extraction.metadata), - durationMs: Date.now() - event.timestamp.getTime(), - enhancement: { - durationMs: Date.now() - enhancementStartedAt, - outcome: "success", - }, - instanceId: event.instanceId, - liteParse, - outcome: "success", + return { + creditsUsed, + durationMs: Date.now() - startedAt, + outcome: "success" as const, pageCount: extraction.pageCount, - params, + queuedMs: getMetadataNumber(extraction.metadata, "queuedMs"), + parseMs: getMetadataNumber(extraction.metadata, "parseMs"), provider: extraction.provider, providerMode: extraction.providerMode, routeReason: extraction.routeReason, - schedule, - }); - - return { outcome: "success" }; - }); - - return result; + }; + } catch (error) { + return { + creditsUsed, + durationMs: Date.now() - startedAt, + error, + outcome: "error" as const, + }; + } } } @@ -291,6 +281,7 @@ function assertWorkflowParams( actorUserId: value.actorUserId ?? null, assetKind: value.assetKind, requestId: value.requestId ?? null, + healing: value.healing === true, }; } @@ -300,5 +291,10 @@ function getErrorMessage(error: unknown) { function getExtractionCreditsUsed(metadata: StagedPageExtractionResult["metadata"]) { // Only LlamaParse reports credits; other providers leave the key absent. - return typeof metadata.creditsUsed === "number" ? metadata.creditsUsed : null; + return getMetadataNumber(metadata, "creditsUsed"); +} + +function getMetadataNumber(metadata: StagedPageExtractionResult["metadata"], key: string) { + const value = metadata[key]; + return typeof value === "number" ? value : null; } diff --git a/src/features/workspaces/extraction/workspace-page-projection.test.ts b/src/features/workspaces/extraction/workspace-page-projection.test.ts index 7a328c3e3..383d9bb76 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.test.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.test.ts @@ -40,11 +40,9 @@ describe("workspace page projections", () => { pages: { requested: "2-3", returned: [2, 3], total: 3 }, }); const prefix = reference.manifestObjectKey.slice(0, -"manifest.json".length); - expect(storage.readKeys).toEqual([ - reference.manifestObjectKey, - getWorkspacePageObjectKey(prefix, 2), - getWorkspacePageObjectKey(prefix, 3), - ]); + // A contiguous selection coalesces into a single ranged read of the packed + // pages object rather than one request per page. + expect(storage.readKeys).toEqual([reference.manifestObjectKey, `${prefix}pages.md`]); }); it("preserves missing page numbers as blank pages", async () => { @@ -162,30 +160,78 @@ describe("workspace page projections", () => { expect(storage.readKeys).toEqual([reference.manifestObjectKey]); }); - it("reads projections published before per-page sizes were added", async () => { + // Projections written before schema version 2 store one object per page and their + // regeneration is billable, so the read path must keep serving them unmigrated. + it("reads legacy per-page projections without migration", async () => { const storage = createObjectStorage(); - const reference = await writeWorkspacePageProjection({ - bucket: storage.bucket, - itemId: "item-1", - pages: [{ pageNumber: 1, markdown: "Page 1" }], - provider: "liteparse", - providerMode: "fast", - runId: "run-1", - sourceHash: "etag-1", - tier: "fast", - workspaceId: "workspace-1", + const prefix = "workspace_file_objects/workspace-1/item-1/extractions/run-1/fast/"; + const manifestObjectKey = `${prefix}manifest.json`; + storage.values.set(getWorkspacePageObjectKey(prefix, 1), "First"); + storage.values.set(getWorkspacePageObjectKey(prefix, 2), "Second"); + storage.values.set( + manifestObjectKey, + JSON.stringify({ + createdAt: new Date().toISOString(), + itemId: "item-1", + markdownBytes: 11, + markdownLength: 11, + metadata: {}, + pageCount: 2, + pages: [ + { markdownBytes: 5, pageNumber: 1 }, + { markdownBytes: 6, pageNumber: 2 }, + ], + provider: "liteparse", + providerMode: "fast", + runId: "run-1", + schemaVersion: 1, + sourceHash: "etag-1", + workspaceId: "workspace-1", + }), + ); + + await expect( + readWorkspacePageProjection({ + bucket: storage.bucket, + expectedSourceHash: "etag-1", + manifestObjectKey, + }), + ).resolves.toEqual({ + content: "## Page 1\n\nFirst", + emptyPages: [], + pages: { requested: "1", returned: [1], total: 2 }, }); - const { pages: _pages, ...manifestWithoutPages } = reference.manifest; + expect(storage.readKeys).toContain(getWorkspacePageObjectKey(prefix, 1)); + }); + + it("reads legacy manifests published before per-page sizes were added", async () => { + const storage = createObjectStorage(); + const prefix = "workspace_file_objects/workspace-1/item-1/extractions/run-1/fast/"; + const manifestObjectKey = `${prefix}manifest.json`; + storage.values.set(getWorkspacePageObjectKey(prefix, 1), "Page 1"); storage.values.set( - reference.manifestObjectKey, - JSON.stringify({ ...manifestWithoutPages, schemaVersion: 1 }), + manifestObjectKey, + JSON.stringify({ + createdAt: new Date().toISOString(), + itemId: "item-1", + markdownBytes: 6, + markdownLength: 6, + metadata: {}, + pageCount: 1, + provider: "liteparse", + providerMode: "fast", + runId: "run-1", + schemaVersion: 1, + sourceHash: "etag-1", + workspaceId: "workspace-1", + }), ); await expect( readWorkspacePageProjection({ bucket: storage.bucket, expectedSourceHash: "etag-1", - manifestObjectKey: reference.manifestObjectKey, + manifestObjectKey, }), ).resolves.toEqual({ content: "## Page 1\n\nPage 1", @@ -290,12 +336,16 @@ function createObjectStorage() { values.delete(key); } }, - async get(key: string) { + async get(key: string, options?: { range?: { offset: number; length: number } }) { readKeys.push(key); const value = values.get(key); if (value === undefined) { return null; } + const fullBytes = new TextEncoder().encode(value); + const bytes = options?.range + ? fullBytes.subarray(options.range.offset, options.range.offset + options.range.length) + : fullBytes; currentOpenBodies += 1; highestOpenBodies = Math.max(highestOpenBodies, currentOpenBodies); let consumed = false; @@ -308,20 +358,25 @@ function createObjectStorage() { return { body: { cancel: async () => consume() }, key, - size: new TextEncoder().encode(value).byteLength, + size: bytes.byteLength, + arrayBuffer: async () => { + consume(); + return bytes.slice().buffer; + }, text: async () => { consume(); - return value; + return new TextDecoder().decode(bytes); }, json: async () => { consume(); - return JSON.parse(value) as unknown; + return JSON.parse(new TextDecoder().decode(bytes)) as unknown; }, }; }, - async put(key: string, value: string) { - values.set(key, value); - return { key, size: new TextEncoder().encode(value).byteLength }; + async put(key: string, value: string | Blob) { + const text = typeof value === "string" ? value : await value.text(); + values.set(key, text); + return { key, size: new TextEncoder().encode(text).byteLength }; }, async list(input: { prefix?: string }) { const objects = Array.from(values.keys()) diff --git a/src/features/workspaces/extraction/workspace-page-projection.ts b/src/features/workspaces/extraction/workspace-page-projection.ts index d6afea258..374127679 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.ts @@ -12,33 +12,50 @@ import { } from "#/features/workspaces/read-page-selection"; import { deleteR2Prefix } from "#/lib/r2"; -const projectionSchemaVersion = 1; +// Version 1 stored one R2 object per page; version 2 stores every page concatenated +// in a single `pages.md` and serves individual pages with ranged reads, using the +// per-page byte counts the manifest already carries as the offset index. Old +// projections are derived data whose regeneration is billable, so both versions stay +// readable — v1 through the per-page path below, with no migration or backfill. +const projectionSchemaVersion = 2; const pageNumberWidth = 6; -const pageWriteConcurrency = 8; const maxPageMarkdownBytes = 1024 * 1024; const maxPageReadBytes = 2 * 1024 * 1024; +// Bounds what a projection may hold in total, which also bounds what the writer and +// the search indexer materialize in memory. Roughly seven times the densest document +// measured in production (a 1,527-page textbook at ~2.4 MB). +const maxProjectionMarkdownBytes = 16 * 1024 * 1024; const workspacePageProjectionManifestPageSchema = z.object({ markdownBytes: z.number().int().nonnegative(), pageNumber: z.number().int().positive(), }); -const workspacePageProjectionManifestSchema = z.object({ +const workspacePageProjectionManifestBaseSchema = z.object({ createdAt: z.string(), itemId: z.string(), markdownBytes: z.number().int().nonnegative(), markdownLength: z.number().int().nonnegative(), metadata: z.record(z.string(), jsonValueSchema), pageCount: z.number().int().positive(), - pages: z.array(workspacePageProjectionManifestPageSchema).optional(), provider: z.string(), providerMode: z.string(), runId: z.string(), - schemaVersion: z.literal(projectionSchemaVersion), sourceHash: z.string(), workspaceId: z.string(), }); +const workspacePageProjectionManifestSchema = z.discriminatedUnion("schemaVersion", [ + workspacePageProjectionManifestBaseSchema.extend({ + schemaVersion: z.literal(1), + pages: z.array(workspacePageProjectionManifestPageSchema).optional(), + }), + workspacePageProjectionManifestBaseSchema.extend({ + schemaVersion: z.literal(2), + pages: z.array(workspacePageProjectionManifestPageSchema), + }), +]); + type WorkspacePageProjectionManifest = z.infer; type WorkspacePageProjectionManifestPage = z.infer< @@ -59,7 +76,7 @@ export async function writeWorkspacePageProjection(input: { }) { const prefix = getWorkspacePageProjectionPrefix(input); const encoder = new TextEncoder(); - const writes: Promise[] = []; + const parts: Uint8Array[] = []; let lastPageNumber = 0; let markdownBytes = 0; let markdownLength = 0; @@ -73,31 +90,38 @@ export async function writeWorkspacePageProjection(input: { throw new Error("Extracted pages must be ordered by unique, increasing page number."); } + // Gaps become zero-byte spans: they keep page numbering stable without + // storing anything. for (let pageNumber = lastPageNumber + 1; pageNumber < page.pageNumber; pageNumber += 1) { - await schedulePageWrite(input.bucket, writes, prefix, pageNumber, ""); pages.push({ markdownBytes: 0, pageNumber }); } - const pageBytes = encoder.encode(page.markdown).byteLength; - if (pageBytes > maxPageMarkdownBytes) { + const pageBytes = encoder.encode(page.markdown); + if (pageBytes.byteLength > maxPageMarkdownBytes) { throw new Error(`Extracted page ${page.pageNumber} exceeds the page size limit.`); } + if (markdownBytes + pageBytes.byteLength > maxProjectionMarkdownBytes) { + throw new Error("Extracted document exceeds the projection size limit."); + } - await schedulePageWrite(input.bucket, writes, prefix, page.pageNumber, page.markdown); - pages.push({ markdownBytes: pageBytes, pageNumber: page.pageNumber }); + parts.push(pageBytes); + pages.push({ markdownBytes: pageBytes.byteLength, pageNumber: page.pageNumber }); lastPageNumber = page.pageNumber; - markdownBytes += pageBytes; + markdownBytes += pageBytes.byteLength; markdownLength += page.markdown.length; if (page.markdown.length > 0) { usablePageCount += 1; } } - await flushPageWrites(writes); if (lastPageNumber === 0 || usablePageCount === 0) { throw new Error("Extraction did not produce usable page Markdown."); } + await input.bucket.put(getWorkspacePagesObjectKey(prefix), new Blob(parts), { + httpMetadata: { contentType: "text/markdown; charset=utf-8" }, + }); + const manifest: WorkspacePageProjectionManifest = { createdAt: new Date().toISOString(), itemId: input.itemId, @@ -120,22 +144,11 @@ export async function writeWorkspacePageProjection(input: { return { manifest, manifestObjectKey }; } catch (error) { - const cleanupErrors: unknown[] = []; - try { - await flushPageWrites(writes); - } catch (cleanupError) { - cleanupErrors.push(cleanupError); - } - try { await deleteR2Prefix(input.bucket, prefix); } catch (cleanupError) { - cleanupErrors.push(cleanupError); - } - - if (cleanupErrors.length > 0) { throw new AggregateError( - [error, ...cleanupErrors], + [error, cleanupError], "Workspace page projection failed and cleanup did not complete.", { cause: error }, ); @@ -170,44 +183,28 @@ export async function readWorkspacePageProjection(input: { } const requested = input.pages?.trim() || "1"; const selectedPageNumbers = parseWorkspacePageRange(requested, manifest.pageCount); - const pageMetadataByNumber = manifest.pages - ? new Map(manifest.pages.map((page) => [page.pageNumber, page] as const)) - : null; - const selectedManifestBytes = pageMetadataByNumber - ? selectedPageNumbers.reduce( - (total, pageNumber) => - total + requireManifestPage(pageMetadataByNumber, pageNumber).markdownBytes, - 0, - ) + const selectedManifestBytes = manifest.pages + ? sumManifestPageBytes(manifest.pages, selectedPageNumbers) : null; if (selectedManifestBytes !== null && selectedManifestBytes > maxPageReadBytes) { throw new WorkspacePageSelectionError("page_selection_too_large"); } const prefix = getManifestPrefix(input.manifestObjectKey); - const pages: Array<{ markdown: string; pageNumber: number }> = []; - let totalBytes = 0; - - // Consume each R2 body before opening the next one; never retain a batch of live responses. - for (const pageNumber of selectedPageNumbers) { - const object = await getWorkspacePageProjectionObject({ - bucket: input.bucket, - pageMetadataByNumber, - pageNumber, - prefix, - }); - - totalBytes += object.size; - if (totalBytes > maxPageReadBytes) { - await object.body.cancel(); - throw new WorkspacePageSelectionError("page_selection_too_large"); - } - - pages.push({ - markdown: await object.text(), - pageNumber, - }); - } + const pages = + manifest.schemaVersion === 2 + ? await readPackedPages({ + bucket: input.bucket, + manifest, + prefix, + selectedPageNumbers, + }) + : await readLegacyPages({ + bucket: input.bucket, + manifest, + prefix, + selectedPageNumbers, + }); return { content: pages.map(formatProjectionPage).join("\n\n"), @@ -232,13 +229,39 @@ export async function* iterateWorkspacePageProjection(input: { throw new Error("Workspace page projection source does not match its published revision."); } + const prefix = getManifestPrefix(input.manifestObjectKey); + + if (manifest.schemaVersion === 2) { + // One read instead of one per page. The whole object is bounded by the write + // limit, so materializing it is cheaper than a thousand sequential round trips. + const object = await input.bucket.get(getWorkspacePagesObjectKey(prefix)); + if (!object) { + throw new Error("Workspace page projection content was not found."); + } + + const bytes = new Uint8Array(await object.arrayBuffer()); + if (bytes.byteLength !== manifest.markdownBytes) { + throw new Error("Workspace page projection content does not match its manifest."); + } + + const decoder = new TextDecoder(); + let offset = 0; + for (const page of manifest.pages) { + yield { + markdown: decoder.decode(bytes.subarray(offset, offset + page.markdownBytes)), + pageNumber: page.pageNumber, + }; + offset += page.markdownBytes; + } + return; + } + const pageMetadataByNumber = manifest.pages ? new Map(manifest.pages.map((page) => [page.pageNumber, page] as const)) : null; - const prefix = getManifestPrefix(input.manifestObjectKey); for (let pageNumber = 1; pageNumber <= manifest.pageCount; pageNumber += 1) { - const object = await getWorkspacePageProjectionObject({ + const object = await getLegacyPageObject({ bucket: input.bucket, pageMetadataByNumber, pageNumber, @@ -252,7 +275,112 @@ export async function* iterateWorkspacePageProjection(input: { } } -async function getWorkspacePageProjectionObject(input: { +/** + * Serves a page selection from the packed `pages.md` object. Contiguous selected + * pages coalesce into one ranged read; page offsets are the running sum of the + * per-page byte counts in the manifest. + */ +async function readPackedPages(input: { + bucket: R2Bucket; + manifest: Extract; + prefix: string; + selectedPageNumbers: number[]; +}) { + const spans = new Map(); + let offset = 0; + for (const page of input.manifest.pages) { + spans.set(page.pageNumber, { offset, length: page.markdownBytes }); + offset += page.markdownBytes; + } + + const markdownByPage = new Map(); + const decoder = new TextDecoder(); + const sorted = [...input.selectedPageNumbers].sort((left, right) => left - right); + + for (const run of coalesceContiguousRuns(sorted)) { + const runSpans = run.map((pageNumber) => { + const span = spans.get(pageNumber); + if (!span) { + throw new Error(`Workspace page projection manifest is missing page ${pageNumber}.`); + } + return span; + }); + const runLength = runSpans.reduce((total, span) => total + span.length, 0); + + if (runLength === 0) { + for (const pageNumber of run) { + markdownByPage.set(pageNumber, ""); + } + continue; + } + + const object = await input.bucket.get(getWorkspacePagesObjectKey(input.prefix), { + range: { offset: runSpans[0].offset, length: runLength }, + }); + if (!object) { + throw new Error("Workspace page projection content was not found."); + } + + const bytes = new Uint8Array(await object.arrayBuffer()); + if (bytes.byteLength !== runLength) { + throw new Error("Workspace page projection content does not match its manifest."); + } + + let runOffset = 0; + for (const [index, pageNumber] of run.entries()) { + const spanLength = runSpans[index].length; + markdownByPage.set( + pageNumber, + decoder.decode(bytes.subarray(runOffset, runOffset + spanLength)), + ); + runOffset += spanLength; + } + } + + return input.selectedPageNumbers.map((pageNumber) => ({ + markdown: markdownByPage.get(pageNumber) ?? "", + pageNumber, + })); +} + +async function readLegacyPages(input: { + bucket: R2Bucket; + manifest: Extract; + prefix: string; + selectedPageNumbers: number[]; +}) { + const pageMetadataByNumber = input.manifest.pages + ? new Map(input.manifest.pages.map((page) => [page.pageNumber, page] as const)) + : null; + const pages: Array<{ markdown: string; pageNumber: number }> = []; + let totalBytes = 0; + + // Consume each R2 body before opening the next one; never retain a batch of live + // responses. + for (const pageNumber of input.selectedPageNumbers) { + const object = await getLegacyPageObject({ + bucket: input.bucket, + pageMetadataByNumber, + pageNumber, + prefix: input.prefix, + }); + + totalBytes += object.size; + if (totalBytes > maxPageReadBytes) { + await object.body.cancel(); + throw new WorkspacePageSelectionError("page_selection_too_large"); + } + + pages.push({ + markdown: await object.text(), + pageNumber, + }); + } + + return pages; +} + +async function getLegacyPageObject(input: { bucket: R2Bucket; pageMetadataByNumber: ReadonlyMap | null; pageNumber: number; @@ -272,15 +400,37 @@ async function getWorkspacePageProjectionObject(input: { return object; } -function requireManifestPage( - pagesByNumber: ReadonlyMap, - pageNumber: number, +function sumManifestPageBytes( + pages: readonly WorkspacePageProjectionManifestPage[], + selectedPageNumbers: readonly number[], ) { - const page = pagesByNumber.get(pageNumber); - if (!page) { - throw new Error(`Workspace page projection manifest is missing page ${pageNumber}.`); + const bytesByPage = new Map(pages.map((page) => [page.pageNumber, page.markdownBytes] as const)); + + return selectedPageNumbers.reduce((total, pageNumber) => { + const bytes = bytesByPage.get(pageNumber); + if (bytes === undefined) { + throw new Error(`Workspace page projection manifest is missing page ${pageNumber}.`); + } + return total + bytes; + }, 0); +} + +function coalesceContiguousRuns(sortedPageNumbers: readonly number[]) { + const runs: number[][] = []; + let currentRun: number[] = []; + let previous: number | null = null; + + for (const pageNumber of sortedPageNumbers) { + if (previous !== null && pageNumber === previous + 1) { + currentRun.push(pageNumber); + } else { + currentRun = [pageNumber]; + runs.push(currentRun); + } + previous = pageNumber; } - return page; + + return runs; } async function readWorkspacePageProjectionManifest( @@ -304,6 +454,11 @@ function getWorkspacePageProjectionPrefix(input: { return `${getWorkspaceFileItemObjectPrefix(input)}extractions/${encodePathPart(input.runId)}/${input.tier}/`; } +function getWorkspacePagesObjectKey(prefix: string) { + return `${prefix}pages.md`; +} + +/** Key layout used by schema version 1, kept for projections published before v2. */ export function getWorkspacePageObjectKey(prefix: string, pageNumber: number) { return `${prefix}pages/${String(pageNumber).padStart(pageNumberWidth, "0")}.md`; } @@ -318,35 +473,13 @@ function parseWorkspacePageProjectionManifest(value: unknown): WorkspacePageProj throw new Error("Workspace page projection manifest is invalid."); } } - return manifest; -} - -async function schedulePageWrite( - bucket: R2Bucket, - writes: Promise[], - prefix: string, - pageNumber: number, - markdown: string, -) { - const write = bucket - .put(getWorkspacePageObjectKey(prefix, pageNumber), markdown, { - httpMetadata: { contentType: "text/markdown; charset=utf-8" }, - }) - .then(() => undefined); - writes.push(write); - - if (writes.length >= pageWriteConcurrency) { - await flushPageWrites(writes); - } -} - -async function flushPageWrites(writes: Promise[]) { - const results = await Promise.allSettled(writes.splice(0)); - const failure = results.find((result) => result.status === "rejected"); - - if (failure?.status === "rejected") { - throw failure.reason; + if ( + manifest.schemaVersion === 2 && + manifest.pages.reduce((total, page) => total + page.markdownBytes, 0) !== manifest.markdownBytes + ) { + throw new Error("Workspace page projection manifest is invalid."); } + return manifest; } function normalizeProjectionPage(page: MarkdownProjectionPage): MarkdownProjectionPage { diff --git a/src/features/workspaces/extraction/workspace-projection-readiness.test.ts b/src/features/workspaces/extraction/workspace-projection-readiness.test.ts index 3670dbe93..323719ddc 100644 --- a/src/features/workspaces/extraction/workspace-projection-readiness.test.ts +++ b/src/features/workspaces/extraction/workspace-projection-readiness.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest"; +import { workspaceExtractionStallThresholdMs } from "#/features/workspaces/extraction/workspace-extraction-budgets"; import { resolveWorkspaceProjectionReadiness } from "#/features/workspaces/extraction/workspace-projection-readiness"; import type { ReadWorkspaceKernelFileProjectionResult } from "#/features/workspaces/kernel/workspace-kernel-types"; @@ -64,21 +65,26 @@ describe("resolveWorkspaceProjectionReadiness", () => { }); it("stalls a processing projection that outlived the retrying extraction budget", () => { + const elapsedMs = workspaceExtractionStallThresholdMs + 60_000; const projection = createProjection({ status: "processing", - updatedAt: new Date(now - 46 * 60_000).toISOString(), + updatedAt: new Date(now - elapsedMs).toISOString(), }); expect(resolveWorkspaceProjectionReadiness(projection, now)).toEqual({ state: "stalled", - elapsedSeconds: 46 * 60, + elapsedSeconds: elapsedMs / 1000, }); }); + // Derived rather than hardcoded: the threshold is the sum of every step budget, so + // a literal here would silently stop testing the boundary the moment a timeout + // moves — and calling a healthy run stalled makes the reconciler queue a duplicate + // workflow, and a duplicate bill, against a document that is still parsing. it("keeps a slow but healthy extraction pending rather than stalling it", () => { const projection = createProjection({ status: "processing", - updatedAt: new Date(now - 31 * 60_000).toISOString(), + updatedAt: new Date(now - (workspaceExtractionStallThresholdMs - 60_000)).toISOString(), }); expect(resolveWorkspaceProjectionReadiness(projection, now)).toMatchObject({ diff --git a/src/features/workspaces/extraction/workspace-projection-readiness.ts b/src/features/workspaces/extraction/workspace-projection-readiness.ts index 2a4bdc7c0..c06399c2e 100644 --- a/src/features/workspaces/extraction/workspace-projection-readiness.ts +++ b/src/features/workspaces/extraction/workspace-projection-readiness.ts @@ -1,15 +1,6 @@ +import { workspaceExtractionStallThresholdMs } from "#/features/workspaces/extraction/workspace-extraction-budgets"; import type { ReadWorkspaceKernelFileProjectionResult } from "#/features/workspaces/kernel/workspace-kernel-types"; -/** - * How long a projection may sit in `processing` before it is treated as stalled. - * - * The enhanced extraction step allows a 10 minute timeout across 3 attempts with - * exponential backoff, and the projection row is not touched between attempts, so - * a healthy run can legitimately stay `processing` for a little over half an hour. - * The threshold sits above that worst case so slow retries are never mislabelled. - */ -export const workspaceExtractionStallThresholdMs = 45 * 60_000; - const minimumRetryAfterSeconds = 15; const maximumRetryAfterSeconds = 120; diff --git a/src/features/workspaces/files/workspace-file-processor.ts b/src/features/workspaces/files/workspace-file-processor.ts index 00d90ed2d..317705328 100644 --- a/src/features/workspaces/files/workspace-file-processor.ts +++ b/src/features/workspaces/files/workspace-file-processor.ts @@ -1,9 +1,13 @@ import { Container, getRandom } from "@cloudflare/containers"; +import { awaitUploadResponse } from "#/lib/http/streaming-upload"; + const workspaceFileProcessorPort = 8080; const workspaceFileProcessorPoolSize = 2; +// Each of these must stay under the workflow step that wraps it, otherwise the step +// dies first and the abort never fires — trading a named error for an opaque timeout. const processorRequestTimeoutMs = { - "/parse/pdf": 10 * 60_000, + "/parse/pdf": 4 * 60_000, "/prepare/pdf": 2 * 60_000, "/preview/image": 2 * 60_000, } as const; @@ -40,7 +44,8 @@ export async function requestWorkspaceFileProcessor( } const body = new FixedLengthStream(input.sizeBytes); - const [response] = await Promise.all([ + + return await awaitUploadResponse( processor.fetch( new Request(`http://workspace-file-processor${input.path}`, { body: body.readable, @@ -51,7 +56,5 @@ export async function requestWorkspaceFileProcessor( } as RequestInit & { duplex: "half" }), ), input.body.pipeTo(body.writable), - ]); - - return response; + ); } diff --git a/src/features/workspaces/kernel/workspace-kernel.ts b/src/features/workspaces/kernel/workspace-kernel.ts index a3d6b8a34..c36e7cd0e 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -542,6 +542,7 @@ export class WorkspaceKernel extends Agent { this.lastExtractionHealingRequestAt = now; this.ctx.waitUntil( reconcileWorkspaceFileExtractions({ + schedule: (task) => this.ctx.waitUntil(task), sql: this.kernelSql, workflow: this.env.WORKSPACE_FILE_EXTRACTION_WORKFLOW, workspaceId: this.name, diff --git a/src/features/workspaces/operations/read-items.ts b/src/features/workspaces/operations/read-items.ts index 295812604..3676c21a6 100644 --- a/src/features/workspaces/operations/read-items.ts +++ b/src/features/workspaces/operations/read-items.ts @@ -5,6 +5,7 @@ import { type WorkspaceReadItemsOutput, } from "#/features/workspaces/content/workspace-content-contract"; import { readWorkspaceContent } from "#/features/workspaces/content/workspace-content-reader"; +import { recordWorkspaceFileReadOutcomes } from "#/features/workspaces/content/workspace-read-observability"; import { createWorkspaceReadReferences } from "#/features/workspaces/content/workspace-read-references"; import { getDocumentSessionFromEnv } from "#/features/workspaces/document-session-access"; import type { WorkspaceAccessContext } from "#/features/workspaces/operations/workspace-access-context"; @@ -33,6 +34,13 @@ export async function readWorkspaceItemsOperation( requests: input.requests, }); + recordWorkspaceFileReadOutcomes({ + operationId: accessContext.operationId, + results, + userId: accessContext.actor.userId, + workspaceId: accessContext.workspaceId, + }); + return { references: createWorkspaceReadReferences(results), results, diff --git a/src/integrations/llamaparse/client.ts b/src/integrations/llamaparse/client.ts index b2f385ab3..84f0c037e 100644 --- a/src/integrations/llamaparse/client.ts +++ b/src/integrations/llamaparse/client.ts @@ -83,3 +83,8 @@ export function getNumberValue(value: unknown, key: string) { const field = getRecordValue(value, key); return typeof field === "number" ? field : null; } + +export function getBooleanValue(value: unknown, key: string) { + const field = getRecordValue(value, key); + return typeof field === "boolean" ? field : null; +} diff --git a/src/integrations/posthog/events.ts b/src/integrations/posthog/events.ts index cdfa49f1f..80e6ed652 100644 --- a/src/integrations/posthog/events.ts +++ b/src/integrations/posthog/events.ts @@ -41,6 +41,30 @@ export interface PostHogEventPropertiesByName { share_method: "link"; shared_role: WorkspaceMembershipRole; }; + workspace_file_read_completed: { + workspace_id: string; + operation_id: string; + item_id: string | null; + status: "ready" | "pending" | "failed"; + /** Extraction phase when the read found the projection still pending. */ + phase: "queued" | "extracting" | null; + failure_code: string | null; + /** True when the content served came from the fast pass and may still improve. */ + provisional: boolean | null; + empty_page_count: number | null; + returned_page_count: number | null; + elapsed_seconds: number | null; + }; + workspace_file_extraction_healing_enqueued: { + workspace_id: string; + /** Candidates enqueued this sweep; duplicate workflow ids are skipped downstream. */ + total: number; + missing: number; + failed: number; + stalled: number; + provisional: number; + unreadable: number; + }; ai_turn_started: { thread_id: string; workspace_id: string; @@ -88,6 +112,10 @@ export interface PostHogEventPropertiesByName { credits_used: number | null; duration_ms: number; enhancement_duration_ms: number; + /** Time the provider job sat queued before running, when the provider reports it. */ + enhancement_queued_ms: number | null; + /** Time the provider job spent actually parsing, when the provider reports it. */ + enhancement_parse_ms: number | null; enhancement_error_message: string | null; enhancement_error_type: string | null; enhancement_outcome: "error" | "success"; diff --git a/src/lib/http/streaming-multipart.ts b/src/lib/http/streaming-multipart.ts index e9cce862d..679617ed3 100644 --- a/src/lib/http/streaming-multipart.ts +++ b/src/lib/http/streaming-multipart.ts @@ -1,3 +1,5 @@ +import { awaitUploadResponse } from "#/lib/http/streaming-upload"; + export function createStreamingMultipartFile(input: { body: ReadableStream; contentType: string; @@ -21,7 +23,11 @@ export function createStreamingMultipartFile(input: { return { body: stream.readable, contentType: `multipart/form-data; boundary=${boundary}`, - done, + // Callers must await their request through this rather than awaiting the body + // pump alongside it. See awaitUploadResponse for why. + awaitResponse(response: Promise): Promise { + return awaitUploadResponse(response, done); + }, }; } diff --git a/src/lib/http/streaming-multipart.worker.test.ts b/src/lib/http/streaming-multipart.worker.test.ts new file mode 100644 index 000000000..45db74d6f --- /dev/null +++ b/src/lib/http/streaming-multipart.worker.test.ts @@ -0,0 +1,46 @@ +import { describe, expect, it } from "vitest"; + +import { createStreamingMultipartFile } from "#/lib/http/streaming-multipart"; + +function createMultipart(body: ReadableStream) { + return createStreamingMultipartFile({ + body, + contentType: "application/pdf", + fileName: "source.pdf", + formFieldName: "file", + sizeBytes: 1024, + }); +} + +describe("createStreamingMultipartFile", () => { + it("resolves from the response when the body pump never finishes", async () => { + // Nothing ever drains multipart.body, so the pump parks forever — the shape of a + // server that answers with a 4xx before it has read the upload. Awaiting the pump + // alongside the response here used to hang until the workflow step timed out. + const multipart = createMultipart( + new ReadableStream({ + pull: () => new Promise(() => {}), + }), + ); + + await expect(multipart.awaitResponse(Promise.resolve("quota exceeded"))).resolves.toBe( + "quota exceeded", + ); + }); + + it("surfaces a body pump failure even when the response never settles", async () => { + const multipart = createMultipart( + new ReadableStream({ + pull: (controller) => controller.error(new Error("source stream failed")), + }), + ); + // Stand in for the fetch that would be reading the request body. Without a + // consumer the pump parks on its first write and never reaches the source error. + const drained = multipart.body.pipeTo(new WritableStream()).catch(() => undefined); + + await expect(multipart.awaitResponse(new Promise(() => {}))).rejects.toThrow( + "source stream failed", + ); + await drained; + }); +}); diff --git a/src/lib/http/streaming-upload.ts b/src/lib/http/streaming-upload.ts new file mode 100644 index 000000000..0302a4ea2 --- /dev/null +++ b/src/lib/http/streaming-upload.ts @@ -0,0 +1,17 @@ +/** + * Awaits a request whose body is still being pumped into it. + * + * An endpoint that answers before draining the body — a 4xx, a redirect, a quota + * rejection — leaves the writer parked on backpressure that never clears, so awaiting + * the pump alongside the response hangs until something further up times out. The + * response decides when the upload is over. + * + * While the request is still consuming the body, a pump failure surfaces by breaking + * the body and failing the request. Once the response has settled, the rest of the + * pump is abandoned unobserved — safe here because a settled response means the + * server already decided without the remaining bytes, and retries build fresh + * streams rather than reusing this one. + */ +export function awaitUploadResponse(response: Promise, body: Promise): Promise { + return Promise.race([response, body.then(() => response)]); +}