@mono-agent/agent-runtime 0.20.11 → 0.21.0

package/src/ai/providers/pi-native.js

@@ -55,12 +55,15 @@ import {
   buildErrorDetails,
   buildErrorResult,
   buildSuccessResult,
+  failureKindForPiError,
   emitCapabilitiesResolved,
   emitUsageCostEvents,
   usageFromMessages,
+  hasMeasuredUsage,
   withSubagentUsage,
 } from "./pi-native/result-builder.js";
 import {
+  captureSessionRecovery,
   cleanupSessionOnThrow,
   commitSession,
   discardUncommittedSession,
@@ -74,18 +77,22 @@ import {
   runReactiveCompaction,
 } from "./pi-native/compaction-driver.js";
 import {
+  activateTurnHarness,
   buildTurnHarness,
   buildTurnTools,
+  createLiveInputPromptEpoch,
   runHarnessPrompt,
   startLiveInput,
   thinkingLevelForEffort,
 } from "./pi-native/turn-runner.js";
 import { resolvePiTransport } from "./pi-native/transport.js";
+import { withOpenCodeSessionHeaders } from "./pi-native/provider-attribution.js";
 
 /**
- * Pi 0.80.6 exposes per-tool execution markers but not AgentHarness's global
- * toolExecution option. Resolve mono-agent's programmatic mode once per run;
- * individual tool builders then mark stateful/mutating tools sequential.
+ * Resolve mono-agent's programmatic mode once per run. Tool builders mark
+ * stateful/mutating tools sequential; the Pi 0.85 harness adapter promotes any
+ * such marker to its global toolExecution setting because mixed scheduling is
+ * no longer available upstream.
  */
 export function resolvePiToolExecutionMode(options = {}) {
   const warnings = [];
@@ -188,23 +195,82 @@ export function createDynamicCredentialStore(apiKeys, resolvePiApiKey, runtimeWa
 // resolved model. `piResolvedModels` is an advanced/test seam mirroring
 // `piResolvedModel`: when supplied it is used verbatim (the model dispatched via
 // `piResolvedModel` may live outside pi's builtin catalog, e.g. a faux model).
-function buildRunModels(runtime, options, runtimeWarnings) {
-  if (options.piResolvedModels) return options.piResolvedModels;
-  const credentials = createDynamicCredentialStore(runtime.apiKeys, options.resolvePiApiKey, runtimeWarnings);
-  if (options.customProvider) {
-    const model = runtime.model;
-    const models = createModels({ credentials });
-    models.setProvider(createProvider({
-      id: model.provider,
-      name: model.name || model.provider,
-      baseUrl: model.baseUrl,
-      auth: { apiKey: envApiKeyAuth(model.name || model.provider, []) },
-      models: [model],
-      api: openAICompletionsApi(),
-    }));
-    return models;
+function buildRunModels(runtime, options, runtimeWarnings, providerAttributionSessionId) {
+  let models;
+  if (options.piResolvedModels) {
+    models = options.piResolvedModels;
+  } else {
+    const credentials = createDynamicCredentialStore(runtime.apiKeys, options.resolvePiApiKey, runtimeWarnings);
+    if (options.customProvider) {
+      const model = runtime.model;
+      models = createModels({ credentials });
+      models.setProvider(createProvider({
+        id: model.provider,
+        name: model.name || model.provider,
+        baseUrl: model.baseUrl,
+        auth: { apiKey: envApiKeyAuth(model.name || model.provider, []) },
+        models: [model],
+        api: openAICompletionsApi(),
+      }));
+    } else {
+      models = builtinModels({
+        credentials,
+        ...(options.providerCheckAuthContext === undefined
+          ? {}
+          : { authContext: options.providerCheckAuthContext }),
+      });
+    }
   }
-  return builtinModels({ credentials });
+  return withProviderCheckOutputCap(
+    withOpenCodeSessionHeaders(models, providerAttributionSessionId),
+    options.providerCheckMaxTokens,
+  );
+}
+
+const PROVIDER_CHECK_REQUEST_METHODS = new Set([
+  "stream",
+  "complete",
+  "streamSimple",
+  "completeSimple",
+  "streamDeferred",
+  "fetchDeferred",
+  "cancelDeferred",
+]);
+
+/**
+ * Pi's Agent resolves the transport model through `Models` again, so capping
+ * only the model handed to the harness does not constrain the actual provider
+ * request. Bind the same cap at the dispatcher boundary used by every request
+ * path. This wrapper is activated only for explicit provider checks.
+ *
+ * @param {import("@earendil-works/pi-ai").Models} models
+ * @param {unknown} requestedCap
+ * @returns {import("@earendil-works/pi-ai").Models}
+ */
+function withProviderCheckOutputCap(models, requestedCap) {
+  const cap = Number(requestedCap);
+  if (!Number.isSafeInteger(cap) || cap <= 0) return models;
+  const wrappers = new Map();
+  return /** @type {import("@earendil-works/pi-ai").Models} */ (new Proxy(models, {
+    get(target, property) {
+      const value = Reflect.get(target, property, target);
+      if (typeof property !== "string" || typeof value !== "function") return value;
+      if (!PROVIDER_CHECK_REQUEST_METHODS.has(property)) return value.bind(target);
+      let wrapper = wrappers.get(property);
+      if (wrapper === undefined) {
+        wrapper = (model, ...args) => {
+          const current = Number(model?.maxTokens);
+          const cappedModel = {
+            ...model,
+            maxTokens: Number.isFinite(current) && current > 0 ? Math.min(current, cap) : cap,
+          };
+          return value.call(target, cappedModel, ...args);
+        };
+        wrappers.set(property, wrapper);
+      }
+      return wrapper;
+    },
+  }));
 }
 
 // Normalize the incoming runtime messages into AgentMessages the harness can
@@ -258,6 +324,14 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
   // an intermediate layer flattens the Error to its message.
   installTransportErrorProbe();
   const resolved = options.model;
+  if (
+    typeof resolved?.provider !== "string"
+    || resolved.provider.length === 0
+    || typeof resolved.model !== "string"
+    || resolved.model.length === 0
+  ) {
+    throw new Error("invalid pi model reference: provider and model are required");
+  }
   const start = Date.now();
   const events = [];
   const runtimeWarnings = [];
@@ -303,7 +377,14 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       diagnostics: {},
     },
     session: null,
+    // Fresh stateless calls own a private repo, even when attribution matches a
+    // live primary session. Cleanup must never delete that primary's record.
+    ephemeralSessionRepo: null,
     sessionEntry: null,
+    // Fresh keep-alive sessions are registered busy until the harness closes;
+    // this prevents another turn from reopening the same repo record while its
+    // first handle is still live.
+    registeredSessionEntry: null,
     // True when a requestedSessionId had no live entry AND no durable transcript,
     // so a fresh durable session was created under that id (cross-restart resume,
     // first turn for the conversation). Distinct from a true resume (sessionEntry
@@ -322,12 +403,17 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
     // (host/runtime-side throws after the session mutated) can roll back too, not
     // just the success path.
     baselineLeafId: null,
+    recoveryOperationId: undefined,
+    recoveryInputIds: null,
+    retainRecoveryTail: false,
   };
 
   const providerSessionId = options.sessionId
     || options.providerSessionId
+    || options.providerAttributionSessionId
     || options.runId
     || randomUUID();
+  const providerAttributionSessionId = options.providerAttributionSessionId || providerSessionId;
   // Prefer the explicit sessionId, but fall back to providerSessionId so a caller
   // that only supplies providerSessionId still resumes the prior session instead
   // of being treated as a fresh run (which would drop prior context).
@@ -352,13 +438,17 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
     }
     emitCaptured(events, options.onEvent, event);
   };
+  // Exec and Monitor both run a command a host configured Bash's risk tier for.
+  // Leaving either unmapped would let the same shell command take a lower
+  // approval tier simply by being started through a different tool.
+  const bashRiskTier = options.toolRiskTiers?.Bash;
   const approvalRiskTiers = {
     ...(options.toolRiskTiers || {}),
-    ...(
-      options.toolRiskTiers?.Exec === undefined && options.toolRiskTiers?.Bash !== undefined
-        ? { Exec: options.toolRiskTiers.Bash }
-        : {}
-    ),
+    ...(bashRiskTier !== undefined
+      ? Object.fromEntries(["Exec", "Monitor"]
+        .filter((name) => options.toolRiskTiers?.[name] === undefined)
+        .map((name) => [name, bashRiskTier]))
+      : {}),
   };
   const approvalManager = options.onToolApprovalRequest
     ? createApprovalManager({
@@ -405,7 +495,7 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
     // ready pi-ai Model (e.g. a registered faux provider model) plus optional
     // capabilities, bypassing the static model-registry lookup. Production
     // callers leave it undefined and resolve through pi-ai's registry.
-    const runtime = options.piResolvedModel
+    let runtime = options.piResolvedModel
       ? {
         model: options.piResolvedModel,
         capabilities: options.piResolvedCapabilities || {
@@ -420,10 +510,18 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
         apiKeys: new Map(),
       }
       : resolvePiRuntimeModel(resolved, options);
+    if (Number.isSafeInteger(options.providerCheckMaxTokens) && options.providerCheckMaxTokens > 0) {
+      runtime = {
+        ...runtime,
+        model: {
+          ...runtime.model,
+          maxTokens: Math.min(runtime.model.maxTokens, options.providerCheckMaxTokens),
+        },
+      };
+    }
     const capabilities = runtime.capabilities || {};
     const effectiveThinkingLevel = thinkingLevelForEffort(options.effort || "medium", capabilities);
-    const reference = resolved.reference
-      || (resolved.sdk === "pi" ? `pi:${resolved.provider}:${resolved.model}` : `${resolved.sdk}:${resolved.model}`);
+    const reference = resolved.reference || `${resolved.provider}:${resolved.model}`;
 
     // Resolve the typed `toolLimits` / `compaction` policy objects against the
     // deprecated `settings` fallback (per-group precedence: a present typed
@@ -486,15 +584,17 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       ? Number(options.maxRetryDelayMs)
       : 60_000;
     // Steering controls user follow-up delivery, not tool scheduling. Keep it
-    // independent from piToolExecutionMode; tools carry their own executionMode.
+    // independent from piToolExecutionMode; the adapter derives Pi 0.85's
+    // global scheduling mode from the tools' executionMode markers.
     const toolSteeringMode = "one-at-a-time";
 
-    const piModels = buildRunModels(runtime, options, runtimeWarnings);
+    const piModels = buildRunModels(runtime, options, runtimeWarnings, providerAttributionSessionId);
 
-    // Construct the harness, subscribe the stream normalizer, and wire the
-    // external abort handler (which sets runState.externalAbort and aborts the
-    // harness). Sets runState.harness + runState.removeAbortHandler.
-    harness = buildTurnHarness(runState, {
+    // Construct the harness and recover any interrupted durable operation.
+    // Stream subscription and the external abort handler are activated only
+    // after prior transcript seeding, so manual append lifecycle events cannot
+    // surface as output from this run.
+    harness = await buildTurnHarness(runState, {
       session: runState.session,
       piModels,
       model: runtime.model,
@@ -506,11 +606,7 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       maxRetries,
       maxRetryDelayMs,
       steeringMode: toolSteeringMode,
-      onEvent,
       options,
-      toolLimits,
-      sdk: resolved.sdk,
-      reference,
     });
 
     // Seed prior transcript (everything before the trailing user turn) into the
@@ -536,10 +632,14 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       try { runState.baselineLeafId = await runState.session.getLeafId(); } catch { /* best-effort */ }
     }
 
-    // Live steering: consume follow-up messages and steer the harness mid-run.
-    // The consumer is tied to run completion so it stops steering once the run
-    // finishes and does not swallow messages meant for a later turn.
-    const liveInput = startLiveInput({ harness, options, onEvent });
+    activateTurnHarness(runState, {
+      harness,
+      onEvent,
+      options,
+      toolLimits,
+      sdk: "pi",
+      reference,
+    });
 
     // Re-check abort right before issuing the provider request. The abort
     // handler is only installed at ~:639, AFTER a long stretch of awaited setup
@@ -582,19 +682,53 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       runtimeWarnings,
     });
 
+    // Arm the main-prompt epoch after proactive compaction so compaction and
+    // transcript seeding cannot be mistaken for live-input consumption.
+    runState.recoveryBaselineTipId = await runState.session.getLeafId();
+    const liveInputEpoch = createLiveInputPromptEpoch({ harness, onEvent });
+    const liveInput = startLiveInput({ harness, options, onEvent, promptEpoch: liveInputEpoch });
+
+    // Report the live harness value, not a downstream recreation of Pi's
+    // thinking-level normalization. This also captures any future adapter-side
+    // adjustment between construction and the provider request.
+    const providerEffectiveEffort = harness.getThinkingLevel();
+    onEvent({
+      type: "provider_execution_config",
+      sdk: "pi",
+      model: reference,
+      ...(options.effort ? { effort: options.effort } : {}),
+      effectiveEffort: providerEffectiveEffort,
+      timestamp: Date.now(),
+    });
+
     onEvent({
       type: "provider_request_started",
-      sdk: resolved.sdk,
+      sdk: "pi",
       model: reference,
       runtime: "pi",
       timestamp: Date.now(),
     });
 
-    let { runError } = await runHarnessPrompt(harness, promptText, promptImages);
-
-    // The run is done: stop the live-steering consumer so it cannot steer a
-    // finished harness or swallow a follow-up meant for the next turn.
-    await liveInput.stop();
+    let runError;
+    try {
+      // Own the operation at admission so steers consumed mid-run settle as
+      // applied immediately; finish() re-checks the settled id at the end.
+      const promptResult = await runHarnessPrompt(harness, promptText, promptImages, {
+        onOperationAdmitted: (operationId) => {
+          runState.recoveryOperationId = operationId;
+          liveInputEpoch.confirm(operationId);
+        },
+      });
+      runError = promptResult.runError;
+      liveInputEpoch.finish(promptResult.operationId);
+    } finally {
+      // Stop joins unresolved native enqueue and reconciles every returned
+      // entry before the exact-operation subscription is removed.
+      try { await liveInput.stop(); } finally {
+        runState.recoveryInputIds = liveInputEpoch.consumedInputIds();
+        liveInputEpoch.close();
+      }
+    }
 
     runState.externalAbort ||= !!options.abortSignal?.aborted;
 
@@ -765,7 +899,7 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       // runtime.compaction.enabled. See docs/reference/feature-registry.md runtime.context-compaction.
       contextCompactionApplied: runState.compaction.policy?.enabled ? runState.compaction.applied : null,
     });
-    emitCapabilitiesResolved(onEvent, { sdk: resolved.sdk, model: reference, capabilitiesUsed });
+    emitCapabilitiesResolved(onEvent, { sdk: "pi", model: reference, capabilitiesUsed });
 
     // Re-check the abort signal: a cancel can land during the post-run work above
     // (live-input teardown, structured-output finalization retry) after the line
@@ -780,6 +914,12 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
     // tracks LIVENESS so disposeProviderSession / idle-TTL eviction can reach
     // native sessions, keep-alive registers the session, and a failed/aborted
     // resumed turn rolls back to its pre-turn leaf.
+    runState.retainRecoveryTail = !!durableRepo && options.sessionKeepAlive === true
+      && !runState.maxTurnsHit
+      && (!errorMessage || failureKindForPiError(errorMessage, diagnostics) === "provider_unavailable")
+      && typeof options.sessionRecovery?.runId === "string" && options.sessionRecovery.runId.length > 0
+      && Number.isSafeInteger(options.sessionRecovery?.revision) && options.sessionRecovery.revision >= 0
+      && typeof runState.recoveryOperationId === "string";
     await commitSession(runState, {
       options,
       requestedSessionId,
@@ -791,21 +931,21 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       onEvent,
     });
 
-    // Final abort guard (durable cancel TOCTOU): if a cancel raced the lifecycle
-    // commit above — landing AFTER the keep-alive/!externalAbort decision but
-    // before this return — the cancelled turn is still in the durable transcript
-    // and (for keep-alive) the live registry. Roll it back so the next resume sees
-    // the pre-turn state (rollbackAbortedTurn: a resumed session moves to its
-    // baseline leaf and drops its live entry; a fresh durable session deletes its
-    // jsonl). The abort re-check + return stay inline with NO await between the
-    // false-branch check and the return, so an external cancel cannot newly fire
-    // past it (I10).
+    // Uncoordinated calls keep their final rollback guard. Coordinated calls
+    // retain the tail and close it before granting a recovery receipt; recheck
+    // cancellation after that await as well.
     if (!runState.externalAbort && options.abortSignal?.aborted) {
       runState.externalAbort = true;
-      await rollbackAbortedTurn(runState, { requestedSessionId, providerSessionId, durableRepo });
+      if (!runState.retainRecoveryTail) await rollbackAbortedTurn(runState, { requestedSessionId, providerSessionId, durableRepo });
     }
 
-    return buildSuccessResult({
+    // A receipt is proof of a successfully closed native operation, never a
+    // promise that the finally block will eventually close it.
+    const providerSessionRecovery = runState.retainRecoveryTail
+      ? await captureSessionRecovery(runState, { options, providerSessionId, modelKey: `${resolved.provider}:${resolved.model}`, model: runtime.model, pending: !!errorMessage || runState.externalAbort })
+      : undefined;
+    if (runState.retainRecoveryTail) runState.externalAbort ||= !!options.abortSignal?.aborted;
+    return { ...buildSuccessResult({
       finalText,
       finalThinking,
       events,
@@ -824,8 +964,10 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       providerSessionId,
       runtimeWarnings,
       capabilitiesUsed,
+      usageMeasured: hasMeasuredUsage(runTranscript),
       structuredResult: runState.structuredResult,
-    });
+      effectiveEffort: providerEffectiveEffort,
+    }), ...(providerSessionRecovery ? { providerSessionRecovery } : {}) };
   } catch (err) {
     runState.externalAbort ||= !!options.abortSignal?.aborted;
     // Drop a just-created fresh durable session, release a create-on-miss
@@ -859,9 +1001,12 @@ export async function generatePiNativeResponse(systemPrompt, options = {}) {
       runtimeWarnings,
       isRetryable,
       piTransport,
+      effectiveEffort: harness?.getThinkingLevel?.(),
     });
   } finally {
+    try { await harness?.close?.(); } catch { /* best-effort */ }
     if (runState.sessionEntry) runState.sessionEntry.busy = false;
+    if (runState.registeredSessionEntry) runState.registeredSessionEntry.busy = false;
     runState.removeAbortHandler?.();
     try {
       await closeRunTools();
@@ -875,6 +1020,6 @@ export const piNativeRuntimeBridge = {
   id: "pi",
   kind: "pi",
   capabilities: runtimeCapabilities("pi"),
-  supports: (ref) => ref?.sdk === "pi",
+  supports: (ref) => Boolean(ref?.provider && ref?.model),
   execute: generatePiNativeResponse,
 };

package/src/ai/runtime/capabilities.js

@@ -1,7 +1,4 @@
-import {
-  TOOL_POLICY_ALLOW_ALL_ONLY,
-  TOOL_POLICY_PROJECTED,
-} from "./tool-policy.js";
+import { TOOL_POLICY_PROJECTED } from "./tool-policy.js";
 
 export const COMMON_CAPABILITIES = {
   streaming: true,
@@ -15,18 +12,9 @@ export const COMMON_CAPABILITIES = {
   supports_live_input: true,
   supports_native_subagents: true,
   supports_request_tool_environment: false,
-  supports_fast_mode: false,
-  tool_policy: TOOL_POLICY_PROJECTED,
 };
 
 export const RUNTIME_CAPABILITIES = {
-  claude: {
-    runtime: "sdk",
-    ...COMMON_CAPABILITIES,
-    // The Claude Agent SDK persists sessions on disk; the bridge resumes
-    // them via queryOptions.resume when options.sessionId is supplied.
-    supports_session_resume: true,
-  },
   pi: {
     runtime: "pi-agent",
     ...COMMON_CAPABILITIES,
@@ -38,50 +26,17 @@ export const RUNTIME_CAPABILITIES = {
     supports_native_subagents: false,
     supports_request_tool_environment: true,
     supports_mcp_apps: true,
-  },
-  codex: {
-    runtime: "cli",
-    ...COMMON_CAPABILITIES,
-    // The codex-app bridge keeps the app-server subprocess + thread alive
-    // when options.sessionKeepAlive is set; resumed turns reuse the thread.
-    supports_session_resume: true,
-    supports_fast_mode: true,
-    tool_policy: TOOL_POLICY_ALLOW_ALL_ONLY,
-  },
-  opencode: {
-    runtime: "cli",
-    ...COMMON_CAPABILITIES,
-    structured_output: false,
-    supports_session_resume: false,
-    supports_mcp: false,
-    supports_skills: false,
-    supports_live_input: false,
-    supports_native_subagents: false,
-    tool_policy: TOOL_POLICY_ALLOW_ALL_ONLY,
-  },
-  acp: {
-    runtime: "acp-stdio",
-    ...COMMON_CAPABILITIES,
-    structured_output: false,
-    supports_session_resume: true,
-    // Runtime request-scoped MCP servers are not projected into ACP sessions.
-    // Profiles may own static ACP MCP configuration, but that is not the
-    // supports_mcp contract advertised to the route capability gate.
-    supports_mcp: false,
-    supports_skills: false,
-    supports_builtin_tools: false,
-    supports_live_input: false,
-    supports_native_subagents: false,
-    supports_request_tool_environment: false,
-    tool_policy: TOOL_POLICY_ALLOW_ALL_ONLY,
+    tool_policy: TOOL_POLICY_PROJECTED,
   },
 };
 
-export function runtimeCapabilities(sdkOrModel) {
-  if (!sdkOrModel) throw new Error("runtimeCapabilities requires a model reference or sdk kind");
-  const sdk = typeof sdkOrModel === "string" ? sdkOrModel : sdkOrModel?.sdk;
-  if (!sdk) throw new Error("runtimeCapabilities: unrecognized argument");
-  const caps = RUNTIME_CAPABILITIES[sdk];
-  if (!caps) throw new Error(`unknown provider sdk: ${sdk}`);
-  return { kind: sdk, ...caps };
+/**
+ * @param {*} [_legacySelection]
+ * @returns {{kind: "pi"} & typeof RUNTIME_CAPABILITIES.pi}
+ */
+export function runtimeCapabilities(_legacySelection) {
+  // Capability lookup no longer selects a backend. The ignored optional value
+  // keeps the still-separate router migration type-safe without allowing its
+  // legacy route ids to influence the sole surviving runtime contract.
+  return { kind: "pi", ...RUNTIME_CAPABILITIES.pi };
 }