feat(compute): integrate worker operation state polling for document parsing
Add worker operation state polling to document parse routes, enabling real-time status and progress updates from the ONNX layout worker. Introduce `fetchWorkerOperationState` utility and status mapping logic to synchronize parse status with worker job state, improving accuracy for pending, running, and failed operations.
This commit is contained in:
parent
bfea95aa35
commit
6f68c169b6
3 changed files with 101 additions and 5 deletions
|
|
@ -4,6 +4,7 @@ import { db } from '@/db';
|
||||||
import { documents } from '@/db/schema';
|
import { documents } from '@/db/schema';
|
||||||
import { requireAuthContext } from '@/lib/server/auth/auth';
|
import { requireAuthContext } from '@/lib/server/auth/auth';
|
||||||
import { getWorkerClientConfigFromEnv } from '@/lib/server/compute/worker';
|
import { getWorkerClientConfigFromEnv } from '@/lib/server/compute/worker';
|
||||||
|
import { fetchWorkerOperationState } from '@/lib/server/compute/worker-op-state';
|
||||||
import { isValidDocumentId } from '@/lib/server/documents/blobstore';
|
import { isValidDocumentId } from '@/lib/server/documents/blobstore';
|
||||||
import { normalizeParseStatus, parseDocumentParseState } from '@/lib/server/documents/parse-state';
|
import { normalizeParseStatus, parseDocumentParseState } from '@/lib/server/documents/parse-state';
|
||||||
import { healStaleDocumentParseState } from '@/lib/server/documents/parse-state-healing';
|
import { healStaleDocumentParseState } from '@/lib/server/documents/parse-state-healing';
|
||||||
|
|
@ -76,12 +77,22 @@ async function toSnapshotState(row: ParseRow): Promise<SnapshotState> {
|
||||||
state: parseDocumentParseState(row.parseState),
|
state: parseDocumentParseState(row.parseState),
|
||||||
});
|
});
|
||||||
const parseStatus = normalizeParseStatus(state.status);
|
const parseStatus = normalizeParseStatus(state.status);
|
||||||
|
const opId = typeof state.opId === 'string' && state.opId.trim() ? state.opId.trim() : null;
|
||||||
|
if (opId && parseStatus !== 'ready') {
|
||||||
|
const workerState = await fetchWorkerOperationState<PdfLayoutJobResult>(opId);
|
||||||
|
if (workerState && workerState.opId === opId) {
|
||||||
|
return {
|
||||||
|
snapshot: snapshotFromWorkerState(workerState),
|
||||||
|
opId,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
return {
|
return {
|
||||||
snapshot: {
|
snapshot: {
|
||||||
parseStatus,
|
parseStatus,
|
||||||
parseProgress: state.progress ?? null,
|
parseProgress: state.progress ?? null,
|
||||||
},
|
},
|
||||||
opId: typeof state.opId === 'string' && state.opId.trim() ? state.opId.trim() : null,
|
opId,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ import { and, eq, inArray } from 'drizzle-orm';
|
||||||
import { db } from '@/db';
|
import { db } from '@/db';
|
||||||
import { documents } from '@/db/schema';
|
import { documents } from '@/db/schema';
|
||||||
import { requireAuthContext } from '@/lib/server/auth/auth';
|
import { requireAuthContext } from '@/lib/server/auth/auth';
|
||||||
|
import { fetchWorkerOperationState } from '@/lib/server/compute/worker-op-state';
|
||||||
import {
|
import {
|
||||||
getParsedDocumentBlob,
|
getParsedDocumentBlob,
|
||||||
getParsedDocumentBlobByKey,
|
getParsedDocumentBlobByKey,
|
||||||
|
|
@ -20,6 +21,7 @@ import { healStaleDocumentParseState } from '@/lib/server/documents/parse-state-
|
||||||
import { getOpenReaderTestNamespace, getUnclaimedUserIdForNamespace } from '@/lib/server/testing/test-namespace';
|
import { getOpenReaderTestNamespace, getUnclaimedUserIdForNamespace } from '@/lib/server/testing/test-namespace';
|
||||||
import { isS3Configured } from '@/lib/server/storage/s3';
|
import { isS3Configured } from '@/lib/server/storage/s3';
|
||||||
import type { ParsedPdfDocument } from '@/types/parsed-pdf';
|
import type { ParsedPdfDocument } from '@/types/parsed-pdf';
|
||||||
|
import type { PdfLayoutJobResult, WorkerOperationState } from '@openreader/compute-core/api-contracts';
|
||||||
|
|
||||||
export const dynamic = 'force-dynamic';
|
export const dynamic = 'force-dynamic';
|
||||||
|
|
||||||
|
|
@ -35,6 +37,21 @@ function hasAnyParsedBlocks(doc: ParsedPdfDocument | null): boolean {
|
||||||
return doc.pages.some((page) => Array.isArray(page.blocks) && page.blocks.length > 0);
|
return doc.pages.some((page) => Array.isArray(page.blocks) && page.blocks.length > 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function mapWorkerStatusToParseStatus(status: WorkerOperationState['status']) {
|
||||||
|
switch (status) {
|
||||||
|
case 'queued':
|
||||||
|
return 'pending' as const;
|
||||||
|
case 'running':
|
||||||
|
return 'running' as const;
|
||||||
|
case 'succeeded':
|
||||||
|
return 'ready' as const;
|
||||||
|
case 'failed':
|
||||||
|
return 'failed' as const;
|
||||||
|
default:
|
||||||
|
return 'pending' as const;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string }> }) {
|
export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string }> }) {
|
||||||
try {
|
try {
|
||||||
if (!isS3Configured()) return s3NotConfiguredResponse();
|
if (!isS3Configured()) return s3NotConfiguredResponse();
|
||||||
|
|
@ -80,7 +97,20 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string
|
||||||
userId: row.userId,
|
userId: row.userId,
|
||||||
state,
|
state,
|
||||||
});
|
});
|
||||||
const effectiveStatus = normalizeParseStatus(state.status);
|
let effectiveStatus = normalizeParseStatus(state.status);
|
||||||
|
let effectiveProgress = state.progress ?? null;
|
||||||
|
const opId = typeof state.opId === 'string' ? state.opId.trim() : '';
|
||||||
|
if (opId && effectiveStatus !== 'ready') {
|
||||||
|
const workerState = await fetchWorkerOperationState<PdfLayoutJobResult>(opId);
|
||||||
|
if (workerState && workerState.opId === opId) {
|
||||||
|
const workerStatus = mapWorkerStatusToParseStatus(workerState.status);
|
||||||
|
// Keep DB/blob as source of truth for "ready"; prefer worker only for active/failed states.
|
||||||
|
if (workerStatus === 'pending' || workerStatus === 'running' || workerStatus === 'failed') {
|
||||||
|
effectiveStatus = workerStatus;
|
||||||
|
effectiveProgress = workerStatus === 'running' ? (workerState.progress ?? null) : null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if (effectiveStatus === 'failed' && retryFailed) {
|
if (effectiveStatus === 'failed' && retryFailed) {
|
||||||
await db
|
await db
|
||||||
|
|
@ -104,7 +134,7 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string
|
||||||
if (effectiveStatus !== 'ready') {
|
if (effectiveStatus !== 'ready') {
|
||||||
return NextResponse.json({
|
return NextResponse.json({
|
||||||
parseStatus: effectiveStatus,
|
parseStatus: effectiveStatus,
|
||||||
parseProgress: state.progress ?? null,
|
parseProgress: effectiveProgress,
|
||||||
}, { status: 202 });
|
}, { status: 202 });
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -190,7 +220,19 @@ export async function POST(req: NextRequest, ctx: { params: Promise<{ id: string
|
||||||
userId: row.userId,
|
userId: row.userId,
|
||||||
state,
|
state,
|
||||||
});
|
});
|
||||||
const effectiveStatus = normalizeParseStatus(state.status);
|
let effectiveStatus = normalizeParseStatus(state.status);
|
||||||
|
let effectiveProgress = state.progress ?? null;
|
||||||
|
const opId = typeof state.opId === 'string' ? state.opId.trim() : '';
|
||||||
|
if (opId && effectiveStatus !== 'ready') {
|
||||||
|
const workerState = await fetchWorkerOperationState<PdfLayoutJobResult>(opId);
|
||||||
|
if (workerState && workerState.opId === opId) {
|
||||||
|
const workerStatus = mapWorkerStatusToParseStatus(workerState.status);
|
||||||
|
if (workerStatus === 'pending' || workerStatus === 'running' || workerStatus === 'failed') {
|
||||||
|
effectiveStatus = workerStatus;
|
||||||
|
effectiveProgress = workerStatus === 'running' ? (workerState.progress ?? null) : null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if (effectiveStatus !== 'running') {
|
if (effectiveStatus !== 'running') {
|
||||||
await db
|
await db
|
||||||
|
|
@ -215,7 +257,7 @@ export async function POST(req: NextRequest, ctx: { params: Promise<{ id: string
|
||||||
return NextResponse.json(
|
return NextResponse.json(
|
||||||
{
|
{
|
||||||
parseStatus: effectiveStatus === 'running' ? 'running' : 'pending',
|
parseStatus: effectiveStatus === 'running' ? 'running' : 'pending',
|
||||||
parseProgress: effectiveStatus === 'running' ? (state.progress ?? null) : null,
|
parseProgress: effectiveStatus === 'running' ? effectiveProgress : null,
|
||||||
},
|
},
|
||||||
{ status: 202 },
|
{ status: 202 },
|
||||||
);
|
);
|
||||||
|
|
|
||||||
43
src/lib/server/compute/worker-op-state.ts
Normal file
43
src/lib/server/compute/worker-op-state.ts
Normal file
|
|
@ -0,0 +1,43 @@
|
||||||
|
import { getWorkerClientConfigFromEnv } from '@/lib/server/compute/worker';
|
||||||
|
import type { WorkerOperationState } from '@openreader/compute-core/api-contracts';
|
||||||
|
|
||||||
|
const WORKER_OP_REQUEST_TIMEOUT_MS = 2_500;
|
||||||
|
|
||||||
|
export async function fetchWorkerOperationState<Result>(
|
||||||
|
opId: string | null | undefined,
|
||||||
|
): Promise<WorkerOperationState<Result> | null> {
|
||||||
|
const normalized = opId?.trim();
|
||||||
|
if (!normalized) return null;
|
||||||
|
|
||||||
|
let cfg: { baseUrl: string; token: string };
|
||||||
|
try {
|
||||||
|
cfg = getWorkerClientConfigFromEnv();
|
||||||
|
} catch {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
const controller = new AbortController();
|
||||||
|
const timeout = setTimeout(() => controller.abort(), WORKER_OP_REQUEST_TIMEOUT_MS);
|
||||||
|
|
||||||
|
try {
|
||||||
|
const res = await fetch(`${cfg.baseUrl}/ops/${encodeURIComponent(normalized)}`, {
|
||||||
|
method: 'GET',
|
||||||
|
headers: {
|
||||||
|
Authorization: `Bearer ${cfg.token}`,
|
||||||
|
Accept: 'application/json',
|
||||||
|
},
|
||||||
|
cache: 'no-store',
|
||||||
|
signal: controller.signal,
|
||||||
|
});
|
||||||
|
|
||||||
|
if (!res.ok) return null;
|
||||||
|
const parsed = await res.json() as WorkerOperationState<Result>;
|
||||||
|
if (!parsed || typeof parsed !== 'object' || parsed.opId !== normalized) return null;
|
||||||
|
return parsed;
|
||||||
|
} catch {
|
||||||
|
return null;
|
||||||
|
} finally {
|
||||||
|
clearTimeout(timeout);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Loading…
Reference in a new issue