@oh-my-pi/pi-ai 17.3.4 → 17.3.7

package/src/oneshot-retry.ts

@@ -0,0 +1,230 @@
+import { extractRetryHint } from "@oh-my-pi/pi-utils";
+import * as AIError from "./error";
+import type { AssistantMessage } from "./types";
+import { getHeadersFromError, getRetryAfterMsFromHeaders, type HeadersLike } from "./utils/retry-after";
+
+/**
+ * Transient-failure retry for **oneshot** (non-agent-loop) completions.
+ *
+ * Why this exists: `streamSimple`/`completeSimple` retry *auth* failures
+ * (credential rotation) but deliberately surface *transient* provider failures
+ * — Anthropic `overloaded_error`, `rate_limit_error`, HTTP 429/500/502/503/529
+ * — as a **resolved** `AssistantMessage` with `stopReason: "error"`. For the
+ * main agent turn that is correct: `TurnRecovery` owns recovery there, and it
+ * must refuse to replay once tool calls or visible text have streamed.
+ *
+ * Oneshots have no such hazard. A summary, title, handoff, or image
+ * description produces no side effects, so re-issuing the whole request is
+ * safe and is almost always what the caller wants. Before this helper every
+ * oneshot call site had to re-implement that decision, and most did not —
+ * failing on the first blip, or swallowing it into `null` so a transient
+ * overload was indistinguishable from a legitimate empty result.
+ *
+ * Classification reuses the existing provider predicates (`AIError`), so the
+ * set of retryable Anthropic failures stays defined in exactly one place.
+ * Usage limits are included: unlike the provider loop — which excludes them so
+ * credential rotation can own them — a oneshot has no rotation layer above it,
+ * and the retry hint the provider supplies (`retry-after`, "try again in ~5m")
+ * is honored, so waiting is the correct response.
+ */
+export interface OneshotRetryOptions {
+	/** Total attempts, including the first. Default 3. Values < 1 are treated as 1. */
+	maxAttempts?: number;
+	/** First backoff step in ms; doubles per attempt. Default 500. */
+	baseDelayMs?: number;
+	/**
+	 * Upper bound for a single wait. Default 30_000. A provider retry hint
+	 * longer than this aborts the retry instead of parking the caller — the
+	 * error surfaces so higher-level recovery (or the user) can decide.
+	 */
+	maxDelayMs?: number;
+	/**
+	 * Stops further attempts. Two distinct paths, both preserving the caller's
+	 * intent: an abort already visible when an attempt settles surfaces that
+	 * attempt's own result (`completeSimple` reports `stopReason: "aborted"`),
+	 * while an abort that lands during the backoff wait rejects with the abort
+	 * reason — a user cancel stays a cancel and is never relabelled as the
+	 * provider failure we happened to be waiting on.
+	 *
+	 * This helper does NOT pass the signal into `run` — cancelling the in-flight
+	 * request is the closure's job, because a per-attempt deadline must be
+	 * rebuilt on every attempt. Construct it inside `run`
+	 * (`signal: AbortSignal.timeout(MS)`, or `AbortSignal.any([outer, perAttempt])`);
+	 * a deadline captured outside would fire once and then abort every retry,
+	 * silently turning this helper into a single attempt.
+	 */
+	signal?: AbortSignal;
+	/**
+	 * Headers of the attempt that just failed, used to honor `retry-after`.
+	 *
+	 * Load-bearing: a transient Anthropic failure arrives as a **resolved**
+	 * `AssistantMessage`, and `AssistantMessage` carries no headers — so without
+	 * this the real `retry-after` / `x-ratelimit-reset` values on a 429/529 are
+	 * invisible and only the (usually hint-free) error text is available.
+	 * Callers that already capture headers via `SimpleStreamOptions.onResponse`
+	 * should return the latest capture here; it is read once per failed attempt.
+	 * Thrown errors need no wiring — headers are recovered from the error itself.
+	 */
+	getResponseHeaders?: () => HeadersLike;
+	/** Observability hook. Fires immediately before sleeping. */
+	onRetry?: (info: OneshotRetryInfo) => void;
+}
+
+export interface OneshotRetryInfo {
+	/** 1-based index of the attempt that just failed. */
+	attempt: number;
+	maxAttempts: number;
+	delayMs: number;
+	/** True when `delayMs` came from a provider retry hint rather than backoff. */
+	fromRetryHint: boolean;
+	errorMessage: string;
+	/** `AIError` classification bits of the failure. */
+	errorId: number;
+}
+
+const DEFAULT_MAX_ATTEMPTS = 3;
+const DEFAULT_BASE_DELAY_MS = 500;
+const DEFAULT_MAX_DELAY_MS = 30_000;
+/** Cap on pure backoff growth. A provider hint may still exceed this, up to `maxDelayMs`. */
+const BACKOFF_CEILING_MS = 8_000;
+const RETRY_AFTER_MS_SUFFIX = /(?:^|\s)retry-after-ms=([0-9]+(?:\.[0-9]+)?)(?=\s|$)/i;
+
+function backoffDelayMs(attempt: number, baseDelayMs: number): number {
+	const growth = Math.min(baseDelayMs * 2 ** (attempt - 1), BACKOFF_CEILING_MS);
+	// 75-100% jitter, matching the provider loop and TurnRecovery, so a fleet of
+	// concurrent oneshots does not re-converge on the same instant.
+	return Math.round(growth * (0.75 + Math.random() * 0.25));
+}
+
+/** Retryable when the provider says transient, or when it says "wait, then retry". */
+function isRetryableOneshotFailure(errorId: number, errorStatus: number | undefined, errorMessage: string): boolean {
+	// llama.cpp reports deterministic tool-call JSON parse failures as HTTP 500.
+	// Replaying the same prompt produces the same malformed output.
+	if (AIError.LLAMA_CPP_TOOL_CALL_PARSE_PATTERN.test(errorMessage)) return false;
+	if (AIError.is(errorId, AIError.Flag.ContentBlocked)) return false;
+	return (
+		AIError.isTransientStatus(errorStatus) ||
+		AIError.is(errorId, AIError.Flag.Transient) ||
+		AIError.is(errorId, AIError.Flag.UsageLimit) ||
+		AIError.retriable(errorId)
+	);
+}
+
+function sleep(delayMs: number, signal?: AbortSignal): Promise<void> {
+	if (delayMs <= 0) return Promise.resolve();
+	const { promise, resolve, reject } = Promise.withResolvers<void>();
+	const timer = setTimeout(() => {
+		signal?.removeEventListener("abort", onAbort);
+		resolve();
+	}, delayMs);
+	const onAbort = () => {
+		clearTimeout(timer);
+		reject(signal?.reason ?? new AIError.AbortError("oneshot retry aborted"));
+	};
+	if (signal) {
+		if (signal.aborted) {
+			clearTimeout(timer);
+			return Promise.reject(signal.reason ?? new AIError.AbortError("oneshot retry aborted"));
+		}
+		signal.addEventListener("abort", onAbort, { once: true });
+	}
+	return promise;
+}
+
+/**
+ * Run a oneshot completion, retrying transient provider failures.
+ *
+ * Handles both failure shapes: a resolved `AssistantMessage` carrying
+ * `stopReason: "error"` (what `completeSimple` produces) and a thrown error
+ * (what the raw HTTP helpers produce). A non-retryable failure is returned or
+ * rethrown unchanged, so existing caller error handling keeps working — this
+ * only removes the *first-blip* failure mode.
+ */
+export async function retryTransientCompletion(
+	run: (attempt: number) => Promise<AssistantMessage>,
+	options?: OneshotRetryOptions,
+): Promise<AssistantMessage> {
+	const maxAttempts = Math.max(1, options?.maxAttempts ?? DEFAULT_MAX_ATTEMPTS);
+	const baseDelayMs = options?.baseDelayMs ?? DEFAULT_BASE_DELAY_MS;
+	const maxDelayMs = options?.maxDelayMs ?? DEFAULT_MAX_DELAY_MS;
+	const signal = options?.signal;
+
+	for (let attempt = 1; ; attempt++) {
+		let message: AssistantMessage | undefined;
+		let thrown: unknown;
+		try {
+			message = await run(attempt);
+			if (message.stopReason !== "error") return message;
+		} catch (error) {
+			thrown = error;
+		}
+		// A caller abort is never a transient failure — surface it immediately so
+		// cancellation stays responsive.
+		if (signal?.aborted) {
+			if (thrown !== undefined) throw thrown;
+			return message as AssistantMessage;
+		}
+
+		const errorId =
+			thrown !== undefined ? AIError.classify(thrown) : AIError.classifyMessage(message as AssistantMessage);
+		if (AIError.is(errorId, AIError.Flag.Abort) || AIError.is(errorId, AIError.Flag.UserInterrupt)) {
+			if (thrown !== undefined) throw thrown;
+			return message as AssistantMessage;
+		}
+		const errorMessage =
+			thrown !== undefined
+				? thrown instanceof Error
+					? thrown.message
+					: String(thrown)
+				: ((message as AssistantMessage).errorMessage ?? "unknown error");
+		const errorStatus = thrown !== undefined ? AIError.status(thrown) : (message as AssistantMessage).errorStatus;
+		const lastAttempt = attempt >= maxAttempts;
+		if (lastAttempt || !isRetryableOneshotFailure(errorId, errorStatus, errorMessage)) {
+			if (thrown !== undefined) throw thrown;
+			return message as AssistantMessage;
+		}
+
+		// Headers first: a real Anthropic 429/529 carries `retry-after` /
+		// `x-ratelimit-reset*` in the response, and the resolved AssistantMessage
+		// has none — the caller supplies them via `getResponseHeaders`. Thrown
+		// errors (e.g. AnthropicApiError) carry their own headers.
+		const headers: HeadersLike = thrown !== undefined ? getHeadersFromError(thrown) : options?.getResponseHeaders?.();
+		const headerHintMs = getRetryAfterMsFromHeaders(headers);
+		const extractedTextHintMs = extractRetryHint(undefined, errorMessage);
+		const suffixValue = RETRY_AFTER_MS_SUFFIX.exec(errorMessage)?.[1];
+		const parsedSuffixMs = suffixValue === undefined ? undefined : Number(suffixValue);
+		const suffixHintMs =
+			parsedSuffixMs !== undefined && Number.isFinite(parsedSuffixMs) && parsedSuffixMs > 0
+				? Math.ceil(parsedSuffixMs)
+				: undefined;
+		const textHintMs =
+			extractedTextHintMs === undefined && suffixHintMs === undefined
+				? undefined
+				: Math.max(extractedTextHintMs ?? 0, suffixHintMs ?? 0);
+		const hintMs =
+			headerHintMs === undefined && textHintMs === undefined
+				? undefined
+				: Math.max(headerHintMs ?? 0, textHintMs ?? 0);
+		// An over-cap hint means "come back much later"; parking a oneshot that
+		// long is worse than surfacing the failure to the caller.
+		if (hintMs !== undefined && hintMs > maxDelayMs) {
+			if (thrown !== undefined) throw thrown;
+			return message as AssistantMessage;
+		}
+		const backoff = backoffDelayMs(attempt, baseDelayMs);
+		const delayMs = Math.min(Math.max(hintMs ?? 0, backoff), maxDelayMs);
+
+		options?.onRetry?.({
+			attempt,
+			maxAttempts,
+			delayMs,
+			fromRetryHint: hintMs !== undefined && hintMs >= backoff,
+			errorMessage,
+			errorId,
+		});
+		// Aborting mid-backoff rejects with the caller's abort reason: a user
+		// cancel must stay a cancel, not get relabelled as the provider failure we
+		// happened to be waiting on.
+		await sleep(delayMs, signal);
+	}
+}

package/src/providers/anthropic-messages-server.ts

@@ -27,7 +27,7 @@ import {
 	type AnthropicUserContentBlock,
 	anthropicMessagesRequestSchema,
 } from "./anthropic-messages-server-schema";
-import { isAnthropicWebSearchHistoryBlock } from "./anthropic-wire";
+import { isAnthropicServerToolHistoryBlock } from "./anthropic-wire";
 
 /**
  * Anthropic Messages API (https://docs.anthropic.com/en/api/messages) ↔ pi-ai
@@ -216,10 +216,10 @@ function walkAssistantContent(
 				break;
 			case "server_tool_use":
 			case "web_search_tool_result":
-				if (isAnthropicWebSearchHistoryBlock(block)) {
-					// Native web-search call/result. Anthropic requires these
-					// replayed verbatim (encrypted_content included), so retain
-					// the block instead of flattening it to text.
+			case "tool_search_tool_result":
+				if (isAnthropicServerToolHistoryBlock(block)) {
+					// Anthropic requires supported server-tool call/results replayed
+					// verbatim, so retain each opaque block instead of flattening it.
 					out.push({ type: "anthropicServerTool", block: { ...block } });
 				} else {
 					// Other server tools use distinct result block types that omp

package/src/providers/anthropic-wire.ts

@@ -77,6 +77,11 @@ export type ServerToolUseBlockParam = {
 /** Web-search server-tool call whose matching result is replayable by omp. */
 export type WebSearchServerToolUseBlockParam = ServerToolUseBlockParam & { name: "web_search" };
 
+/** Tool-search server-tool call whose matching result is replayable by omp. */
+export type ToolSearchServerToolUseBlockParam = ServerToolUseBlockParam & {
+	name: "tool_search_tool_regex" | "tool_search_tool_bm25";
+};
+
 /** Native web-search result replayed inside an assistant turn. */
 export type WebSearchToolResultBlockParam = {
 	type: "web_search_tool_result";
@@ -85,18 +90,37 @@ export type WebSearchToolResultBlockParam = {
 	[key: string]: unknown;
 };
 
-/** True for the complete native web-search history variants omp can replay. */
-export function isAnthropicWebSearchHistoryBlock(block: {
+/** Native tool-search result replayed inside an assistant turn. */
+export type ToolSearchToolResultBlockParam = {
+	type: "tool_search_tool_result";
+	tool_use_id: string;
+	content: unknown;
+	[key: string]: unknown;
+};
+
+/** Anthropic server-tool history variants omp can replay atomically. */
+export type AnthropicServerToolHistoryBlockParam =
+	| WebSearchServerToolUseBlockParam
+	| WebSearchToolResultBlockParam
+	| ToolSearchServerToolUseBlockParam
+	| ToolSearchToolResultBlockParam;
+
+/** True when a block is complete Anthropic server-tool history omp can replay. */
+export function isAnthropicServerToolHistoryBlock(block: {
 	type: string;
 	name?: unknown;
 	id?: unknown;
 	tool_use_id?: unknown;
 	content?: unknown;
-}): block is WebSearchServerToolUseBlockParam | WebSearchToolResultBlockParam {
+}): block is AnthropicServerToolHistoryBlockParam {
 	if (block.type === "server_tool_use") {
-		return block.name === "web_search" && typeof block.id === "string" && block.id.length > 0;
+		const supportedName =
+			block.name === "web_search" ||
+			block.name === "tool_search_tool_regex" ||
+			block.name === "tool_search_tool_bm25";
+		return supportedName && typeof block.id === "string" && block.id.length > 0;
 	}
-	if (block.type === "web_search_tool_result") {
+	if (block.type === "web_search_tool_result" || block.type === "tool_search_tool_result") {
 		return typeof block.tool_use_id === "string" && block.tool_use_id.length > 0 && Object.hasOwn(block, "content");
 	}
 	return false;
@@ -132,6 +156,7 @@ export type ContentBlockParam =
 	| ToolResultBlockParam
 	| ServerToolUseBlockParam
 	| WebSearchToolResultBlockParam
+	| ToolSearchToolResultBlockParam
 	| ThinkingBlockParam
 	| RedactedThinkingBlockParam
 	| FallbackBlockParam;
@@ -319,6 +344,7 @@ export type ResponseContentBlock =
 	| { type: "tool_use"; id: string; name: string; input?: Record<string, unknown> | null }
 	| ServerToolUseBlockParam
 	| WebSearchToolResultBlockParam
+	| ToolSearchToolResultBlockParam
 	| { type: "fallback"; from: { model: string }; to: { model: string } };
 
 export type ContentBlockDelta =

package/src/providers/anthropic.ts

@@ -79,7 +79,7 @@ import {
 	type Usage as AnthropicWireUsage,
 	type ContentBlockParam,
 	type FallbackParam,
-	isAnthropicWebSearchHistoryBlock,
+	isAnthropicServerToolHistoryBlock,
 	type MessageCreateParams,
 	type MessageCreateParamsStreaming,
 	type MessageParam,
@@ -2404,7 +2404,7 @@ const streamAnthropicOnce = (
 								streamedReplayUnsafeContent = true;
 								const block: Block = {
 									type: "thinking",
-									thinking: "",
+									thinking: event.content_block.thinking ?? "",
 									thinkingSignature: "",
 									[kStreamingBlockIndex]: event.index,
 								};
@@ -2416,6 +2416,14 @@ const streamAnthropicOnce = (
 									contentIndex,
 									partial: output,
 								});
+								if (block.thinking) {
+									stream.push({
+										type: "thinking_delta",
+										contentIndex,
+										delta: block.thinking,
+										partial: output,
+									});
+								}
 							} else if (event.content_block.type === "redacted_thinking") {
 								streamedReplayUnsafeContent = true;
 								const block: Block = {
@@ -2429,8 +2437,11 @@ const streamAnthropicOnce = (
 									kind: "redactedThinking",
 								});
 							} else if (
-								isAnthropicWebSearchHistoryBlock(event.content_block) &&
-								umansGatewayWebSearchHeader === undefined
+								isAnthropicServerToolHistoryBlock(event.content_block) &&
+								(umansGatewayWebSearchHeader === undefined ||
+									(event.content_block.type === "server_tool_use"
+										? event.content_block.name !== "web_search"
+										: event.content_block.type !== "web_search_tool_result"))
 							) {
 								streamedReplayUnsafeContent = true;
 								const block: Block = {

package/src/providers/cursor.ts

@@ -293,11 +293,23 @@ const CURSOR_PROXY_TUNNEL_TIMEOUT_MS = 30_000;
  * model reads it and should route around the capability, not retry the call.
  */
 const NOT_IMPLEMENTED_SUFFIX = "not implemented by this client";
+/** Bare gRPC `resource_exhausted` end-streams (also inside a Connect error message). */
+const RESOURCE_EXHAUSTED_PATTERN = /resource.?exhausted/i;
 const NOT_IMPLEMENTED = `Not implemented by this client`;
 
 const conversationStateCache = new Map<string, ConversationStateStructure>();
 const conversationBlobStores = new Map<string, Map<string, Uint8Array>>();
 const warnedCursorKimiK3ReplayMessages = new Set<string>();
+/**
+ * Base conversation id → rotated wire id (#8345). Cursor's backend can pin a
+ * per-conversation rejection (bare `resource_exhausted`, zero tokens) to one
+ * conversationId forever; the session is then unusable until /fork mints a
+ * new id. On the first such failure the id is rotated once and the cached
+ * state migrates, so the retry loop's next attempt starts a fresh
+ * conversation — the same recovery /fork performs. Keyed by the base id the
+ * caller derived, so a failed rotation is never repeated.
+ */
+const rotatedConversationIds = new Map<string, string>();
 
 export interface CursorOptions extends StreamOptions {
 	customSystemPrompt?: string;
@@ -587,13 +599,19 @@ export const streamCursor: StreamFunction<"cursor-agent"> = (
 			h2Completion.resolve();
 		};
 
+		// Hoisted out of the try block: the #8345 rotation in the catch path
+		// needs both ids, and the catch block cannot see try-scoped consts.
+		let baseConversationId: string | undefined;
+		let conversationId: string | undefined;
+		let usageState: UsageState | undefined;
 		try {
 			const apiKey = options?.apiKey;
 			if (!apiKey) {
 				throw new AIError.MissingApiKeyError(undefined, "Cursor API key (access token) is required");
 			}
 
-			const conversationId = options?.conversationId ?? options?.sessionId ?? crypto.randomUUID();
+			baseConversationId = options?.conversationId ?? options?.sessionId ?? crypto.randomUUID();
+			conversationId = rotatedConversationIds.get(baseConversationId) ?? baseConversationId;
 			const blobStore = conversationBlobStores.get(conversationId) ?? new Map<string, Uint8Array>();
 			conversationBlobStores.set(conversationId, blobStore);
 			const cachedState = conversationStateCache.get(conversationId);
@@ -667,7 +685,7 @@ export const streamCursor: StreamFunction<"cursor-agent"> = (
 			let currentThinkingBlock: (ThinkingContent & { [kStreamingBlockIndex]: number }) | null = null;
 			let currentToolCall: ToolCallState | null = null;
 			const resolvedMcpToolCallIds = new Set<string>();
-			const usageState: UsageState = { sawTokenDelta: false };
+			usageState = { sawTokenDelta: false };
 
 			const state: BlockState = {
 				get currentTextBlock() {
@@ -702,7 +720,7 @@ export const streamCursor: StreamFunction<"cursor-agent"> = (
 			openBlockState = state;
 
 			const onConversationCheckpoint = (checkpoint: ConversationStateStructure) => {
-				conversationStateCache.set(conversationId, checkpoint);
+				conversationStateCache.set(conversationId!, checkpoint);
 			};
 
 			h2Request.on("response", headers => {
@@ -759,7 +777,7 @@ export const streamCursor: StreamFunction<"cursor-agent"> = (
 							h2Request!,
 							options?.execHandlers,
 							options?.onToolResult,
-							usageState,
+							usageState!,
 							requestContextTools,
 							onConversationCheckpoint,
 						).catch(error => {
@@ -873,6 +891,28 @@ export const streamCursor: StreamFunction<"cursor-agent"> = (
 				flushOpenToolCalls(output, stream, openBlockState);
 			}
 			const result = await AIError.finalize(error, { api: model.api, signal: options?.signal });
+			// #8345: a server-side per-conversation rejection surfaces as a bare
+			// resource_exhausted with zero tokens — the conversation is poisoned,
+			// not the account (sibling conversations keep working). Rotate the
+			// wire id once and migrate the cached state so the next attempt (the
+			// caller's retry loop) starts a fresh conversation, exactly like
+			// /fork. Only the first failure rotates; repeated failures keep the
+			// rotated id so a genuine account-level exhaustion is not hidden.
+			if (
+				conversationId !== undefined &&
+				baseConversationId !== undefined &&
+				usageState !== undefined &&
+				!usageState.sawTokenDelta &&
+				RESOURCE_EXHAUSTED_PATTERN.test(result.message) &&
+				!rotatedConversationIds.has(baseConversationId)
+			) {
+				const rotated = crypto.randomUUID();
+				rotatedConversationIds.set(baseConversationId, rotated);
+				const state = conversationStateCache.get(conversationId);
+				if (state) conversationStateCache.set(rotated, state);
+				const blobs = conversationBlobStores.get(conversationId);
+				if (blobs) conversationBlobStores.set(rotated, blobs);
+			}
 			output.stopReason = result.stopReason;
 			output.errorStatus = result.status;
 			output.errorId = result.id;