deepline 0.3.228 → 0.3.230

package/dist/bundling-sources/shared_libs/play-runtime/context.ts

@@ -26,6 +26,7 @@ import {
   serializePlayDatasetCell,
 } from '../plays/dataset';
 import type { PlayDataset, PlayDatasetInput } from '../plays/dataset';
+import type { PlayDatasetColumnTypes } from '../play-data-plane/contracts';
 import {
   compileRequestsWithStrategy,
   executeChunkedRequests,
@@ -201,8 +202,13 @@ import {
   TOOL_EXECUTE_RUNNER_ACCEPTS_VALUE,
   TOOL_EXECUTE_PROVIDER_RATE_SCOPE_HEADER,
   TOOL_EXECUTE_PROVIDER_RATE_RETRY_SNAPSHOT_HEADER,
+  TOOL_EXECUTE_PROVIDER_THROTTLE_HEADER,
   classifyToolExecuteHttpFailure,
   createToolExecuteHttpFailureAttemptTracker,
+  TOOL_EXECUTE_BARE_RATE_LIMIT_MAX_ATTEMPTS,
+  TOOL_EXECUTE_RATE_LIMIT_MAX_ATTEMPTS,
+  TOOL_EXECUTE_RETRY_DELAY_FALLBACK_MS,
+  TOOL_EXECUTE_RATE_LIMIT_RETRY_DELAY_MAX_MS,
   parseToolExecuteRetryAfterMs,
   parseToolExecuteRetryAtMs,
   parseToolExecuteAuthScopeChangedError,
@@ -680,7 +686,7 @@ async function runCtxFetchPhase<T>(input: {
   callerSignal?: AbortSignal | null;
   controller: AbortController;
   deadline: CtxFetchDeadline;
-  phase: 'headers' | 'body';
+  phase: 'headers' | 'body' | 'total';
   run: () => Promise<T>;
 }): Promise<T> {
   const remainingMs = remainingCtxFetchMs(input.deadline);
@@ -695,7 +701,9 @@ async function runCtxFetchPhase<T>(input: {
   const phaseLimitMs =
     input.phase === 'headers'
       ? input.deadline.headersMs
-      : input.deadline.bodyMs;
+      : input.phase === 'body'
+        ? input.deadline.bodyMs
+        : remainingMs;
   const timeoutMs = Math.min(phaseLimitMs, remainingMs);
   const timeoutPhase = timeoutMs >= remainingMs ? 'total' : input.phase;
   let timeout: ReturnType<typeof setTimeout> | null = null;
@@ -744,19 +752,42 @@ async function sleepWithinCtxFetchDeadline(input: {
   delayMs: number;
 }): Promise<void> {
   const controller = new AbortController();
-  await runCtxFetchPhase({
-    callerSignal: input.callerSignal,
-    controller,
-    deadline: input.deadline,
-    phase: 'headers',
-    run: () =>
-      new Promise<void>((resolve) => {
-        setTimeout(
-          resolve,
-          Math.min(input.delayMs, remainingCtxFetchMs(input.deadline)),
-        );
-      }),
-  });
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  try {
+    await runCtxFetchPhase({
+      callerSignal: input.callerSignal,
+      controller,
+      deadline: input.deadline,
+      phase: 'total',
+      run: () =>
+        new Promise<void>((resolve) => {
+          timer = setTimeout(resolve, input.delayMs);
+        }),
+    });
+  } finally {
+    if (timer !== undefined) clearTimeout(timer);
+  }
+}
+
+function isCtxFetchBackpressure(response: Response, bodyText: string): boolean {
+  if (response.status === 429) return true;
+  if (response.status !== 503) return false;
+  // This Salesforce endpoint reports a rejected burst as 503. Match the final
+  // destination and exact response contract; arbitrary 503s remain failures.
+  try {
+    const destination = new URL(response.url);
+    const body: unknown = JSON.parse(bodyText);
+    return (
+      destination.protocol === 'https:' &&
+      destination.hostname.endsWith('.salesforce-sites.com') &&
+      body !== null &&
+      typeof body === 'object' &&
+      !Array.isArray(body) &&
+      (body as Record<string, unknown>).status === 'burst_circuit_open'
+    );
+  } catch {
+    return false;
+  }
 }
 
 async function readCtxFetchBody(input: {
@@ -1673,6 +1704,69 @@ function rowsFromUnknown(value: unknown): Record<string, unknown>[] {
   );
 }
 
+function customerDbColumnTypes(
+  columns: unknown,
+  rows: readonly Record<string, unknown>[],
+): PlayDatasetColumnTypes | undefined {
+  if (!Array.isArray(columns)) return undefined;
+  // Core PostgreSQL OIDs only. Extension/domain identities are database-local,
+  // and source SQL types never authorize coercing the values returned by a tool.
+  const coreValueTypes: Record<
+    number,
+    PlayDatasetColumnTypes[string]['valueType']
+  > = {
+    16: 'boolean',
+    20: 'string',
+    21: 'number',
+    23: 'number',
+    25: 'string',
+    114: 'json',
+    700: 'number',
+    701: 'number',
+    1042: 'string',
+    1043: 'string',
+    1082: 'string',
+    1083: 'string',
+    1114: 'string',
+    1184: 'string',
+    1266: 'string',
+    1700: 'string',
+    2950: 'string',
+    3802: 'json',
+  };
+  const entries: Array<[string, PlayDatasetColumnTypes[string]]> = [];
+  for (const value of columns) {
+    const column = recordOrNull(value);
+    const postgresTypeId = column?.data_type_id;
+    if (
+      typeof column?.name !== 'string' ||
+      typeof postgresTypeId !== 'number' ||
+      !Object.hasOwn(coreValueTypes, postgresTypeId)
+    )
+      continue;
+    const observed = new Set<PlayDatasetColumnTypes[string]['valueType']>();
+    for (const row of rows) {
+      const cell = row[column.name];
+      if (cell === null || cell === undefined) continue;
+      observed.add(
+        typeof cell === 'boolean'
+          ? 'boolean'
+          : typeof cell === 'string'
+            ? 'string'
+            : typeof cell === 'number' && Number.isFinite(cell)
+              ? 'number'
+              : 'json',
+      );
+    }
+    const valueType =
+      postgresTypeId === 114 || postgresTypeId === 3802 || observed.size > 1
+        ? 'json'
+        : ([...observed][0] ?? coreValueTypes[postgresTypeId]!);
+    entries.push([column.name, { postgresTypeId, valueType }]);
+  }
+  return entries.length > 0 ? Object.fromEntries(entries) : undefined;
+}
+
 function finiteNonNegativeInteger(value: unknown): number | null {
   if (typeof value !== 'number' || !Number.isFinite(value) || value < 0) {
     return null;
@@ -2541,9 +2635,7 @@ const WAITING_ROW = Symbol('deepline.waiting_row');
 const FAILED_ROW = Symbol('deepline.failed_row');
 const COMPLETED_ROW = Symbol('deepline.completed_row');
 type MapRowExecutionResult =
-  | typeof COMPLETED_ROW
-  | typeof WAITING_ROW
-  | typeof FAILED_ROW;
+  typeof COMPLETED_ROW | typeof WAITING_ROW | typeof FAILED_ROW;
 
 function normalizeFetchHeaders(
   headers: RequestInit['headers'],
@@ -4341,8 +4433,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
     artifactTableNamespace: string;
     mapNodeId?: string | null;
     explicitKey?:
-      | ((row: Record<string, unknown>, index: number) => string)
-      | null;
+      ((row: Record<string, unknown>, index: number) => string) | null;
   }): MapExecutionScope {
     const mapInvocationId = `${input.logicalNamespace}:${this.mapInvocationIndex}`;
     this.mapInvocationIndex += 1;
@@ -4592,12 +4683,10 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
       // run actually owns. The map/cohort publisher keeps the claim-free
       // batched write, where the pointer genuinely shares the canonical
       // result's single bounded request.
-      if (
-        !(
-          this.#options.claimRuntimeStepReceipt ||
-          this.#options.claimRuntimeStepReceipts
-        )
-      ) {
+      if (!(
+        this.#options.claimRuntimeStepReceipt ||
+        this.#options.claimRuntimeStepReceipts
+      )) {
         return;
       }
       const existing = await this.getRuntimeStepReceipts(
@@ -7914,6 +8003,9 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
         rows.length >= effectiveCount ? rows.slice(0, effectiveCount) : null,
       sourceLabel: 'query result rows',
       tableNamespace: null,
+      columnTypes: isCustomerDbDatasetTool(toolId)
+        ? customerDbColumnTypes(raw?.columns, rows)
+        : undefined,
       resolvers: {
         count: async () => Math.min(totalRows, datasetLimit),
         at: async (index) => (await fetchPage(index, 1))[0],
@@ -8378,8 +8470,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
     _key: string,
     _items: PlayDatasetInput<T>,
     _input?:
-      | MapFieldDefinition<T, Record<string, unknown>>
-      | RuntimeStepProgram,
+      MapFieldDefinition<T, Record<string, unknown>> | RuntimeStepProgram,
     _options?: RuntimeDatasetOptions<T>,
   ): never {
     void _key;
@@ -8472,6 +8563,18 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
     let itemsToProcess: Array<Record<string, unknown>> = [];
     let itemOriginalIndexes: number[] = [];
     const datasetColumnNames = Object.keys(input);
+    // Generated fields no longer describe the source column, even when they
+    // retain its name. Every writer still validates the actual incoming values.
+    const inputColumnTypes = isPlayDataset(items)
+      ? items.columnTypes
+      : undefined;
+    const unchangedColumnTypes = inputColumnTypes
+      ? Object.fromEntries(
+          Object.entries(inputColumnTypes).filter(
+            ([field]) => !datasetColumnNames.includes(field),
+          ),
+        )
+      : undefined;
     const stripFieldOutputs = (row: Record<string, unknown>) => {
       const stripped = cloneCsvAliasedRow(row);
       for (const fieldName of datasetColumnNames) {
@@ -8491,8 +8594,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
 
     const userKeyOption = options?.key;
     let explicitKeyResolver:
-      | ((row: Record<string, unknown>, index: number) => string)
-      | null = null;
+      ((row: Record<string, unknown>, index: number) => string) | null = null;
     if (userKeyOption !== undefined) {
       explicitKeyResolver = (row, index) => {
         const inputRow = stripFieldOutputs(row);
@@ -8863,6 +8965,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
             forceRefresh: this.#options.cachePolicy?.forceToolRefresh === true,
             inputOffset: page.offset,
             mode: options?.mode,
+            columnTypes: unchangedColumnTypes,
           },
         );
         resolvedTableNamespace = normalizeTableNamespace(
@@ -9424,6 +9527,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
             ? immediateMaterializedRows
             : null,
         tableNamespace: resolvedTableNamespace,
+        columnTypes: unchangedColumnTypes,
         workProgress: {
           total: totalInputCount,
           executed: executedCount,
@@ -9666,6 +9770,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
           staticPipeline: this.currentStaticPipeline,
           forceRefresh: this.#options.cachePolicy?.forceToolRefresh === true,
           mode: options?.mode,
+          columnTypes: unchangedColumnTypes,
         },
       );
       resolvedTableNamespace = normalizeTableNamespace(
@@ -10175,6 +10280,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
         .map((row) => this.toMaterializedOutputRow(row)),
       residentRows: terminalRows,
       tableNamespace: resolvedTableNamespace,
+      columnTypes: unchangedColumnTypes,
       workProgress: {
         total: totalInputCount,
         executed: executedCount,
@@ -12895,6 +13001,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
     const callIndex = this.toolCallIndexByKey.get(callIndexKey) ?? 0;
     this.toolCallIndexByKey.set(callIndexKey, callIndex + 1);
     const logicalCallId = stableDigest(`${callIndexKey}:${callIndex}`);
+    this.resourceGovernor.recordLogicalToolCall();
     let executionAuthScopeDigest =
       (await this.resolveToolAuthScopeDigest(toolId))?.trim() ?? null;
     const eventWaitHandler =
@@ -14991,16 +15098,30 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
           let response: Response | null = null;
           let bodyText: string | null = null;
           try {
-            const canRetryTransport = ['GET', 'HEAD', 'OPTIONS'].includes(
-              method,
-            );
+            // Only replay a body whose bytes are stable. In particular, a
+            // consumed ReadableStream cannot be retried even with a write key.
+            const replayableBody =
+              init.body == null || typeof init.body === 'string';
+            const canRetryTransport =
+              replayableBody && ['GET', 'HEAD', 'OPTIONS'].includes(method);
             const { safePublicFetch } = await loadSafeFetch();
-            for (
-              let attempt = 1;
-              attempt <= FETCH_TRANSPORT_MAX_ATTEMPTS;
-              attempt += 1
-            ) {
+            let transportAttempt = 1;
+            let rateLimitAttempt = 0;
+            while (true) {
               const attemptController = new AbortController();
+              let permit: Awaited<
+                ReturnType<RuntimeResourceGovernor['acquireProviderPermit']>
+              > | null = null;
+              let dispatched = false;
+              const releasePhysicalPermit = async () => {
+                const current = permit;
+                permit = null;
+                if (!current) return;
+                if (dispatched) current.release();
+                else await current.cancel();
+                dispatched = false;
+              };
+              let transportRetry = false;
               try {
                 const attemptInit = {
                   ...fetchInit,
@@ -15012,10 +15133,53 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
                   callerSignal: init.signal,
                   controller: attemptController,
                   deadline: fetchDeadline,
-                  phase: 'headers',
+                  phase: 'total',
                   run: () =>
                     safePublicFetch(url, attemptInit, {
-                      fetchImpl: this.#options.fetchImpl,
+                      // safePublicFetch validates each URL before invoking this
+                      // wrapper. Redirect hops are physical dispatches too.
+                      fetchImpl: async (physicalUrl, physicalInit) => {
+                        await releasePhysicalPermit();
+                        const signal = attemptInit.signal;
+                        permit = await runCtxFetchPhase({
+                          callerSignal: init.signal,
+                          controller: attemptController,
+                          deadline: fetchDeadline,
+                          phase: 'total',
+                          run: async () => {
+                            for (;;) {
+                              const acquired =
+                                await this.resourceGovernor.acquireProviderPermit(
+                                  { toolId: CTX_FETCH_EGRESS_TOOL_ID, signal },
+                                );
+                              try {
+                                signal.throwIfAborted();
+                                const valid =
+                                  !acquired.validate ||
+                                  (await acquired.validate());
+                                signal.throwIfAborted();
+                                if (valid) return acquired;
+                              } catch (error) {
+                                await acquired.cancel();
+                                throw error;
+                              }
+                              await acquired.cancel();
+                            }
+                          },
+                        });
+                        return await runCtxFetchPhase({
+                          callerSignal: init.signal,
+                          controller: attemptController,
+                          deadline: fetchDeadline,
+                          phase: 'headers',
+                          run: () => {
+                            dispatched = true;
+                            return (
+                              this.#options.fetchImpl ?? globalThis.fetch
+                            )(physicalUrl, physicalInit);
+                          },
+                        });
+                      },
                       sensitiveHeaders: Object.keys(secretHeaderMarkers),
                       stripHeadersOnCrossOriginRedirect: true,
                       headersOnCrossOriginRedirect: {
@@ -15029,7 +15193,6 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
                   deadline: fetchDeadline,
                   response,
                 });
-                break;
               } catch (error) {
                 void response?.body?.cancel(error).catch(() => undefined);
                 response = null;
@@ -15045,22 +15208,67 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
                   error instanceof Error ? error.message : String(error);
                 if (
                   canRetryTransport &&
-                  attempt < FETCH_TRANSPORT_MAX_ATTEMPTS
+                  transportAttempt < FETCH_TRANSPORT_MAX_ATTEMPTS
                 ) {
                   this.runtimeLog(
-                    `ctx.fetch(${method} ${url}) transport failed on attempt ${attempt}/${FETCH_TRANSPORT_MAX_ATTEMPTS}; retrying: ${message}`,
+                    `ctx.fetch(${method} ${url}) transport failed on attempt ${transportAttempt}/${FETCH_TRANSPORT_MAX_ATTEMPTS}; retrying: ${message}`,
+                  );
+                  transportRetry = true;
+                } else {
+                  throw new Error(
+                    `ctx.fetch(${method} ${url}) failed on attempt ${transportAttempt}/${FETCH_TRANSPORT_MAX_ATTEMPTS}: ${message}`,
                   );
-                  await sleepWithinCtxFetchDeadline({
-                    callerSignal: init.signal,
-                    deadline: fetchDeadline,
-                    delayMs: FETCH_TRANSPORT_RETRY_DELAY_MS * attempt,
-                  });
-                  continue;
                 }
-                throw new Error(
-                  `ctx.fetch(${method} ${url}) failed on attempt ${attempt}/${FETCH_TRANSPORT_MAX_ATTEMPTS}: ${message}`,
-                );
+              } finally {
+                await releasePhysicalPermit();
+              }
+              if (transportRetry) {
+                await sleepWithinCtxFetchDeadline({
+                  callerSignal: init.signal,
+                  deadline: fetchDeadline,
+                  delayMs: FETCH_TRANSPORT_RETRY_DELAY_MS * transportAttempt,
+                });
+                transportAttempt += 1;
+                continue;
               }
+              if (
+                !response ||
+                bodyText === null ||
+                !replayableBody ||
+                !isCtxFetchBackpressure(response, bodyText)
+              )
+                break;
+              rateLimitAttempt += 1;
+              const retryAfter = response.headers.get('retry-after');
+              const attemptCap = retryAfter?.trim()
+                ? TOOL_EXECUTE_RATE_LIMIT_MAX_ATTEMPTS
+                : TOOL_EXECUTE_BARE_RATE_LIMIT_MAX_ATTEMPTS;
+              if (rateLimitAttempt >= attemptCap) break;
+              const delayMs = Math.max(
+                parseToolExecuteRetryAfterMs(retryAfter),
+                Math.min(
+                  TOOL_EXECUTE_RATE_LIMIT_RETRY_DELAY_MAX_MS,
+                  TOOL_EXECUTE_RETRY_DELAY_FALLBACK_MS * rateLimitAttempt,
+                ),
+              );
+              this.runtimeLog(
+                `[runtime.ctx_fetch_backpressure] ${JSON.stringify({
+                  key: normalizedKey,
+                  method,
+                  origin: new URL(response.url || url).origin,
+                  http_status: response.status,
+                  attempt: rateLimitAttempt,
+                  max_attempts: attemptCap,
+                  retry_delay_ms: delayMs,
+                })}`,
+              );
+              // Destination feedback must not cool the shared generic_http
+              // bucket used by unrelated websites. Pace this logical request.
+              await sleepWithinCtxFetchDeadline({
+                callerSignal: init.signal,
+                deadline: fetchDeadline,
+                delayMs,
+              });
             }
             if (!response) {
               throw new Error(
@@ -15283,6 +15491,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
       // Terminal logs are intentionally bounded; keep the small aggregate in
       // stats so every retained terminal snapshot can explain queueing.
       runtimeResources: this.getRuntimeResourceSnapshot(),
+      runtimeWorkCounters: this.getRuntimeWorkCounters(),
     };
   }
 
@@ -15291,6 +15500,15 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
     return this.resourceGovernor.snapshot();
   }
 
+  /** Attempt-local counts without serializing the diagnostic resource map. */
+  getRuntimeWorkCounters(): {
+    logicalCalls: number;
+    providerRequests: number;
+    provider429s: number;
+  } {
+    return this.resourceGovernor.workCounters();
+  }
+
   /**
    * Record what one receipt gateway response said about itself.
    *
@@ -16467,8 +16685,8 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
               },
             })),
             buildPhysicalCall: async ({ members }) => {
-              const owned = members.map(
-                (member) => batchRequestByCallId.get(member.memberId)!,
+              const owned = members.map((member) =>
+                batchRequestByCallId.get(member.memberId)!,
               );
               const compiled = compileOwnedPhysicalBatch(owned);
               const providerActionVersion =
@@ -18794,9 +19012,7 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
     // running for.
     const activeRowScope = rowContext.getStore();
     const memoryLedgerRowScopes: readonly (
-      | MemoryLedgerRowScope
-      | null
-      | undefined
+      MemoryLedgerRowScope | null | undefined
     )[] = options?.memoryLedgerRowScopes ?? [
       { rowKey: activeRowScope?.rowKey, fieldName: activeRowScope?.fieldName },
     ];
@@ -19639,6 +19855,9 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
                           logProvider429RetryPhase('http_dispatch_started');
                           activeProviderRequest =
                             this.beginProviderRequest(provider);
+                          // Count at dispatch, including ambiguous transport
+                          // failures that never produce a residence sample.
+                          this.resourceGovernor.recordProviderDispatch();
                           const dispatchedResponse = await fetch(url, {
                             method: 'POST',
                             signal: abortController?.signal,
@@ -19694,21 +19913,13 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
                           // it has been read and verified below.
                           response = dispatchedResponse;
                           if (dispatchedResponse.status === 429) {
-                            provider429RetryDiagnosticStartedAt =
-                              performance.now();
-                            provider429RetryDiagnosticNextInvocationAttempt =
-                              invocationAttempt + 1;
-                            logProvider429RetryPhase(
-                              'http_429_response_received',
-                            );
-                            providerRequestsInFlightAt429 =
-                              activeProviderRequest?.inFlightCount() ?? null;
                             const observedAt = Date.now();
                             const retryAfter =
                               dispatchedResponse.headers.get('retry-after');
-                            // Without a provider deadline the classifier owns
-                            // the conservative bare-429 cooldown. Do not replace
-                            // it with the generic one-second HTTP retry fallback.
+                            // Preserve absolute provider deadlines even when the
+                            // response did not carry the gateway's counter
+                            // attestation. Classification below still decides
+                            // whether this 429 may update shared provider pacing.
                             providerRetryAtMs = retryAfter?.trim()
                               ? parseToolExecuteRetryAtMs(
                                   retryAfter,
@@ -19716,6 +19927,23 @@ export class PlayContextImpl implements ScalarPlayAuthoringRuntimeContext {
                                 )
                               : undefined;
                           }
+                          if (
+                            dispatchedResponse.status === 429 &&
+                            dispatchedResponse.headers.get(
+                              TOOL_EXECUTE_PROVIDER_THROTTLE_HEADER,
+                            ) === '1'
+                          ) {
+                            this.resourceGovernor.recordProvider429();
+                            provider429RetryDiagnosticStartedAt =
+                              performance.now();
+                            provider429RetryDiagnosticNextInvocationAttempt =
+                              invocationAttempt + 1;
+                            logProvider429RetryPhase(
+                              'http_429_response_received',
+                            );
+                            providerRequestsInFlightAt429 =
+                              activeProviderRequest?.inFlightCount() ?? null;
+                          }
                           const dispatchedFetchToHeadersMs =
                             Date.now() - fetchedAt;
                           // The response writer owns these named wall times.

package/dist/bundling-sources/shared_libs/play-runtime/ctx-types.ts

@@ -7,6 +7,7 @@ import type {
 } from './cell-provenance';
 import type { PlayBundleArtifact } from '../plays/artifact-types';
 import type { PlayRunContractSnapshot } from '../plays/contracts';
+import type { PlayDatasetColumnTypes } from '../play-data-plane/contracts';
 import type { PlayStructuredDefinition } from '../plays/definition';
 import type { PlayStaticPipeline } from '../plays/static-pipeline';
 import type {
@@ -891,6 +892,7 @@ export interface ContextOptions {
       forceRefresh?: boolean;
       inputOffset?: number;
       mode?: RuntimeDatasetRunOptions['mode'];
+      columnTypes?: PlayDatasetColumnTypes;
     },
   ) => Promise<MapStartResult>;
   /**
@@ -1169,6 +1171,9 @@ export interface ContextOptions {
 }
 
 export interface PlayCheckpoint {
+  /** Last completed runner invocation within the current scheduler attempt. */
+  workCounterInvocationAttempt?: number;
+  workCounterInvocationSeq?: number;
   /**
    * Clock captured when this logical run first creates its context. All
    * staleAfterSeconds receipt buckets use this value across durable resumes.

package/dist/bundling-sources/shared_libs/play-runtime/deployment/fly-machines.ts

@@ -566,6 +566,7 @@ export class FlyMachinesClient {
     method: string,
     path: string,
     body?: unknown,
+    leaseNonce?: string,
   ): Promise<unknown> {
     const retryDelaysMs =
       this.rateLimitOptions.retryDelaysMs ??
@@ -585,6 +586,7 @@ export class FlyMachinesClient {
             headers: {
               authorization: `Bearer ${this.token}`,
               'content-type': 'application/json',
+              ...(leaseNonce ? { 'fly-machine-lease-nonce': leaseNonce } : {}),
             },
             body: body === undefined ? undefined : JSON.stringify(body),
             signal,
@@ -689,10 +691,43 @@ export class FlyMachinesClient {
     );
   }
 
-  async deleteMachine(id: string, force: boolean): Promise<void> {
+  /** Hold exclusive provider ownership while rechecking and reclaiming a
+   * stopped Machine. A concurrent rollback cannot update its hold or restart
+   * it until this lease is released or expires. */
+  async acquireMachineLease(
+    id: string,
+    options: { ttl: number; description: string },
+  ): Promise<string> {
+    const result = await this.request('POST', `/machines/${id}/lease`, options);
+    const nonce =
+      result &&
+      typeof result === 'object' &&
+      'data' in result &&
+      result.data &&
+      typeof result.data === 'object' &&
+      'nonce' in result.data
+        ? result.data.nonce
+        : undefined;
+    if (typeof nonce !== 'string' || nonce.trim() === '') {
+      throw new Error('Fly Machines API did not return a lease nonce.');
+    }
+    return nonce;
+  }
+
+  async releaseMachineLease(id: string, nonce: string): Promise<void> {
+    await this.request('DELETE', `/machines/${id}/lease`, undefined, nonce);
+  }
+
+  async deleteMachine(
+    id: string,
+    force: boolean,
+    leaseNonce?: string,
+  ): Promise<void> {
     await this.request(
       'DELETE',
       `/machines/${id}?force=${force ? 'true' : 'false'}`,
+      undefined,
+      leaseNonce,
     );
   }
 
@@ -718,8 +753,9 @@ export class FlyMachinesClient {
   async stopMachine(
     id: string,
     options?: { signal?: string; timeout?: string },
+    leaseNonce?: string,
   ): Promise<void> {
-    await this.request('POST', `/machines/${id}/stop`, options);
+    await this.request('POST', `/machines/${id}/stop`, options, leaseNonce);
   }
 
   async machinesInLane(releaseId: string): Promise<FlyMachine[]> {