@farukada/aws-langgraph-dynamodb-ts 1.0.0-rc.2 → 1.0.0

package/dist/esm/checkpointer/internal/pending-writes.js

@@ -0,0 +1,229 @@
+import { collectS3Keys } from '../../shared/codec/codec.js';
+import { cleanUpS3Orphans } from '../../shared/codec/s3/offloader.js';
+import { mapWithConcurrency } from '../../shared/concurrency.js';
+import { abortErrorFrom } from '../../shared/dynamodb/abort.js';
+import { commitRow, isConditionalCheckFailed, OVERWRITE_CAS_MAX_ATTEMPTS, readRow, rejectedRow, revisionGuard, settledVerdict, verdictFor, verifyRow, } from '../../shared/dynamodb/idempotent-write.js';
+import { withDynamoDBRetry, retryFor } from '../../shared/dynamodb/retry.js';
+import { PARTITION_KEY_ATTRIBUTE, rowKeyOf } from '../../shared/dynamodb/table-schema.js';
+import { truncateForLog } from '../../shared/logging/truncate.js';
+import { WRITE_GROUP_ATTRIBUTE } from './rows.js';
+export async function commitPendingWrites(context, batch) {
+    const special = batch.items.filter((item) => item.index < 0);
+    const regular = batch.items.filter((item) => item.index >= 0);
+    const [specialError, regularOutcome] = await Promise.all([
+        writeSpecialRowsWithCleanup(context, batch.threadId, special, batch.signal),
+        writeRegularRows(context, regular, batch.signal),
+    ]);
+    await releaseDeadUploads(context, regularOutcome.deadUploads);
+    const firstError = specialError ?? regularOutcome.error;
+    if (firstError)
+        throw firstError;
+}
+async function releaseDeadUploads(context, dead) {
+    if (!context.offloader)
+        return;
+    await cleanUpS3Orphans(context.offloader, {
+        keys: collectS3Keys(dead.map((item) => item.value)),
+        operation: 'putWrites',
+        logger: context.logger,
+    });
+}
+async function commitSpecialRow(context, item, guard, signal) {
+    await commitRow(context, item, item.value, { guard, signal });
+}
+async function attemptCasWrites(context, item, initial, signal) {
+    let observed = initial;
+    for (let attempt = 1; attempt <= OVERWRITE_CAS_MAX_ATTEMPTS; attempt++) {
+        if (signal?.aborted) {
+            return { done: true, outcome: { committed: false, error: abortErrorFrom(signal) } };
+        }
+        const attempted = observed;
+        try {
+            await commitSpecialRow(context, item, revisionGuard(WRITE_GROUP_ATTRIBUTE, attempted), signal);
+            return { done: true, outcome: { committed: true, superseded: attempted.value } };
+        }
+        catch (error) {
+            const verified = await verifyAfterFailure(context, item, attempted, error);
+            if (!verified.observed || !isConditionalCheckFailed(error)) {
+                return { done: true, outcome: verified.outcome };
+            }
+            observed = verified.observed;
+        }
+    }
+    return { done: false, observed };
+}
+async function overwriteUnconditionally(context, item, observed, signal) {
+    if (signal?.aborted)
+        return { committed: false, error: abortErrorFrom(signal) };
+    try {
+        await commitSpecialRow(context, item, undefined, signal);
+        return { committed: true, superseded: observed.value };
+    }
+    catch (error) {
+        return (await verifyAfterFailure(context, item, observed, error, false)).outcome;
+    }
+}
+async function writeWithoutOffloader(context, item, signal) {
+    try {
+        await withDynamoDBRetry((request) => context.client.put({ TableName: context.tableName, Item: item }, request), retryFor(context, signal));
+        return { committed: true };
+    }
+    catch (error) {
+        return { committed: false, error: error };
+    }
+}
+export async function writeSpecialRow(context, item, signal) {
+    if (!context.offloader)
+        return writeWithoutOffloader(context, item, signal);
+    try {
+        const initial = await readSpecialRow(context, item);
+        const attempt = await attemptCasWrites(context, item, initial, signal);
+        if (attempt.done)
+            return attempt.outcome;
+        context.logger.warn('putWrites: special-write compare-and-swap exhausted; overwriting unconditionally, which ' +
+            'can orphan one S3 object under a concurrent call (reclaimed by ensureS3LifecycleRule)', { sortKey: item.SK, channel: item.channel, attempts: OVERWRITE_CAS_MAX_ATTEMPTS });
+        return await overwriteUnconditionally(context, item, attempt.observed, signal);
+    }
+    catch (error) {
+        return { committed: true, error: error };
+    }
+}
+export function writeRowProbe(item) {
+    return {
+        key: rowKeyOf(item),
+        kind: 'attribute',
+        attribute: WRITE_GROUP_ATTRIBUTE,
+        expected: item.writeGroup,
+        also: ['value'],
+    };
+}
+function stateOf(row) {
+    if (!row)
+        return { exists: false };
+    return {
+        exists: true,
+        value: row.value,
+        revision: row[WRITE_GROUP_ATTRIBUTE],
+    };
+}
+export async function readSpecialRow(context, item) {
+    return stateOf(await readRow(context, writeRowProbe(item)));
+}
+export async function verifyAfterFailure(context, item, attempted, error, pinned = true) {
+    const probe = writeRowProbe(item);
+    const rejected = isConditionalCheckFailed(error) ? rejectedRow(error) : undefined;
+    const read = rejected
+        ? { verdict: verdictFor(probe, rejected), row: rejected }
+        : await verifyRow(context, probe);
+    const verdict = rejected === undefined && (!pinned || stillPinned(attempted, read.row))
+        ? settledVerdict(read.verdict, error)
+        : read.verdict;
+    if (verdict === 'landed')
+        return { outcome: { committed: true, superseded: attempted.value } };
+    if (verdict === 'unverified')
+        return { outcome: { committed: true, error } };
+    return { outcome: { committed: false, error }, observed: stateOf(read.row) };
+}
+function stillPinned(attempted, row) {
+    const current = stateOf(row);
+    if (current.revision !== attempted.revision)
+        return false;
+    if (current.exists === attempted.exists)
+        return true;
+    return !current.exists && attempted.revision === undefined;
+}
+async function deleteDescriptors(context, descriptors, label, scope) {
+    if (!context.offloader)
+        return;
+    await cleanUpS3Orphans(context.offloader, {
+        keys: collectS3Keys(descriptors.filter((ref) => Boolean(ref))),
+        operation: label,
+        logger: context.logger,
+        ...(scope === undefined ? {} : { scope }),
+    });
+}
+export async function writeSpecialRowsWithCleanup(context, threadId, items, signal) {
+    if (items.length === 0)
+        return undefined;
+    const outcomes = await Promise.all(items.map(async (item) => [
+        item,
+        await writeSpecialRow(context, item, signal),
+    ]));
+    await deleteDescriptors(context, outcomes.filter(([, o]) => o.committed).map(([, o]) => o.superseded), 'putWrites.special.previous', [threadId]);
+    await deleteDescriptors(context, outcomes.filter(([, o]) => !o.committed).map(([item]) => item.value), 'putWrites.special.newUpload');
+    return outcomes.find(([, o]) => o.error)?.[1].error;
+}
+const FIRST_WRITE_WINS = {
+    ConditionExpression: `attribute_not_exists(${PARTITION_KEY_ATTRIBUTE})`,
+    ReturnValuesOnConditionCheckFailure: 'ALL_OLD',
+};
+async function commitWriteRow(context, item, signal) {
+    await commitRow(context, item, item.value, { guard: FIRST_WRITE_WINS, signal });
+}
+async function verifyFailure(context, item) {
+    if (!context.offloader)
+        return { verdict: 'not-landed' };
+    return verifyRow(context, writeRowProbe(item));
+}
+export const REGULAR_WRITE_CONCURRENCY = 32;
+export async function writeRegularRows(context, items, signal) {
+    const settled = await mapWithConcurrency(items, REGULAR_WRITE_CONCURRENCY, async (item) => {
+        if (signal?.aborted)
+            return { ok: false, reason: abortErrorFrom(signal), sent: false };
+        try {
+            await commitWriteRow(context, item, signal);
+            return { ok: true };
+        }
+        catch (error) {
+            return { ok: false, reason: error, sent: true };
+        }
+    });
+    const outcome = { deadUploads: [] };
+    for (const [index, result] of settled.entries()) {
+        if (result.ok)
+            continue;
+        const item = items[index];
+        const { reason } = result;
+        if (!result.sent) {
+            outcome.deadUploads.push(item);
+            outcome.error = outcome.error ?? reason;
+            continue;
+        }
+        if (isConditionalCheckFailed(reason)) {
+            reportGuardRejection(context, item, reason);
+            if (rejectionProvesForeignRow(item, reason))
+                outcome.deadUploads.push(item);
+            continue;
+        }
+        const { verdict, row } = await verifyFailure(context, item);
+        const settledNow = row === undefined ? settledVerdict(verdict, reason) : verdict;
+        if (settledNow === 'landed')
+            continue;
+        if (settledNow === 'not-landed')
+            outcome.deadUploads.push(item);
+        outcome.error = outcome.error ?? reason;
+    }
+    return outcome;
+}
+function rejectedChannel(error) {
+    return rejectedRow(error)?.channel;
+}
+export function reportGuardRejection(context, item, error) {
+    const found = rejectedChannel(error);
+    if (found !== undefined && found !== item.channel) {
+        context.logger.warn('putWrites: write row held by an unexpected channel; write not persisted', {
+            sortKey: item.SK,
+            expected: item.channel,
+            found: truncateForLog(found),
+        });
+        return;
+    }
+    context.logger.debug('putWrites: skipped a write already committed for this task and channel', {
+        sortKey: item.SK,
+        channel: item.channel,
+    });
+}
+export function rejectionProvesForeignRow(item, error) {
+    const group = rejectedRow(error)?.writeGroup;
+    return group !== undefined && group !== item.writeGroup;
+}

package/dist/esm/checkpointer/internal/read.d.ts

@@ -0,0 +1,135 @@
+/**
+ * Hides which rows a checkpoint read issues, with what consistency, and how a
+ * checkpoint an older release wrote is read back.
+ *
+ * A tuple is a META row, a PAYLOAD row and a checkpoint's WRITE rows, read
+ * strongly for `getTuple` and eventually for `list`, with expired rows absent
+ * however long DynamoDB's sweep lags, and with a pre-v4 checkpoint's `Send`s
+ * migrated from its parent's pending writes. The caller asks for a tuple; which
+ * reads that takes, and in what order, is decided here.
+ */
+import { type Checkpoint, type CheckpointMetadata, type CheckpointPendingWrite, type CheckpointTuple } from '@langchain/langgraph-checkpoint';
+import type { ThreadAddress } from './parse.js';
+import { type CheckpointLocation, type CheckpointMetaRow, type CheckpointPayloadRow, type MetaReader } from './rows.js';
+import type { CheckpointerContext } from './setup.js';
+/** The thread and namespace a read was asked for. */
+export interface ThreadLocation {
+    readonly threadId: string;
+    readonly checkpointNs: string;
+}
+/** A checkpoint's parent, whose pending writes a pre-v4 checkpoint's migration reads. */
+export interface PendingSendsSource extends ThreadLocation {
+    readonly parentCheckpointId: string | undefined;
+}
+/** Per-read options for the payload and writes reads. */
+export interface ReadOptions {
+    signal?: AbortSignal;
+    /** `false` for bulk reads (`list`) that trade read-your-writes for half the read cost. */
+    consistent?: boolean;
+}
+/**
+ * The META row a read is about: the one `checkpointId` names, else the newest
+ * in the namespace.
+ *
+ * Accepts: `address` — parsed; its `checkpointId` names the row, and its
+ * absence asks for the newest in the namespace. `signal` — aborts the read.
+ * `reader` — the public method this read serves, named when a foreign row is
+ * reported; `getTuple` unless the caller says otherwise.
+ *
+ * Returns: the row, or undefined when there is none — including when every row
+ * in the namespace has expired, or when the only rows there belong to another
+ * writer.
+ *
+ * Throws: `FORMAT_UNSUPPORTED` for a row of ours written by a newer version;
+ * whatever the read throws.
+ *
+ * Guarantees: strongly consistent, and expiry is judged here rather than waited
+ * for, so a checkpoint past its ttl is absent to every reader however long
+ * DynamoDB's sweep lags. The newest-first read returns the first live row of
+ * ours and stops there, stepping over expired and foreign rows
+ * {@link LATEST_META_PAGE_SIZE} at a time when the adapter has a `ttl`, and
+ * reading one row per page when it has none.
+ */
+export declare function fetchTargetMeta(context: CheckpointerContext, address: ThreadAddress, signal?: AbortSignal, reader?: MetaReader): Promise<CheckpointMetaRow | undefined>;
+/**
+ * The PAYLOAD row of one checkpoint.
+ *
+ * Accepts: `at` — the checkpoint's location. `read.consistent` — defaults to
+ * true; `list` passes false and accepts replica lag, since a listing tolerates
+ * what a read-your-writes `getTuple` does not.
+ *
+ * Returns: the row, or undefined when it is not there — the two rows land in
+ * one transaction, so that is a thread being deleted row by row, or DynamoDB's
+ * expiry sweep reaching the PAYLOAD row before the META row — which the caller
+ * answers as "no checkpoint".
+ *
+ * Throws: `FORMAT_UNSUPPORTED` for a row a newer release wrote; whatever the
+ * read throws.
+ */
+export declare function fetchPayload(context: CheckpointerContext, at: CheckpointLocation, read?: ReadOptions): Promise<CheckpointPayloadRow | undefined>;
+/**
+ * Every pending write stored for one checkpoint, decoded, in write order.
+ *
+ * Accepts: `at` — the checkpoint's location. `read.consistent` — omitted reads
+ * strongly consistently; `list` passes `false` explicitly and accepts replica
+ * lag, `getTuple` passes `true`.
+ *
+ * Returns: the writes after `dropSupersededWrites` has resolved
+ * first-write-wins; a checkpoint with none returns an empty array.
+ *
+ * Throws: `FORMAT_UNSUPPORTED` for a row a newer release wrote; whatever the
+ * query or the payload decode throws.
+ *
+ * Guarantees: the read is deliberately uncapped. It must be complete to be
+ * correct — a `Send` fan-out retried with a changed write order leaves
+ * superseded rows behind that would count toward any cap — so past
+ * {@link LIST_SCAN_WARN_THRESHOLD} rows the read still succeeds and an operator
+ * is told the checkpoint is unusually heavy.
+ */
+export declare function fetchPendingWrites(context: CheckpointerContext, at: CheckpointLocation, read?: ReadOptions): Promise<CheckpointPendingWrite[]>;
+/** How a tuple is assembled: cancellation, read consistency, and metadata already decoded by the caller. */
+export interface AssembleOptions {
+    signal?: AbortSignal;
+    /** `true` for `getTuple` (read-your-writes), `false` for the eventually-consistent `list` path. */
+    consistent: boolean;
+    /** Metadata the caller decoded to apply a filter, so it is not decoded (or downloaded) twice. */
+    metadata?: CheckpointMetadata;
+}
+/**
+ * A full {@link CheckpointTuple} built from a META row.
+ *
+ * Accepts: `thread` — the **caller's** location, never the row's: `thread.threadId`
+ * scopes which S3 object the row may point at, so it must come from the
+ * partition the caller asked for; `thread.checkpointNs` is its namespace.
+ * `options.metadata` — already decoded by a filtered `list`, so a filtered
+ * listing decodes and downloads each metadata blob once, not twice.
+ * `options.consistent` — true for `getTuple`, false for `list`.
+ *
+ * Returns: the tuple, with `parentConfig` set only when the row names a parent;
+ * `undefined` when the PAYLOAD row is absent — a thread mid-delete, or an
+ * expiry sweep that reached the PAYLOAD row first — the same answer a caller
+ * gets for a checkpoint that does not exist.
+ *
+ * Throws: whatever the reads and decodes throw.
+ */
+export declare function assembleTuple(context: CheckpointerContext, thread: ThreadLocation, meta: CheckpointMetaRow, options: AssembleOptions): Promise<CheckpointTuple | undefined>;
+/**
+ * Pre-v4 checkpoints kept a task's `Send`s as `__pregel_tasks` pending writes
+ * on the parent rather than in the checkpoint itself. Reading such a
+ * checkpoint rebuilds `channel_values[TASKS]` from those writes and stamps
+ * the channel with the highest version the checkpoint already carries (or
+ * the first version when it carries none), exactly as the reference savers
+ * do, so a thread written before LangGraph 0.2 still resumes.
+ *
+ * Accepts: `checkpoint` — any version. `parent` — the checkpoint's own thread
+ * and namespace, plus `parent.parentCheckpointId`, the parent whose pending
+ * writes hold the sends; absent means there is nowhere to migrate from.
+ * `read` — cancellation and consistency for the pending-writes read.
+ *
+ * Returns: the checkpoint untouched when it is v4 or has no parent, and
+ * otherwise a copy with `channel_values[TASKS]` rebuilt. The input is never
+ * mutated.
+ *
+ * Throws: whatever reading the parent's pending writes throws.
+ */
+export declare function migratePendingSends(context: CheckpointerContext, checkpoint: Checkpoint, parent: PendingSendsSource, read: ReadOptions): Promise<Checkpoint>;

package/dist/esm/checkpointer/internal/read.js

@@ -0,0 +1,125 @@
+import { maxChannelVersion, TASKS, } from '@langchain/langgraph-checkpoint';
+import { nowSeconds } from '../../shared/clock.js';
+import { LIST_SCAN_WARN_THRESHOLD, paginateQuery } from '../../shared/dynamodb/paginate.js';
+import { withDynamoDBRetry, retryFor } from '../../shared/dynamodb/retry.js';
+import { isExpiredRow, withoutExpired, assertReadableRow, } from '../../shared/dynamodb/table-schema.js';
+import { truncateForLog } from '../../shared/logging/truncate.js';
+import { beginsWithQuery, metaRowKey, metaSortKeyPrefix, parseHeadRow, parseWriteRow, partitionKey, payloadRowKey, readCheckpoint, readMetadata, toPendingWrites, writeSortKeyPrefix, } from './rows.js';
+const LATEST_META_PAGE_SIZE = 50;
+export async function fetchTargetMeta(context, address, signal, reader = 'getTuple') {
+    const { threadId, checkpointNs, checkpointId } = address;
+    const now = nowSeconds();
+    if (checkpointId !== undefined) {
+        const result = await withDynamoDBRetry((request) => context.client.get({
+            TableName: context.tableName,
+            Key: metaRowKey({ threadId, checkpointNs, checkpointId }),
+            ConsistentRead: true,
+        }, request), retryFor(context, signal));
+        const meta = parseHeadRow(context, result.Item, reader);
+        return meta && !isExpiredRow(meta, now) ? meta : undefined;
+    }
+    const params = beginsWithQuery(context.tableName, partitionKey(threadId), metaSortKeyPrefix(checkpointNs), {
+        limit: context.ttl === undefined ? 1 : LATEST_META_PAGE_SIZE,
+        consistent: true,
+    });
+    const rows = paginateQuery({
+        retry: retryFor(context, signal),
+        signal,
+        client: context.client,
+        params: withoutExpired(params, now),
+        maxItems: Number.POSITIVE_INFINITY,
+        maxIterations: Number.POSITIVE_INFINITY,
+    });
+    for await (const raw of rows) {
+        const meta = parseHeadRow(context, raw, reader);
+        if (meta && !isExpiredRow(meta, now))
+            return meta;
+    }
+    return undefined;
+}
+export async function fetchPayload(context, at, read = {}) {
+    const result = await withDynamoDBRetry((request) => context.client.get({
+        TableName: context.tableName,
+        Key: payloadRowKey(at),
+        ConsistentRead: read.consistent ?? true,
+    }, request), retryFor(context, read.signal));
+    const item = result.Item;
+    if (item !== undefined)
+        assertReadableRow(item, 'checkpoint payload');
+    return item;
+}
+export async function fetchPendingWrites(context, at, read = {}) {
+    const params = beginsWithQuery(context.tableName, partitionKey(at.threadId), writeSortKeyPrefix(at.checkpointNs, at.checkpointId), { ascending: true, consistent: read.consistent ?? true });
+    const items = [];
+    for await (const item of paginateQuery({
+        retry: retryFor(context, read.signal),
+        signal: read.signal,
+        client: context.client,
+        params,
+        maxItems: Number.POSITIVE_INFINITY,
+        maxIterations: Number.POSITIVE_INFINITY,
+    })) {
+        const row = parseWriteRow(item);
+        if (row === undefined) {
+            context.logger.warn('getTuple: skipped a row that is not a pending-write item', {
+                sortKey: truncateForLog(item.SK),
+            });
+            continue;
+        }
+        items.push(row);
+    }
+    if (items.length >= LIST_SCAN_WARN_THRESHOLD) {
+        context.logger.warn('getTuple: a checkpoint carries very many pending-write rows; the read is complete but slow', {
+            threadId: at.threadId,
+            checkpointId: at.checkpointId,
+            rows: items.length,
+        });
+    }
+    return toPendingWrites(context, items, at.threadId, read.signal);
+}
+function configFor(threadId, checkpointNs, checkpointId) {
+    return {
+        configurable: { thread_id: threadId, checkpoint_ns: checkpointNs, checkpoint_id: checkpointId },
+    };
+}
+export async function assembleTuple(context, thread, meta, options) {
+    const read = { signal: options.signal, consistent: options.consistent };
+    const at = { ...thread, checkpointId: meta.checkpointId };
+    const payload = await fetchPayload(context, at, read);
+    if (!payload)
+        return undefined;
+    const [checkpoint, metadata, pendingWrites] = await Promise.all([
+        readCheckpoint(context, payload, thread.threadId, options.signal).then((stored) => migratePendingSends(context, stored, { ...thread, parentCheckpointId: meta.parentCheckpointId }, read)),
+        options.metadata ?? readMetadata(context, meta, thread.threadId, options.signal),
+        fetchPendingWrites(context, at, read),
+    ]);
+    const tuple = {
+        config: configFor(thread.threadId, thread.checkpointNs, meta.checkpointId),
+        checkpoint,
+        metadata,
+        pendingWrites,
+    };
+    if (meta.parentCheckpointId !== undefined) {
+        tuple.parentConfig = configFor(thread.threadId, thread.checkpointNs, meta.parentCheckpointId);
+    }
+    return tuple;
+}
+export async function migratePendingSends(context, checkpoint, parent, read) {
+    if (checkpoint.v >= 4 || parent.parentCheckpointId === undefined)
+        return checkpoint;
+    const writes = await fetchPendingWrites(context, {
+        threadId: parent.threadId,
+        checkpointNs: parent.checkpointNs,
+        checkpointId: parent.parentCheckpointId,
+    }, read);
+    const sends = writes.filter(([, channel]) => channel === TASKS).map(([, , value]) => value);
+    const versions = Object.values(checkpoint.channel_versions);
+    return {
+        ...checkpoint,
+        channel_values: { ...checkpoint.channel_values, [TASKS]: sends },
+        channel_versions: {
+            ...checkpoint.channel_versions,
+            [TASKS]: versions.length > 0 ? maxChannelVersion(...versions) : 1,
+        },
+    };
+}