@nredd/pi-subagents 0.22.0-nredd.6 → 0.22.0-nredd.8

package/src/nested-tools.ts

@@ -17,8 +17,7 @@ import {
 } from "./agent-types.js";
 import { loadCustomAgents } from "./custom-agents.js";
 import { isolationParam, resolveAgentInvocationConfig } from "./invocation-config.js";
-import { resolveModel } from "./model-resolver.js";
-import { checkModelScope } from "./model-scope.js";
+import { type ModelSelection, modelSelectionNote, selectModel } from "./model-selection.js";
 import {
   createOutputFilePath,
   getOutputTranscriptDefault,
@@ -75,8 +74,10 @@ export const NESTED_WORKFLOW_NO_AGENT_ERROR =
   + "Do the work directly.";
 
 interface NestedSpawnOptions {
+  homeDir?: string;
   description: string;
   model?: Model<any>;
+  modelSelection?: ModelSelection;
   maxTurns?: number;
   isolated?: boolean;
   inheritContext?: boolean;
@@ -120,6 +121,8 @@ export interface NestedAgentManager {
 }
 
 export interface NestedToolContext {
+  /** Home shared by this branch's user-resource discovery. */
+  homeDir?: string;
   manager: NestedAgentManager;
   pi: ExtensionAPI;
   parentAgentId: string;
@@ -160,6 +163,13 @@ function ownsRecord(record: AgentRecord | undefined, parentAgentId: string): rec
 type ResultPosition = "inline" | "fetched";
 
 function formatRecord(record: AgentRecord, position: ResultPosition): string {
+  // Read off the record, so a fetched or resumed child reports what ran, not
+  // what the current call would have configured.
+  const note = modelSelectionNote(record);
+  return note ? `${note}\n\n${formatOutcome(record, position)}` : formatOutcome(record, position);
+}
+
+function formatOutcome(record: AgentRecord, position: ResultPosition): string {
   if (record.status === "error") {
     return `Agent failed: ${record.error ?? "unknown error"}${partialOutputSuffix(record)}`;
   }
@@ -181,7 +191,7 @@ export function createNestedSubagentTools(context: NestedToolContext): ToolDefin
   // Agents resolve from a registry built for THIS branch's config root (under
   // worktree isolation, the copy). Never via registerAgents — that is
   // process-global state shared with the main session and every other agent.
-  const loadRegistry = () => buildAgentRegistry(loadCustomAgents(context.configCwd, false, { projectTrusted: context.projectTrusted }));
+  const loadRegistry = () => buildAgentRegistry(loadCustomAgents(context.configCwd, false, { projectTrusted: context.projectTrusted, homeDir: context.homeDir }));
   const allowedTypesIn = (registry: Map<string, AgentConfig>): Set<string> | undefined =>
     context.allowedSubagents === "all"
       ? undefined
@@ -201,7 +211,7 @@ export function createNestedSubagentTools(context: NestedToolContext): ToolDefin
       prompt: Type.String({ description: "Self-contained task for the nested agent." }),
       description: Type.String({ description: "Short 3-5 word task description." }),
       subagent_type: Type.String({ description: `Allowed nested agent type. Available: ${availableIn(loadRegistry()).join(", ") || "none"}.` }),
-      model: Type.Optional(Type.String({ description: "Optional provider/model override." })),
+      model: Type.Optional(Type.String({ description: "Optional configured @tier or provider/model override; explicit values win over agentModels and the definition." })),
       thinking: Type.Optional(Type.String({ description: "Optional thinking level." })),
       max_turns: Type.Optional(Type.Number({ minimum: 1 })),
       run_in_background: Type.Optional(
@@ -221,6 +231,15 @@ export function createNestedSubagentTools(context: NestedToolContext): ToolDefin
         if (!ownsRecord(existing, context.parentAgentId)) {
           return textResult(`Nested agent not found or not owned by this parent: "${params.resume}".`, true);
         }
+        // Two resumes of one child in the same turn (parallel tool calls) would
+        // run two prompts on one session; the manager refuses the second, and
+        // this says why instead of a generic failure.
+        if (existing.status === "running" || existing.status === "queued") {
+          return textResult(
+            `Nested agent "${params.resume}" is still ${existing.status}; it can only be resumed once its current run finishes.`,
+            true,
+          );
+        }
         // Frontmatter authoritative, as on the spawn path; the manager falls
         // back to the limit the child was spawned with, then the default.
         const resumed = await context.manager.resume(params.resume, params.prompt, signal, {
@@ -264,40 +283,26 @@ export function createNestedSubagentTools(context: NestedToolContext): ToolDefin
       // Foreground regardless of `backgroundByDefault` — see the reasoning on
       // ResolveOptions. An explicit `true` here still opts in.
       const invocation = resolveAgentInvocationConfig(config, params, {
-        agentType: resolvedType,
         worktreeAllowed: isWorktreeIsolationEnabled(),
         defaultRunInBackground: false,
       });
-      let model = ctx.model;
-      if (invocation.modelInput) {
-        const resolvedModel = resolveModel(invocation.modelInput, ctx.modelRegistry);
-        if (typeof resolvedModel === "string") {
-          if (invocation.modelFromParams) return textResult(resolvedModel, true);
-        } else {
-          model = resolvedModel;
-        }
-      }
-
-      // Same scopeModels policy as the top-level Agent tool — a nested spawn
-      // must not escape the allowlist. A "warn" verdict proceeds silently:
-      // child sessions have no UI surface to toast to.
-      const scopeVerdict = checkModelScope({
-        model,
-        cwd: context.configCwd,
-        modelRegistry: ctx.modelRegistry,
-        callerSupplied: invocation.modelFromParams,
-        agentLabel: config?.displayName ?? resolvedType,
-        modelInput: invocation.modelInput,
+      const selection = selectModel({
+        homeDir: context.homeDir,
+        agentType: resolvedType, definitionModel: config?.model, paramModel: params.model,
+        parentModel: ctx.model, registry: ctx.modelRegistry, cwd: context.configCwd, agentLabel: config?.displayName,
       });
-      if (scopeVerdict.kind === "error") return textResult(scopeVerdict.message, true);
+      if (typeof selection === "string") return textResult(selection, true);
+      const model = selection.model;
 
       // The whole branch shares the root session's transcript directory; read it
       // off the owning parent rather than this child session's own id.
       const rootSessionId = context.manager.getRecord(context.parentAgentId)?.rootSessionId;
       const childDepth = context.depth + 1;
       const options: NestedSpawnOptions = {
+        homeDir: context.homeDir,
         description: params.description,
         model,
+        modelSelection: selection,
         maxTurns: invocation.maxTurns,
         isolated: invocation.isolated,
         inheritContext: invocation.inheritContext,

package/src/output-file.ts

@@ -5,10 +5,12 @@
  * matching Claude Code's task output file format.
  */
 
-import { appendFileSync, chmodSync, mkdirSync, writeFileSync } from "node:fs";
+import { appendFileSync, chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import type { AgentSession, AgentSessionEvent } from "@earendil-works/pi-coding-agent";
+import { type ModelFallback, modelFallbackNote } from "./model-resolver.js";
+import { type TierDisclosure, tierDisclosure, tierNote } from "./model-tiers.js";
 
 /**
  * Project/global default for writing a subagent's `.output` transcript; a custom
@@ -80,6 +82,26 @@ export function ensureOutputFile(path: string): void {
   } catch { /* ignore — streaming writes are best-effort too */ }
 }
 
+/** Append the model-fallback disclosure as a transcript line (best effort, like streaming writes). */
+export function appendModelFallbackEntry(path: string, agentId: string, fallback: ModelFallback): void {
+  try {
+    // A resumed agent reuses its transcript, which may already say it.
+    if (existsSync(path) && readFileSync(path, "utf-8").includes('"type":"model-fallback"')) return;
+    appendFileSync(
+      path,
+      `${JSON.stringify({ isSidechain: true, agentId, type: "model-fallback", modelFallback: fallback, note: modelFallbackNote(fallback), timestamp: new Date().toISOString() })}\n`,
+      "utf-8",
+    );
+  } catch { /* ignore */ }
+}
+
+/** Append tier provenance separately from assistant messages for durable replay. */
+export function appendModelTierEntry(path: string, agentId: string, modelTier: TierDisclosure): void {
+  try {
+    appendFileSync(path, `${JSON.stringify({ isSidechain: true, agentId, type: "model-tier", modelTier: tierDisclosure(modelTier), note: tierNote(modelTier), timestamp: new Date().toISOString() })}\n`, "utf-8");
+  } catch { /* transcript writes are best effort */ }
+}
+
 /** Write the initial user prompt entry. */
 export function writeInitialEntry(path: string, agentId: string, prompt: string, cwd: string): void {
   const entry = {

package/src/package-resources.ts

@@ -55,7 +55,8 @@
 
 import { existsSync, readFileSync, realpathSync, statSync } from "node:fs";
 import { basename, dirname, isAbsolute, join, relative, resolve, sep } from "node:path";
-import { DefaultPackageManager, getAgentDir, SettingsManager } from "@earendil-works/pi-coding-agent";
+import { DefaultPackageManager, SettingsManager } from "@earendil-works/pi-coding-agent";
+import { normalizeHomePath, resolveAgentDir } from "./home-paths.js";
 
 /** Resource kinds a package may declare. Agents are `.md`, workflows are `.js`. */
 export type PackageResourceKind = "agents" | "workflows";
@@ -103,6 +104,8 @@ export interface PiPackage {
 export type PackageGate = boolean | string[] | undefined;
 
 export interface PackageDiscoveryOptions {
+  /** Home directory for user settings and installed resources, resolved at call time. */
+  homeDir?: string;
   /**
    * Pi's project-trust state for `cwd`, from `ctx.isProjectTrusted()`. When
    * false, `.pi/settings.json` is not read — so a package configured only by an
@@ -376,14 +379,26 @@ interface ConfiguredPackageRow {
  * so a later pi that drops or renames this API contributes no packages instead
  * of failing activation.
  */
-function configuredPackages(cwd: string, projectTrusted: boolean): ConfiguredPackageRow[] {
+function configuredPackages(cwd: string, projectTrusted: boolean, homeDir?: string): ConfiguredPackageRow[] {
   try {
     if (typeof DefaultPackageManager !== "function") return [];
-    const agentDir = getAgentDir();
+    const agentDir = resolveAgentDir(homeDir);
     const settingsManager = SettingsManager.create(cwd, agentDir, { projectTrusted });
     const manager = new DefaultPackageManager({ cwd, agentDir, settingsManager });
     if (typeof manager.listConfiguredPackages !== "function") return [];
-    return manager.listConfiguredPackages();
+    if (homeDir === undefined) return manager.listConfiguredPackages();
+    // Pi's enumerator expands local package sources through its own ambient
+    // home. Reuse its public installed-path resolver with normalized inputs,
+    // preserving the original source for package-gate matching and diagnostics.
+    const rows: ConfiguredPackageRow[] = [];
+    for (const scope of ["user", "project"] as const) {
+      const settings = scope === "user" ? settingsManager.getGlobalSettings() : settingsManager.getProjectSettings();
+      for (const pkg of settings.packages ?? []) {
+        const source = typeof pkg === "string" ? pkg : pkg.source;
+        rows.push({ source, scope, installedPath: manager.getInstalledPath(normalizeHomePath(source, homeDir), scope) });
+      }
+    }
+    return rows;
   } catch {
     return [];
   }
@@ -394,10 +409,10 @@ function configuredPackages(cwd: string, projectTrusted: boolean): ConfiguredPac
  * order (global entries first, then project). Duplicates by root are collapsed —
  * a package listed in both scopes is one package on disk.
  */
-export function listPiPackages(cwd: string, projectTrusted = false): PiPackage[] {
+export function listPiPackages(cwd: string, projectTrusted = false, homeDir?: string): PiPackage[] {
   const seen = new Set<string>();
   const out: PiPackage[] = [];
-  for (const entry of configuredPackages(cwd, projectTrusted)) {
+  for (const entry of configuredPackages(cwd, projectTrusted, homeDir)) {
     const root = entry.installedPath;
     if (!root || seen.has(root)) continue;
     const manifest = readPackageJson(root);
@@ -484,9 +499,9 @@ export function invalidatePackageCache(): void {
   cache.clear();
 }
 
-function scan(cwd: string, projectTrusted: boolean, gate: PackageGate): CacheEntry {
+function scan(cwd: string, projectTrusted: boolean, gate: PackageGate, homeDir?: string): CacheEntry {
   const entry = emptyEntry();
-  for (const pkg of listPiPackages(cwd, projectTrusted)) {
+  for (const pkg of listPiPackages(cwd, projectTrusted, homeDir)) {
     if (!packageAllowed(pkg, gate)) continue;
     // The root is normalized here, not at lookup time: the declared paths below
     // go through `resolve(root, entry)`, so a root that arrived with a trailing
@@ -520,10 +535,10 @@ function cacheEntryFor(cwd: string, kind: PackageResourceKind, opts: PackageDisc
   // `\u0000` written as an escape, not a raw byte: a literal NUL makes this
   // whole file read as binary to grep, ripgrep and git, which silently drops
   // it from search results.
-  const key = `${cwd}\u0000${projectTrusted ? 1 : 0}\u0000${JSON.stringify(gate ?? true)}`;
+  const key = `${cwd}\u0000${opts.homeDir ?? ""}\u0000${resolveAgentDir(opts.homeDir)}\u0000${projectTrusted ? 1 : 0}\u0000${JSON.stringify(gate ?? true)}`;
   let hit = cache.get(key);
   if (!hit) {
-    hit = scan(cwd, projectTrusted, gate);
+    hit = scan(cwd, projectTrusted, gate, opts.homeDir);
     cache.set(key, hit);
   }
   return hit;

package/src/process-budget.ts

@@ -0,0 +1,346 @@
+/**
+ * process-budget.ts -- one ceiling on running subagents for the whole pi process.
+ *
+ * `maxConcurrent` and `maxConcurrentForeground` are per `AgentManager`, and
+ * there can be more than one manager per process: `/new`, `/resume`, `/fork`
+ * and `/reload` activate a fresh one and leave the previous manager alive until
+ * its background agents finish. (Subagent sessions do not activate this
+ * extension; nested children run in the parent's manager.) Those pools also
+ * exempt nested children, workflow children, scheduled jobs and resumes. The budget is the one count all of them share, kept on `globalThis`
+ * under `Symbol.for` so every manager and every module instance (a `/reload`
+ * loads a fresh copy of this file) sees the same state.
+ *
+ * Scope is the PROCESS, not the machine: two pi processes each get their own
+ * budget. It bounds running agents, not CPU, memory or commands.
+ *
+ * Every run holds one lease from the moment it is admitted until it actually
+ * settles. A parent keeps its lease while it waits on a child: releasing it
+ * around a nested call would be wrong, because pi runs a turn's tool calls in
+ * parallel and the parent's other tools are still executing.
+ *
+ * Holding parent leases is what makes a plain cap deadlock: twelve parents
+ * each waiting on a child that can never start. The reserve prevents it. With
+ * `L` the policy's `maxDepth` (one value for the whole process, NOT a branch's
+ * own cap) and `limit(j) = maxActive - nestedReserve * (L - j)`, a run at depth
+ * `d` is admitted only while, for EVERY level `j` from `d` to `L`,
+ *
+ *   active at depth <= j      < limit(j)
+ *
+ * (`j = L` is the plain `active < maxActive`). So with `maxActive: 12,
+ * nestedReserve: 2, L: 2`, at most ten top-level runs hold slots at once, and
+ * the last two only ever go to depth-2 children. Every admission keeps
+ * `count(<= j) <= limit(j)` for all `j`, releases only lower counts, and
+ * policy changes that could break it wait for an idle process (see
+ * `setBudgetPolicy`). Suppose every holder is a parent blocked on a child and
+ * `k < L` is the deepest level holding any: then `active = count(<= k) <=
+ * limit(k) < limit(k + 1) <= maxActive`, so a depth-`k+1` child fits at every
+ * level. Leaves (depth `L`) never wait on anyone, so they finish. That needs
+ * `limit` strictly increasing and `limit(1) >= 1` -- `isCoherent` -- and a run
+ * deeper than `L` refused outright (`assertWithinBudgetDepth`), since nothing
+ * reserves a slot for it. Top-level, workflow and scheduled runs all count as
+ * depth 1.
+ *
+ * Waiters are served first-come-first-served among the ELIGIBLE: a deeper
+ * waiter may pass a shallower one whose threshold is not met, which is what the
+ * reserve is for. Without that, a root at the head of the line would block the
+ * child that its own siblings are waiting on.
+ */
+
+const BUDGET_KEY = Symbol.for("pi-subagents:process-budget");
+
+/** Bumped on any incompatible change to `BudgetState`'s shape. */
+const BUDGET_VERSION = 2;
+
+/**
+ * Queued starts the budget will hold while it is enabled. Only a bound on
+ * memory and on runaway fan-out -- a full queue fails the spawn with
+ * `BUDGET_QUEUE_FULL_ERROR` instead of parking work nobody will see start.
+ */
+export const MAX_PENDING_SUBAGENTS = 128;
+
+export const DEFAULT_BUDGET_POLICY: Readonly<BudgetPolicy> = {
+  maxActive: 0,
+  nestedReserve: 2,
+  maxDepth: 2,
+};
+
+export interface BudgetPolicy {
+  /** Process-wide ceiling on running subagents. `0` = off (counted, never enforced). */
+  maxActive: number;
+  /** Slots held back per nesting level below a run, so children can always start. */
+  nestedReserve: number;
+  /**
+   * The deepest level the reserve covers: the GLOBAL `maxSubagentDepth`, so a
+   * project file cannot reshape a process-wide policy. A deeper run is refused.
+   */
+  maxDepth: number;
+}
+
+/** One admitted run. Released exactly once; later releases are no-ops. */
+export interface BudgetLease {
+  readonly depth: number;
+  released: boolean;
+}
+
+/** A start waiting for admission. */
+export interface BudgetWaiter {
+  /** Nesting depth of the run: top-level = 1. */
+  depth: number;
+  /** Local conditions the owner also needs (its own pool limits). */
+  ready: () => boolean;
+  /** Takes ownership of the lease. Must not throw; a throw is swallowed and the lease released. */
+  grant: (lease: BudgetLease) => void;
+}
+
+interface BudgetState {
+  version: number;
+  policy: BudgetPolicy;
+  /** A tightening change, applied once no lease is active (`setBudgetPolicy`). */
+  pendingPolicy: BudgetPolicy | undefined;
+  active: Set<BudgetLease>;
+  waiters: BudgetWaiter[];
+  pumping: boolean;
+  repump: boolean;
+}
+
+export const BUDGET_QUEUE_FULL_ERROR =
+  `Subagent queue is full (${MAX_PENDING_SUBAGENTS} waiting for the process-wide maxActiveSubagents limit). `
+  + "Wait for running agents to finish, or do the work directly.";
+
+function freshState(): BudgetState {
+  return {
+    version: BUDGET_VERSION,
+    policy: { ...DEFAULT_BUDGET_POLICY },
+    pendingPolicy: undefined,
+    active: new Set(),
+    waiters: [],
+    pumping: false,
+    repump: false,
+  };
+}
+
+/**
+ * The process's budget state. An incompatible version (a `/reload` onto a
+ * build that changed the shape) refuses budget operations until pi restarts.
+ * Never fork the state: two private counts could each admit the whole cap.
+ */
+function state(): BudgetState {
+  const g = globalThis as Record<symbol, unknown>;
+  const existing = g[BUDGET_KEY] as BudgetState | undefined;
+  if (existing === undefined) {
+    const created = freshState();
+    g[BUDGET_KEY] = created;
+    return created;
+  }
+  if (existing.version === BUDGET_VERSION) return existing;
+  throw new Error(
+    `[pi-subagents] Incompatible process budget v${String(existing.version)} already loaded (this build uses v${String(BUDGET_VERSION)}). `
+    + "Restart pi before starting more subagents; no private budget will be created.",
+  );
+}
+
+/** The deepest level `policy` reserves for; nesting off still means one level. */
+const topLevel = (policy: BudgetPolicy): number => Math.max(1, policy.maxDepth);
+
+/**
+ * How many runs at depth <= `level` may hold leases before a run at `level` or
+ * shallower has to wait. No floor: `isCoherent` keeps it at least one.
+ */
+export function admissionLimit(policy: BudgetPolicy, level: number): number {
+  if (policy.maxActive <= 0) return Number.POSITIVE_INFINITY;
+  return policy.maxActive - policy.nestedReserve * (topLevel(policy) - level);
+}
+
+/**
+ * Whether `policy` is deadlock-free: off, or `limit` strictly increasing from
+ * at least one at depth 1 -- a reserve of one or more, and room for one
+ * top-level run after reserving for every level below it.
+ */
+export function isCoherent(policy: BudgetPolicy): boolean {
+  if (policy.maxActive <= 0 || topLevel(policy) === 1) return true;
+  return policy.nestedReserve >= 1 && admissionLimit(policy, 1) >= 1;
+}
+
+/** Every level's limit is at least what it was, so no held lease can block a child that fit before. */
+function relaxes(from: BudgetPolicy, to: BudgetPolicy): boolean {
+  if (to.maxActive <= 0) return true;
+  if (from.maxActive <= 0 || topLevel(from) !== topLevel(to)) return false;
+  for (let level = 1; level <= topLevel(to); level++) {
+    if (admissionLimit(to, level) < admissionLimit(from, level)) return false;
+  }
+  return true;
+}
+
+export function getBudgetPolicy(): BudgetPolicy {
+  return { ...state().policy };
+}
+
+/**
+ * Replace the policy. Last writer wins: every activation in the process
+ * shares it. An incoherent policy is refused with a warning. A change that
+ * could strand a waiting child -- turning the budget on, lowering a limit,
+ * changing `maxDepth` -- is held until no lease is active, since running
+ * parents were admitted under the old limits; turning it off or raising
+ * limits applies at once.
+ */
+export function setBudgetPolicy(next: BudgetPolicy): void {
+  const policy: BudgetPolicy = {
+    maxActive: Math.max(0, Math.floor(next.maxActive)),
+    nestedReserve: Math.max(0, Math.floor(next.nestedReserve)),
+    maxDepth: Math.max(0, Math.floor(next.maxDepth)),
+  };
+  if (!isCoherent(policy)) {
+    console.warn(
+      `[pi-subagents] Ignoring maxActiveSubagents ${String(policy.maxActive)} with nestedSubagentReserve ${String(policy.nestedReserve)} `
+      + `at maxSubagentDepth ${String(policy.maxDepth)}: it could deadlock. Nesting needs a reserve of at least 1 and `
+      + "maxActiveSubagents >= nestedSubagentReserve * (maxSubagentDepth - 1) + 1.",
+    );
+    return;
+  }
+  const s = state();
+  if (s.active.size === 0 || relaxes(s.policy, policy)) {
+    s.policy = policy;
+    s.pendingPolicy = undefined;
+    pumpBudget();
+  } else {
+    s.pendingPolicy = policy;
+    console.warn(
+      `[pi-subagents] Process budget change deferred until all running subagents settle: `
+      + `maxActiveSubagents=${String(policy.maxActive)}, nestedSubagentReserve=${String(policy.nestedReserve)}, `
+      + `maxSubagentDepth=${String(policy.maxDepth)}. The previous policy remains in force; continuous work can delay this indefinitely.`,
+    );
+  }
+}
+
+/** The policy waiting for the process to go idle, if any. For diagnostics and tests. */
+export function getPendingBudgetPolicy(): BudgetPolicy | undefined {
+  const pending = state().pendingPolicy;
+  return pending === undefined ? undefined : { ...pending };
+}
+
+/**
+ * Throw when the budget is on and a run at `depth` is deeper than its policy
+ * reserves for: it could wait forever behind its own ancestors. Pass/fail is
+ * fixed for the run's life, since `maxDepth` changes only while idle.
+ */
+export function assertWithinBudgetDepth(depth: number): void {
+  const { policy } = state();
+  if (policy.maxActive > 0 && depth > topLevel(policy)) {
+    throw new Error(
+      `Nested subagent at depth ${String(depth)} exceeds the process budget's depth ${String(topLevel(policy))} `
+      + "(global maxSubagentDepth). Complete the task directly.",
+    );
+  }
+}
+
+/** Running leases and waiting starts, for diagnostics and tests. */
+export function budgetUsage(): { active: number; pending: number; maxActive: number } {
+  const s = state();
+  return { active: s.active.size, pending: s.waiters.length, maxActive: s.policy.maxActive };
+}
+
+function admits(s: BudgetState, depth: number): boolean {
+  const { policy } = s;
+  if (policy.maxActive <= 0) return true;
+  const top = topLevel(policy);
+  const level = Math.max(1, depth);
+  if (level > top) return false;
+  // Leases deeper than `top` cannot exist while it is on; clamp so they still count in the total.
+  const atDepth = new Array<number>(top + 1).fill(0);
+  for (const lease of s.active) {
+    const index = Math.min(top, Math.max(1, lease.depth));
+    atDepth[index] = atDepth[index]! + 1;
+  }
+  let atOrAbove = 0;
+  for (let j = 1; j <= top; j++) {
+    atOrAbove += atDepth[j]!;
+    if (j >= level && atOrAbove >= admissionLimit(policy, j)) return false;
+  }
+  return true;
+}
+
+function newLease(s: BudgetState, depth: number): BudgetLease {
+  const lease: BudgetLease = { depth, released: false };
+  s.active.add(lease);
+  return lease;
+}
+
+/**
+ * Take a lease now, or undefined when this depth is at its limit or an earlier
+ * eligible waiter has first claim. Completion listeners can spawn synchronously
+ * between release and drain; re-entrant grant callbacks can do the same.
+ */
+export function tryAcquireBudget(depth: number): BudgetLease | undefined {
+  const s = state();
+  if (!admits(s, depth) || s.waiters.some(waiter => admits(s, waiter.depth) && waiter.ready())) return undefined;
+  return newLease(s, depth);
+}
+
+/** Whether another waiter fits. Always true while the budget is off. */
+export function budgetHasPendingRoom(): boolean {
+  const s = state();
+  return s.policy.maxActive <= 0 || s.waiters.length < MAX_PENDING_SUBAGENTS;
+}
+
+/** Join the queue. Returns a cancel function; cancelling a granted waiter is a no-op. */
+export function enqueueBudget(waiter: BudgetWaiter): () => void {
+  state().waiters.push(waiter);
+  return () => {
+    const s = state();
+    const i = s.waiters.indexOf(waiter);
+    if (i !== -1) s.waiters.splice(i, 1);
+  };
+}
+
+/**
+ * Give a lease back. Idempotent; does not pump, so callers keep their own
+ * ordering. The last release applies a held policy change.
+ */
+export function releaseBudget(lease: BudgetLease | undefined): void {
+  if (lease === undefined || lease.released) return;
+  lease.released = true;
+  const s = state();
+  s.active.delete(lease);
+  if (s.active.size === 0 && s.pendingPolicy !== undefined) {
+    s.policy = s.pendingPolicy;
+    s.pendingPolicy = undefined;
+  }
+}
+
+/**
+ * Admit every waiter that fits, earliest first, re-scanning after each grant
+ * because a grant changes the counts (`ready()` reads the owner's pools).
+ * Re-entrant calls from inside a grant are folded into the running loop.
+ */
+export function pumpBudget(): void {
+  const s = state();
+  if (s.pumping) {
+    s.repump = true;
+    return;
+  }
+  s.pumping = true;
+  try {
+    do {
+      s.repump = false;
+      for (;;) {
+        const i = s.waiters.findIndex(w => admits(s, w.depth) && w.ready());
+        if (i === -1) break;
+        const waiter = s.waiters.splice(i, 1)[0]!;
+        const lease = newLease(s, waiter.depth);
+        try {
+          waiter.grant(lease);
+        } catch {
+          releaseBudget(lease);
+        }
+      }
+    } while (s.repump);
+  } finally {
+    s.pumping = false;
+  }
+}
+
+/** Test-only: forget all leases and waiters and restore the default policy. */
+export function resetBudgetForTests(): void {
+  const g = globalThis as Record<symbol, unknown>;
+  g[BUDGET_KEY] = freshState();
+}

package/src/quota-admission.ts

@@ -4,7 +4,7 @@
  * Quota tracking lives in a separate extension (nredd/pi-subscription-router),
  * reached over `pi.events` with the same envelope as `cross-extension-rpc.ts`:
  * emit `router:rpc:admit` `{ requestId, model }`, reply on
- * `router:rpc:admit:reply:<requestId>` as `{ success, data?: { ok, resolvedModel?, resetAt?, wait? } }` (`resolvedModel` is ignored).
+ * `router:rpc:admit:reply:<requestId>` as `{ success, data?: { ok, resolvedModel?, resetAt?, wait? } }` (`resolvedModel` labels the admitted physical model).
  *
  * Admission fails open. No router installed, a slow router, a malformed reply
  * or a router error all admit the dispatch: quota is advisory here, and the
@@ -13,6 +13,7 @@
 
 import { nanoid } from "nanoid";
 import { type Clock, realClock } from "./clock.js";
+import type { TierAdmission } from "./model-tiers.js";
 
 export const ADMIT_CHANNEL = "router:rpc:admit";
 
@@ -50,34 +51,45 @@ export function modelRef(model: { provider: string; id: string }): string {
 }
 
 /** Resolve to the block when the router refuses `model`, else undefined. */
-export function checkAdmission(
+export async function checkAdmission(
   events: AdmissionEvents,
   model: string,
   timeoutMs = ADMIT_TIMEOUT_MS,
   clock: Clock = realClock,
 ): Promise<QuotaBlock | undefined> {
+  return (await admissionOutcome(events, model, timeoutMs, clock)).block;
+}
+
+/** Retain the physical admission label and whether the router supplied a signal. */
+export function admissionOutcome(
+  events: AdmissionEvents,
+  model: string,
+  timeoutMs = ADMIT_TIMEOUT_MS,
+  clock: Clock = realClock,
+): Promise<TierAdmission> {
   const requestId = nanoid();
   return new Promise(resolve => {
     let done = false;
-    const finish = (block: QuotaBlock | undefined) => {
+    const finish = (outcome: TierAdmission) => {
       if (done) return;
       done = true;
       timer.cancel();
       unsubscribe();
-      resolve(block);
+      resolve(outcome);
     };
     const unsubscribe = events.on(`${ADMIT_CHANNEL}:reply:${requestId}`, raw => {
-      const reply = raw as { success?: unknown; data?: { ok?: unknown; resetAt?: unknown; wait?: unknown } } | undefined;
-      if (reply?.success !== true || reply.data?.ok !== false) return finish(undefined);
+      const reply = raw as { success?: unknown; data?: { ok?: unknown; resolvedModel?: unknown; resetAt?: unknown; wait?: unknown } } | undefined;
+      if (reply?.success !== true || typeof reply.data?.ok !== "boolean") return finish({ quotaSignal: false });
+      if (reply.data.ok) return finish({ quotaSignal: true, resolvedModel: typeof reply.data.resolvedModel === "string" ? reply.data.resolvedModel : undefined });
       const resetAt = reply.data.resetAt;
       const wait = reply.data.wait;
-      finish({
+      finish({ quotaSignal: true, block: {
         model,
         resetAt: typeof resetAt === "number" && Number.isFinite(resetAt) ? resetAt : undefined,
         ...(typeof wait === "boolean" ? { wait } : {}),
-      });
+      } });
     });
-    const timer = clock.setTimeout(() => finish(undefined), timeoutMs);
+    const timer = clock.setTimeout(() => finish({ quotaSignal: false }), timeoutMs);
     events.emit(ADMIT_CHANNEL, { requestId, model });
   });
 }