fix: stop retrying oversized requests and recover with one forced compaction

packages/agent-core/src/pi-turn-runner/context-overflow-recovery.ts

@@ -0,0 +1,128 @@
+import type { StreamFn } from '@earendil-works/pi-agent-core';
+import {
+  isContextOverflow,
+  type AssistantMessage,
+  type AssistantMessageEvent,
+  type AssistantMessageEventStream,
+  type Message,
+} from '@earendil-works/pi-ai';
+import { isLLMRequestOversized, normalizeLLMError } from '@mavis/shared/llm-error-classifier';
+
+/**
+ * Per-turn slot connecting the provider stream (built before the Agent exists)
+ * with the Agent-owned context lifecycle that can shrink the context.
+ */
+export interface ContextOverflowRecoveryState {
+  /**
+   * Shrinks the context of the request that just failed and returns the
+   * provider-bound messages to resend, or `undefined` when nothing changed.
+   * Installed by `setLLMHook` only when recovery hooks exist.
+   */
+  recover?: (signal: AbortSignal | undefined) => Promise<Message[] | undefined>;
+  /** Recovery runs at most once per turn, so a second rejection is terminal. */
+  attempted: boolean;
+}
+
+export function newContextOverflowRecoveryState(): ContextOverflowRecoveryState {
+  return { attempted: false };
+}
+
+/**
+ * True when a failed provider response means the request was too large to be
+ * accepted as sent: HTTP 413 / `request_too_large`, the local request-body
+ * check, or a context-window overflow recognized by pi-ai.
+ */
+export function isRecoverableContextOverflow(
+  message: AssistantMessage,
+  contextWindow: number | undefined,
+): boolean {
+  if (message.stopReason !== 'error') return false;
+  if (isContextOverflow(message, contextWindow)) return true;
+  return isLLMRequestOversized(
+    normalizeLLMError({
+      finishReason: message.stopReason,
+      ...(message.errorMessage ? { errorMessage: message.errorMessage } : {}),
+    }),
+  );
+}
+
+/**
+ * When the provider rejects a main-agent request as too large before producing
+ * any output, compact once and resend once. Successful streams only pay for
+ * buffering their first non-`start` event. A second rejection, a recovery that
+ * changes nothing, or any other failure is surfaced exactly as received.
+ */
+export function withContextOverflowRecovery(
+  inner: StreamFn,
+  state: ContextOverflowRecoveryState,
+): StreamFn {
+  return (async (model, context, options) => {
+    const stream = await inner(model, context, options);
+    if (!state.recover || state.attempted) return stream;
+    const iterator = stream[Symbol.asyncIterator]();
+    const buffered: AssistantMessageEvent[] = [];
+    try {
+      for (;;) {
+        const next = await iterator.next();
+        if (next.done) return replay(stream, buffered, undefined);
+        buffered.push(next.value);
+        if (next.value.type !== 'start') break;
+      }
+    } catch (error) {
+      return replay(stream, buffered, undefined, { error });
+    }
+    const first = buffered.at(-1);
+    if (
+      first?.type !== 'error' ||
+      options?.signal?.aborted ||
+      !isRecoverableContextOverflow(first.error, model.contextWindow)
+    ) {
+      return replay(stream, buffered, iterator);
+    }
+    state.attempted = true;
+    let messages: Message[] | undefined;
+    try {
+      messages = await state.recover(options?.signal);
+    } catch {
+      // A failing recovery must not mask the provider rejection it tried to fix.
+      return replay(stream, buffered, iterator);
+    }
+    if (!messages || options?.signal?.aborted) return replay(stream, buffered, iterator);
+    try {
+      void Promise.resolve(iterator.return?.()).catch(() => undefined);
+    } catch {
+      // Best-effort cleanup of the rejected attempt.
+    }
+    return inner(model, { ...context, messages }, options);
+  }) as StreamFn;
+}
+
+function replay(
+  stream: AssistantMessageEventStream,
+  buffered: readonly AssistantMessageEvent[],
+  iterator: AsyncIterator<AssistantMessageEvent> | undefined,
+  pending?: { readonly error: unknown },
+): AssistantMessageEventStream {
+  return {
+    [Symbol.asyncIterator]() {
+      let index = 0;
+      let thrown = false;
+      return {
+        async next(): Promise<IteratorResult<AssistantMessageEvent>> {
+          if (index < buffered.length) return { value: buffered[index++]!, done: false };
+          if (pending && !thrown) {
+            thrown = true;
+            throw pending.error;
+          }
+          if (!iterator) return { value: undefined, done: true };
+          return iterator.next();
+        },
+        async return(value?: unknown): Promise<IteratorResult<AssistantMessageEvent>> {
+          await iterator?.return?.(value);
+          return { value: undefined, done: true };
+        },
+      };
+    },
+    result: () => stream.result(),
+  } as unknown as AssistantMessageEventStream;
+}

packages/agent-core/src/pi-turn-runner/hooks.ts

@@ -26,6 +26,12 @@ export interface PiBeforeLlmCallHookInput {
   sessionId: string;
   turnId: string;
   phase: PiBeforeLlmCallPhase;
+  /**
+   * Set only when the runner re-runs `contextOverflowRecoveryHook` after the
+   * provider rejected the request as too large (HTTP 413 / context overflow).
+   * Compaction hooks must compact regardless of their own size estimate.
+   */
+  trigger?: 'context_overflow_recovery';
   /** Current provider request view after request-only hook transforms. */
   messages: AgentMessage[];
   /** Current durable history, updated only by canonical replacements. */
@@ -251,6 +257,15 @@ export type PiOnTurnEndHook = PiOnStepEndHook;
 
 export interface PiTurnHooks {
   beforeLlmCallHook?: readonly PiBeforeLlmCallHook[];
+  /**
+   * Hooks that may shrink the context after the provider rejected a main-agent
+   * request as too large (HTTP 413, `request_too_large`, the local request-body
+   * check, or a context-window overflow) before producing any output. The
+   * runner runs them once per turn with `trigger: 'context_overflow_recovery'`
+   * and, when they replace the context, resends the request once. Only
+   * compaction belongs here: reminder or append hooks would run twice.
+   */
+  contextOverflowRecoveryHook?: readonly PiBeforeLlmCallHook[];
   /** Best-effort observers of actual logical main-agent provider calls. */
   onLlmCallPreparedHook?: readonly PiOnLlmCallPreparedHook[];
   afterLlmCallHook?: readonly PiAfterLlmCallHook[];

packages/agent-core/src/pi-turn-runner/llm.ts

@@ -5,6 +5,7 @@ import { resolveLLMStreamTimeouts, withLLMStreamTimeouts } from './llm-stream-ti
 import type {
   PiAfterLlmReplacementCommit,
   PiBeforeLlmCallAppendMessage,
+  PiBeforeLlmCallHook,
   PiBeforeLlmCallReplaceMetadata,
 } from './hooks.js';
 import type { LLMModelConfig } from './types.js';
@@ -41,40 +42,24 @@ export function composeStreamFn(resolved: LLMModelConfig): StreamFn {
 }
 
 export function setLLMHook(agent: Agent, turn: turnState, history: turnHistory): void {
-  if (turn.hooks.beforeLLM.length === 0 && turn.hooks.afterLLM.length === 0) return;
+  const recoveryHooks = turn.hooks.contextOverflowRecovery;
+  if (
+    turn.hooks.beforeLLM.length === 0 &&
+    turn.hooks.afterLLM.length === 0 &&
+    recoveryHooks.length === 0
+  ) {
+    return;
+  }
 
   let checkpointCallCount = 0;
-  agent.transformContext = async (messages, signal) => {
-    const rejectedMessages = messages.filter((message) =>
-      turn.rejectedAssistantMessages.has(message),
-    );
-    if (rejectedMessages.length > 0) {
-      messages.splice(
-        0,
-        messages.length,
-        ...messages.filter((message) => !turn.rejectedAssistantMessages.has(message)),
-      );
-      for (const message of rejectedMessages) turn.rejectedAssistantMessages.delete(message);
-    }
+  // Pi's live loop context of the latest main-agent request; overflow recovery
+  // compacts exactly this context and keeps the loop in sync through it.
+  let lastRequest: { messages: AgentMessage[]; phase: 'initial' | 'iteration' } | undefined;
 
-    if (turn.hooks.beforeLLM.length === 0) return messages;
-    const phase = checkpointCallCount === 0 ? 'initial' : 'iteration';
-    checkpointCallCount += 1;
-    const checkpoint = await runBeforeLLM(turn, messages, phase, signal);
-    if (checkpoint.type === 'abort') {
-      turn.metrics.tryRecordTerminalFailure({
-        errorSource: 'before_llm_hook',
-        errorKind: 'hook_error',
-      });
-      throw new Error(`before_llm_checkpoint_aborted: ${checkpoint.reason}`);
-    }
-    if (checkpoint.type === 'respond') {
-      turn.syntheticResponse.pending = {
-        text: checkpoint.text,
-        reason: checkpoint.reason,
-      };
-      return messages;
-    }
+  const applyCheckpoint = async (
+    messages: AgentMessage[],
+    checkpoint: Exclude<BeforeLLMResult, { type: 'abort' } | { type: 'respond' }>,
+  ): Promise<AgentMessage[]> => {
     if (checkpoint.type === 'requestOnly') {
       return checkpoint.messages;
     }
@@ -130,13 +115,89 @@ export function setLLMHook(agent: Agent, turn: turnState, history: turnHistory):
     }
     return checkpoint.messages;
   };
+
+  agent.transformContext = async (messages, signal) => {
+    const rejectedMessages = messages.filter((message) =>
+      turn.rejectedAssistantMessages.has(message),
+    );
+    if (rejectedMessages.length > 0) {
+      messages.splice(
+        0,
+        messages.length,
+        ...messages.filter((message) => !turn.rejectedAssistantMessages.has(message)),
+      );
+      for (const message of rejectedMessages) turn.rejectedAssistantMessages.delete(message);
+    }
+
+    const phase = checkpointCallCount === 0 ? 'initial' : 'iteration';
+    lastRequest = { messages, phase };
+    if (turn.hooks.beforeLLM.length === 0) return messages;
+    checkpointCallCount += 1;
+    const checkpoint = await runBeforeLLM(turn, messages, phase, signal, turn.hooks.beforeLLM);
+    if (checkpoint.type === 'abort') {
+      turn.metrics.tryRecordTerminalFailure({
+        errorSource: 'before_llm_hook',
+        errorKind: 'hook_error',
+      });
+      throw new Error(`before_llm_checkpoint_aborted: ${checkpoint.reason}`);
+    }
+    if (checkpoint.type === 'respond') {
+      turn.syntheticResponse.pending = {
+        text: checkpoint.text,
+        reason: checkpoint.reason,
+      };
+      return messages;
+    }
+    return applyCheckpoint(messages, checkpoint);
+  };
+
+  if (recoveryHooks.length === 0) return;
+  turn.contextOverflowRecovery.recover = async (signal) => {
+    const request = lastRequest;
+    if (!request) return undefined;
+    const logContext = {
+      session_id: turn.input.sessionId,
+      turn_id: turn.input.turnId,
+      phase: request.phase,
+      provider: turn.llm.model.provider,
+    };
+    const checkpoint = await runBeforeLLM(
+      turn,
+      request.messages,
+      request.phase,
+      signal,
+      recoveryHooks,
+      'context_overflow_recovery',
+    );
+    if (
+      checkpoint.type === 'abort' ||
+      checkpoint.type === 'respond' ||
+      checkpoint.type === 'unchanged'
+    ) {
+      turn.logger.warn(
+        { ...logContext, outcome: checkpoint.type },
+        '[pi-turn-runner] context overflow recovery left the request unchanged; surfacing the provider rejection',
+      );
+      return undefined;
+    }
+    const messages = await applyCheckpoint(request.messages, checkpoint);
+    turn.logger.info(
+      { ...logContext, outcome: checkpoint.type, message_count: messages.length },
+      '[pi-turn-runner] context overflow recovery compacted the request; resending once',
+    );
+    return await agent.convertToLlm(messages);
+  };
 }
 
+type BeforeLLMResult = Awaited<ReturnType<typeof runBeforeLLM>>;
+
 async function runBeforeLLM(
   turn: turnState,
   messages: AgentMessage[],
   phase: 'initial' | 'iteration',
-  signal?: AbortSignal,
+  signal: AbortSignal | undefined,
+  hooks: readonly PiBeforeLlmCallHook[],
+  trigger?: 'context_overflow_recovery',
 ): Promise<
   | { type: 'unchanged'; messages: AgentMessage[] }
   | { type: 'requestOnly'; messages: AgentMessage[] }
@@ -172,13 +233,14 @@ async function runBeforeLLM(
         afterCommit?: PiAfterLlmReplacementCommit;
       }
     | undefined;
-  for (const hook of turn.hooks.beforeLLM) {
+  for (const hook of hooks) {
     let decision: Awaited<ReturnType<typeof hook>>;
     try {
       decision = await hook({
         sessionId: turn.input.sessionId,
         turnId: turn.input.turnId,
         phase,
+        ...(trigger ? { trigger } : {}),
         messages: [...currentMessages],
         canonicalMessages: [...canonicalMessages],
         model: turn.llm.model,