deepline 0.1.201 → 0.1.203

package/dist/bundling-sources/shared_libs/play-runtime/governor/app-runtime-rate-state-backend.ts

@@ -23,7 +23,20 @@ const MAX_ACQUIRE_SLEEP_MS = 5_000;
 const STORE_FAILURE_RETRY_DELAYS_MS = [100, 250] as const;
 
 interface LeasedBlock {
-  remaining: number;
+  /**
+   * Concurrency reservations, one per locally-cached permit. Only rules with a
+   * `maxConcurrency` mint lease tokens; each id maps to a durable slot that must
+   * be released (or TTL-expired) so the shared concurrency count drains.
+   */
+  leaseIds: string[];
+  /**
+   * Pure-rate (RPS-only) permits with no concurrency slot behind them. The
+   * server decremented the provider window by this count already, so drawing
+   * one locally is a no-op that needs no release token. Tracked separately from
+   * `leaseIds` so a pure-RPS block (the common `defaultPacingForTool` shape)
+   * actually amortizes its whole grant across acquires instead of burning it.
+   */
+  windowPermits: number;
   expiresAt: number;
   rulesKey: string;
   rules: PacingRule[];
@@ -92,10 +105,43 @@ export class AppRuntimeRateStateBackend implements RateStateBackend {
         });
         consecutiveFailures = 0;
         if (response.granted > 0) {
-          const remaining = response.granted - 1;
-          if (remaining > 0)
-            this.mergeBlock(bucketId, remaining, rulesKey, ordered);
-          return this.permitFor(bucketId, ordered);
+          const hasConcurrency = ordered.some(
+            (rule) => rule.maxConcurrency != null,
+          );
+          const leaseIds = Array.isArray(response.leaseIds)
+            ? response.leaseIds.filter(
+                (leaseId): leaseId is string =>
+                  typeof leaseId === 'string' && leaseId.trim().length > 0,
+              )
+            : [];
+          if (hasConcurrency && leaseIds.length !== response.granted) {
+            throw new Error(
+              `Rate-state store granted ${response.granted} concurrency permits for ${bucketId} with ${leaseIds.length} lease tokens.`,
+            );
+          }
+          if (!hasConcurrency && leaseIds.length > 0) {
+            throw new Error(
+              `Rate-state store returned ${leaseIds.length} lease tokens for pure-rate bucket ${bucketId}.`,
+            );
+          }
+          // Serve this acquire from the grant, then cache the rest so the whole
+          // block amortizes into a single round trip. Concurrency rules draw a
+          // lease token (needs release); pure-rate rules draw a window permit
+          // (already decremented server-side, no release).
+          const currentLeaseId = hasConcurrency ? leaseIds[0]! : null;
+          const remainingLeaseIds = hasConcurrency ? leaseIds.slice(1) : [];
+          const remainingWindowPermits = hasConcurrency
+            ? 0
+            : response.granted - 1;
+          if (remainingLeaseIds.length > 0 || remainingWindowPermits > 0)
+            this.mergeBlock(
+              bucketId,
+              remainingLeaseIds,
+              remainingWindowPermits,
+              rulesKey,
+              ordered,
+            );
+          return this.permitFor(bucketId, ordered, currentLeaseId);
         }
         await this.sleep(
           Math.max(1, Math.min(response.waitMs, MAX_ACQUIRE_SLEEP_MS)),
@@ -137,7 +183,11 @@ export class AppRuntimeRateStateBackend implements RateStateBackend {
     });
   }
 
-  private permitFor(bucketId: string, rules: PacingRule[]): PacingPermit {
+  private permitFor(
+    bucketId: string,
+    rules: PacingRule[],
+    leaseId: string | null,
+  ): PacingPermit {
     if (!rules.some((rule) => rule.maxConcurrency != null)) {
       return noopPacingPermit();
     }
@@ -146,14 +196,15 @@ export class AppRuntimeRateStateBackend implements RateStateBackend {
       release: () => {
         if (released) return;
         released = true;
-        this.releaseReserved(bucketId, rules, 1);
+        this.releaseReserved(bucketId, rules, leaseId ? [leaseId] : []);
       },
     };
   }
 
   private mergeBlock(
     bucketId: string,
-    remaining: number,
+    leaseIds: string[],
+    windowPermits: number,
     rulesKey: string,
     rules: PacingRule[],
   ): void {
@@ -164,14 +215,16 @@ export class AppRuntimeRateStateBackend implements RateStateBackend {
       existing.rulesKey === rulesKey &&
       existing.expiresAt > this.now()
     ) {
-      existing.remaining += remaining;
+      existing.leaseIds.push(...leaseIds);
+      existing.windowPermits += windowPermits;
       existing.expiresAt = Math.min(existing.expiresAt, freshExpiresAt);
       return;
     }
     if (existing)
-      this.releaseReserved(bucketId, existing.rules, existing.remaining);
+      this.releaseReserved(bucketId, existing.rules, existing.leaseIds);
     this.blocks.set(bucketId, {
-      remaining,
+      leaseIds,
+      windowPermits,
       expiresAt: freshExpiresAt,
       rulesKey,
       rules,
@@ -186,32 +239,46 @@ export class AppRuntimeRateStateBackend implements RateStateBackend {
     if (!block) return null;
     if (block.rulesKey !== rulesKey || block.expiresAt <= this.now()) {
       this.blocks.delete(bucketId);
-      this.releaseReserved(bucketId, block.rules, block.remaining);
+      // Unused window permits are forfeited on TTL expiry (the server already
+      // decremented the provider window for them, and the LEASE_BLOCK_TTL_MS
+      // window is short so at most one block's worth is stranded per bucket).
+      // Concurrency leases must be released so the shared slot count drains.
+      this.releaseReserved(bucketId, block.rules, block.leaseIds);
       return null;
     }
-    block.remaining -= 1;
-    if (block.remaining <= 0) this.blocks.delete(bucketId);
-    return this.permitFor(bucketId, block.rules);
+    const hasConcurrency = block.rules.some(
+      (rule) => rule.maxConcurrency != null,
+    );
+    // Concurrency blocks draw a lease token (release semantics); pure-rate
+    // blocks draw an untracked window permit that the server already accounted.
+    const leaseId = hasConcurrency ? (block.leaseIds.pop() ?? null) : null;
+    if (!hasConcurrency && block.windowPermits > 0) block.windowPermits -= 1;
+    if (block.leaseIds.length === 0 && block.windowPermits === 0)
+      this.blocks.delete(bucketId);
+    return this.permitFor(bucketId, block.rules, leaseId);
   }
 
   private drainBlock(bucketId: string): void {
     const block = this.blocks.get(bucketId);
     if (!block) return;
     this.blocks.delete(bucketId);
-    this.releaseReserved(bucketId, block.rules, block.remaining);
+    this.releaseReserved(bucketId, block.rules, block.leaseIds);
   }
 
   private releaseReserved(
     bucketId: string,
     rules: readonly PacingRule[],
-    count: number,
+    leaseIds: readonly string[],
   ): void {
-    if (count <= 0 || !rules.some((rule) => rule.maxConcurrency != null))
+    if (
+      leaseIds.length === 0 ||
+      !rules.some((rule) => rule.maxConcurrency != null)
+    )
       return;
     void releaseRateStateViaAppRuntime(this.context, {
       bucketId,
       rules: [...rules],
-      count,
+      leaseIds: [...leaseIds],
       schedulerSchema: this.schedulerSchema,
     }).catch((error) => {
       const normalized = normalizeError(error);

package/dist/bundling-sources/shared_libs/play-runtime/governor/block-reserving-budget-state-backend.ts

@@ -0,0 +1,141 @@
+import type {
+  PlayRunnerBudgetCharge,
+  PlayRunnerBudgetChargeResult,
+} from '../protocol';
+import type { BudgetStateBackend } from './budget-state-backend';
+import { BudgetStateLimitError } from './budget-state-backend';
+
+const RESERVATION_BLOCKS: Record<string, number> = {
+  toolCall: 256,
+  retry: 16,
+  waterfallStep: 256,
+  playCall: 16,
+  descendant: 16,
+};
+
+export type BudgetStateChargePort = (input: {
+  rootRunId: string;
+  reservationId: string;
+  charges: readonly PlayRunnerBudgetCharge[];
+}) => Promise<PlayRunnerBudgetChargeResult>;
+
+/**
+ * Reserves durable budget in coarse blocks while charging each local operation
+ * exactly once. The durable port remains the authority across sandboxes and
+ * isolates; this cache only removes one RPC from the hot path per operation.
+ */
+export class BlockReservingBudgetStateBackend implements BudgetStateBackend {
+  private readonly remaining = new Map<string, number>();
+  private chargeTail: Promise<void> = Promise.resolve();
+
+  constructor(private readonly chargePort: BudgetStateChargePort) {}
+
+  charge(input: {
+    rootRunId: string;
+    charges: readonly PlayRunnerBudgetCharge[];
+  }): Promise<PlayRunnerBudgetChargeResult> {
+    const result = this.chargeTail.then(() => this.chargeSerial(input));
+    this.chargeTail = result.then(
+      () => undefined,
+      () => undefined,
+    );
+    return result;
+  }
+
+  private async chargeSerial(input: {
+    rootRunId: string;
+    charges: readonly PlayRunnerBudgetCharge[];
+  }): Promise<PlayRunnerBudgetChargeResult> {
+    const reservations: PlayRunnerBudgetCharge[] = [];
+    const exactReservations: PlayRunnerBudgetCharge[] = [];
+    const previousRemaining = new Map<string, number | undefined>();
+    for (const charge of input.charges) {
+      const remaining = this.remaining.get(charge.key) ?? 0;
+      previousRemaining.set(charge.key, this.remaining.get(charge.key));
+      if (remaining >= charge.amount) {
+        this.remaining.set(charge.key, remaining - charge.amount);
+        continue;
+      }
+      const shortfall = charge.amount - remaining;
+      const block = Math.min(
+        charge.limit,
+        Math.max(shortfall, RESERVATION_BLOCKS[charge.key] ?? 1),
+      );
+      reservations.push({ ...charge, amount: block });
+      exactReservations.push({ ...charge, amount: shortfall });
+      this.remaining.set(charge.key, block - shortfall);
+    }
+    if (reservations.length === 0) return { counters: {} };
+
+    const restorePreviousRemaining = () => {
+      for (const [key, remaining] of previousRemaining) {
+        if (remaining === undefined) this.remaining.delete(key);
+        else this.remaining.set(key, remaining);
+      }
+    };
+    const request = async (
+      charges: readonly PlayRunnerBudgetCharge[],
+      reservationId: string,
+    ) =>
+      await this.chargePort({
+        rootRunId: input.rootRunId,
+        reservationId,
+        charges,
+      });
+    const requestWithIdempotentRetry = async (
+      charges: readonly PlayRunnerBudgetCharge[],
+      reservationId: string,
+    ) => {
+      try {
+        return await request(charges, reservationId);
+      } catch (error) {
+        const normalized = normalizeBudgetError(error);
+        if (normalized instanceof BudgetStateLimitError) throw normalized;
+        return await request(charges, reservationId).catch((retryError) => {
+          throw normalizeBudgetError(retryError);
+        });
+      }
+    };
+    try {
+      return await requestWithIdempotentRetry(
+        reservations,
+        crypto.randomUUID(),
+      );
+    } catch (error) {
+      const normalized = normalizeBudgetError(error);
+      if (
+        normalized instanceof BudgetStateLimitError &&
+        reservations.some(
+          (reserved, index) =>
+            reserved.amount >
+            (exactReservations[index]?.amount ?? reserved.amount),
+        )
+      ) {
+        for (const reservation of exactReservations) {
+          this.remaining.set(reservation.key, 0);
+        }
+        try {
+          return await requestWithIdempotentRetry(
+            exactReservations,
+            crypto.randomUUID(),
+          );
+        } catch (exactError) {
+          restorePreviousRemaining();
+          throw normalizeBudgetError(exactError);
+        }
+      }
+      restorePreviousRemaining();
+      throw normalized;
+    }
+  }
+}
+
+function normalizeBudgetError(error: unknown): Error {
+  const normalized = error instanceof Error ? error : new Error(String(error));
+  const match = /Play execution (.+) budget exceeded \((\d+)\/(\d+)\)/i.exec(
+    normalized.message,
+  );
+  return match
+    ? new BudgetStateLimitError(match[1]!, Number(match[2]), Number(match[3]))
+    : normalized;
+}

package/dist/bundling-sources/shared_libs/play-runtime/governor/budget-state-backend.ts

@@ -0,0 +1,43 @@
+import type {
+  PlayRunnerBudgetCharge,
+  PlayRunnerBudgetChargeResult,
+} from '../protocol';
+
+export interface BudgetStateBackend {
+  charge(input: {
+    rootRunId: string;
+    charges: readonly PlayRunnerBudgetCharge[];
+  }): Promise<PlayRunnerBudgetChargeResult>;
+}
+
+export class BudgetStateLimitError extends Error {
+  constructor(
+    readonly key: string,
+    readonly observed: number,
+    readonly limit: number,
+  ) {
+    super(`Play execution ${key} budget exceeded (${observed}/${limit}).`);
+    this.name = 'BudgetStateLimitError';
+  }
+}
+
+export class InMemoryBudgetStateBackend implements BudgetStateBackend {
+  private readonly counters = new Map<string, Record<string, number>>();
+
+  async charge(input: {
+    rootRunId: string;
+    charges: readonly PlayRunnerBudgetCharge[];
+  }): Promise<PlayRunnerBudgetChargeResult> {
+    const current = this.counters.get(input.rootRunId) ?? {};
+    const next = { ...current };
+    for (const charge of input.charges) {
+      const observed = (next[charge.key] ?? 0) + charge.amount;
+      if (observed > charge.limit) {
+        throw new BudgetStateLimitError(charge.key, observed, charge.limit);
+      }
+      next[charge.key] = observed;
+    }
+    this.counters.set(input.rootRunId, next);
+    return { counters: { ...next } };
+  }
+}

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

@@ -24,6 +24,11 @@ import {
   createInMemoryAdaptiveAdmission,
   type RuntimeAdaptiveAdmission,
 } from './adaptive-admission';
+import {
+  BudgetStateLimitError,
+  InMemoryBudgetStateBackend,
+  type BudgetStateBackend,
+} from './budget-state-backend';
 
 export interface WorkLease {
   /** Free the slot / pacing permit. Idempotent. MUST be called in a finally. */
@@ -121,7 +126,7 @@ export interface PlayExecutionGovernor {
   suggestedParallelism(toolId: string, fallback: number): Promise<number>;
 
   /** Increment a monotonic budget counter; throws GovernorBudgetError on breach. */
-  chargeBudget(kind: BudgetKind, amount?: number): void;
+  chargeBudget(kind: BudgetKind, amount?: number): Promise<void>;
 
   /**
    * Reserve depth + per-parent + descendant budget for a child play and return
@@ -140,7 +145,7 @@ export interface PlayExecutionGovernor {
   forkChild(input: {
     childPlayName: string;
     childRunId: string;
-  }): GovernanceSnapshot;
+  }): Promise<GovernanceSnapshot>;
 
   /** Effective row concurrency: explicit request clamped to [1, rowMax], else default. */
   resolveRowConcurrency(requested?: number): number;
@@ -156,6 +161,11 @@ export interface PlayExecutionGovernor {
     provider: string;
     latencyMs?: number | null;
   }): void;
+  /** Feed success back using the same tool-to-provider resolution as acquire. */
+  observeToolSuccess(input: {
+    toolId: string;
+    latencyMs?: number | null;
+  }): void;
 
   snapshot(): GovernanceSnapshot;
 }
@@ -164,6 +174,7 @@ interface GovernorInput {
   adapter: AdapterId;
   scope: { orgId: string; rootRunId: string };
   rateState: RateStateBackend;
+  budgetState?: BudgetStateBackend;
   resolvePacing: PacingResolver;
   adaptiveAdmission?: RuntimeAdaptiveAdmission;
   resume?: GovernanceSnapshot;
@@ -258,6 +269,8 @@ export function createPlayExecutionGovernor(
       initialRequestsPerSecond: policy.pacing.defaultProviderRequestsPerSecond,
       maxRequestsPerSecond: policy.pacing.suggestedMaxParallelism,
     });
+  const budgetState = input.budgetState ?? new InMemoryBudgetStateBackend();
+  const providerByTool = new Map<string, string>();
 
   const bucketId = (provider: string) => `${input.scope.orgId}:${provider}`;
 
@@ -268,63 +281,85 @@ export function createPlayExecutionGovernor(
     const resolvedPacing = await input.resolvePacing(toolId);
     const declaredPacing =
       resolvedPacing && resolvedPacing.rules.length > 0 ? resolvedPacing : null;
-    return adaptiveAdmission.resolvePacing({
+    const pacing = adaptiveAdmission.resolvePacing({
       orgId: input.scope.orgId,
       toolId,
       declaredPacing,
       fallbackPacing: defaultPacingForTool(toolId, policy),
     });
+    providerByTool.set(toolId, pacing.provider);
+    return pacing;
   }
 
-  function chargeBudget(kind: BudgetKind, amount = 1): void {
+  const budgetLimit = (kind: BudgetKind): number => {
+    switch (kind) {
+      case 'playCall':
+        return policy.budgets.maxPlayCallCount;
+      case 'toolCall':
+        return policy.budgets.maxToolCallCount;
+      case 'retry':
+        return policy.budgets.maxRetryCount;
+      case 'descendant':
+        return policy.budgets.maxDescendants;
+      case 'waterfallStep':
+        return policy.budgets.maxWaterfallStepExecutions;
+    }
+  };
+
+  async function reserveBudgets(
+    charges: Array<{ kind: BudgetKind; amount: number }>,
+  ): Promise<void> {
+    try {
+      await budgetState.charge({
+        rootRunId: state.rootRunId,
+        charges: charges.map((charge) => ({
+          key: charge.kind,
+          amount: charge.amount,
+          limit: budgetLimit(charge.kind),
+        })),
+      });
+    } catch (error) {
+      if (error instanceof BudgetStateLimitError) {
+        throw new GovernorBudgetError(
+          error.key as BudgetKind,
+          error.observed,
+          error.limit,
+        );
+      }
+      throw error;
+    }
+  }
+
+  async function chargeBudget(kind: BudgetKind, amount = 1): Promise<void> {
+    const current =
+      kind === 'playCall'
+        ? state.playCallCount
+        : kind === 'toolCall'
+          ? state.toolCallCount
+          : kind === 'retry'
+            ? state.retryCount
+            : kind === 'descendant'
+              ? state.descendantCount
+              : state.waterfallStepExecutions;
+    if (current + amount > budgetLimit(kind)) {
+      throw new GovernorBudgetError(kind, current + amount, budgetLimit(kind));
+    }
+    await reserveBudgets([{ kind, amount }]);
     switch (kind) {
       case 'playCall':
         state.playCallCount += amount;
-        if (state.playCallCount > policy.budgets.maxPlayCallCount)
-          throw new GovernorBudgetError(
-            'playCall',
-            state.playCallCount,
-            policy.budgets.maxPlayCallCount,
-          );
         return;
       case 'toolCall':
         state.toolCallCount += amount;
-        if (state.toolCallCount > policy.budgets.maxToolCallCount)
-          throw new GovernorBudgetError(
-            'toolCall',
-            state.toolCallCount,
-            policy.budgets.maxToolCallCount,
-          );
         return;
       case 'retry':
         state.retryCount += amount;
-        if (state.retryCount > policy.budgets.maxRetryCount)
-          throw new GovernorBudgetError(
-            'retry',
-            state.retryCount,
-            policy.budgets.maxRetryCount,
-          );
         return;
       case 'descendant':
         state.descendantCount += amount;
-        if (state.descendantCount > policy.budgets.maxDescendants)
-          throw new GovernorBudgetError(
-            'descendant',
-            state.descendantCount,
-            policy.budgets.maxDescendants,
-          );
         return;
       case 'waterfallStep':
         state.waterfallStepExecutions += amount;
-        if (
-          state.waterfallStepExecutions >
-          policy.budgets.maxWaterfallStepExecutions
-        )
-          throw new GovernorBudgetError(
-            'waterfallStep',
-            state.waterfallStepExecutions,
-            policy.budgets.maxWaterfallStepExecutions,
-          );
         return;
     }
   }
@@ -356,7 +391,7 @@ export function createPlayExecutionGovernor(
       // 3. charge the budget only once the call is actually cleared to run, so a
       // failed/aborted acquisition never permanently consumes tool budget.
       try {
-        chargeBudget('toolCall');
+        await chargeBudget('toolCall');
       } catch (error) {
         permit.release();
         slot.release();
@@ -388,7 +423,7 @@ export function createPlayExecutionGovernor(
 
     chargeBudget,
 
-    forkChild(childInput) {
+    async forkChild(childInput) {
       if (state.ancestryPlayIds.includes(childInput.childPlayName)) {
         const chain = [...state.ancestryPlayIds, childInput.childPlayName].join(
           ' -> ',
@@ -410,11 +445,53 @@ export function createPlayExecutionGovernor(
           nextParent,
           policy.budgets.maxChildPlayCallsPerParent,
         );
+      if (state.playCallCount + 1 > policy.budgets.maxPlayCallCount)
+        throw new GovernorBudgetError(
+          'playCall',
+          state.playCallCount + 1,
+          policy.budgets.maxPlayCallCount,
+        );
+      if (state.descendantCount + 1 > policy.budgets.maxDescendants)
+        throw new GovernorBudgetError(
+          'descendant',
+          state.descendantCount + 1,
+          policy.budgets.maxDescendants,
+        );
       // Charge the run-wide play/descendant budgets on the parent. Charged at
       // fork time (not after the caller's child-play slot acquire) and never
       // refunded if that acquire fails — see the forkChild interface doc.
-      chargeBudget('playCall');
-      chargeBudget('descendant');
+      try {
+        await budgetState.charge({
+          rootRunId: state.rootRunId,
+          charges: [
+            {
+              key: 'playCall',
+              amount: 1,
+              limit: policy.budgets.maxPlayCallCount,
+            },
+            {
+              key: 'descendant',
+              amount: 1,
+              limit: policy.budgets.maxDescendants,
+            },
+            {
+              key: `childPerParent:${state.currentRunId}`,
+              amount: 1,
+              limit: policy.budgets.maxChildPlayCallsPerParent,
+            },
+          ],
+        });
+      } catch (error) {
+        if (error instanceof BudgetStateLimitError) {
+          const budget = error.key.startsWith('childPerParent:')
+            ? 'childPerParent'
+            : (error.key as BudgetKind);
+          throw new GovernorBudgetError(budget, error.observed, error.limit);
+        }
+        throw error;
+      }
+      state.playCallCount += 1;
+      state.descendantCount += 1;
       state.parentChildCalls[parentKey] = nextParent;
       // Child seeds from the parent's accumulated counters → lineage-global budget.
       return {
@@ -456,6 +533,17 @@ export function createPlayExecutionGovernor(
       });
     },
 
+    observeToolSuccess(success) {
+      const provider =
+        providerByTool.get(success.toolId) ??
+        defaultPacingForTool(success.toolId, policy).provider;
+      adaptiveAdmission.observeProviderSuccess({
+        orgId: input.scope.orgId,
+        provider,
+        latencyMs: success.latencyMs,
+      });
+    },
+
     snapshot: () => ({
       ...state,
       ancestryPlayIds: [...state.ancestryPlayIds],